select for update 1 or N as json

test.sql

DROP TABLE test;
CREATE TABLE test (
  id SERIAL NOT NULL PRIMARY KEY,
  a TEXT DEFAULT 'A',
  b TEXT DEFAULT 'B',
  i INTEGER DEFAULT 0
);
INSERT INTO test DEFAULT VALUES ;
INSERT INTO test DEFAULT VALUES ;
INSERT INTO test DEFAULT VALUES ;
SELECT * FROM test;

test1.sql

CREATE OR REPLACE FUNCTION test1()
  RETURNS JSONB AS
$BODY$
DECLARE
  _json JSONB;
BEGIN

  WITH lck AS (
      SELECT
        id, a, b, i
      FROM test
      -- WHERE id = _id ...
      -- ORDER BY ...
      LIMIT 1
      FOR UPDATE
  ), upd AS (
    UPDATE test
    SET
      a = test.a||'A',
      b = test.b||'B',
      i = test.i + 1
    FROM lck
    WHERE test.id = lck.id
    RETURNING test.*
  ) SELECT row_to_json(f) from (select * from upd) f INTO _json;

  RETURN _json;

END;
$BODY$
LANGUAGE plpgsql VOLATILE;

testN.sql

CREATE OR REPLACE FUNCTION testN()
  RETURNS JSONB AS
$BODY$
DECLARE
  _json JSONB;
BEGIN

  WITH lck AS (
      SELECT
        id, a, b, i
      FROM test
      -- WHERE id = _id ...
      -- ORDER BY ...
      -- LIMIT N
      FOR UPDATE
  ), upd AS (
    UPDATE test
    SET
      a = test.a||'A',
      b = test.b||'B',
      i = test.i + 1
    FROM lck
    WHERE test.id = lck.id
    RETURNING test.*
  ) SELECT coalesce(json_agg(row_to_json(f)),'[]') from (select * from upd) f INTO _json;

  RETURN _json;

END;
$BODY$
LANGUAGE plpgsql VOLATILE;