Skip to content

Commit 7f0336e

Browse files
pyramationpyramation
committed
fix: recover RETURN NEXT variable when retvarno is lost
Thread enclosingForVarName through FOR loop contexts and add inferReturnNextVar() to recover the variable from: 1. Enclosing FOR loop variable (most common pattern) 2. Single user-declared variable (SETOF scalar, no loop) Only infers when returnInfo is provided (kind='setof'), preventing false inference on OUT-param functions where bare RETURN NEXT is correct. Also adds extractReturnInfo() to the test infrastructure so round-trip tests can extract the function's return type from the CREATE FUNCTION SQL and pass it to the deparser. Removes fixtures 47-48 from KNOWN_FAILING — they now pass. Updates snapshot tests from bug-demonstration to fix-verification.
1 parent b8b7643 commit 7f0336e

5 files changed

Lines changed: 184 additions & 28 deletions

File tree

‎packages/plpgsql-deparser/__tests__/__snapshots__/deparser-fixes.test.ts.snap‎

Lines changed: 25 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -183,7 +183,18 @@ exports[`plpgsql-deparser bug fixes PERFORM SELECT fix should strip SELECT keywo
183183
END"
184184
`;
185185

186-
exports[`plpgsql-deparser bug fixes RETURN NEXT retvarno loss (demonstrates bug) should lose the variable in RETURN NEXT (libpg-query does not serialize retvarno) 1`] = `
186+
exports[`plpgsql-deparser bug fixes RETURN NEXT retvarno recovery should leave RETURN NEXT bare for OUT-param functions 1`] = `
187+
"BEGIN
188+
FOR i IN 1..5 LOOP
189+
x := i;
190+
y := 'item_' || i::text;
191+
RETURN NEXT;
192+
END LOOP;
193+
RETURN;
194+
END"
195+
`;
196+
197+
exports[`plpgsql-deparser bug fixes RETURN NEXT retvarno recovery should leave RETURN NEXT bare when returnInfo is not provided 1`] = `
187198
"DECLARE
188199
v_entry jsonb;
189200
BEGIN
@@ -194,11 +205,22 @@ BEGIN
194205
END"
195206
`;
196207

197-
exports[`plpgsql-deparser bug fixes RETURN NEXT retvarno loss (demonstrates bug) should lose the variable in RETURN NEXT for declared var (no FOR loop) 1`] = `
208+
exports[`plpgsql-deparser bug fixes RETURN NEXT retvarno recovery should recover FOR loop variable in RETURN NEXT for SETOF functions 1`] = `
209+
"DECLARE
210+
v_entry jsonb;
211+
BEGIN
212+
FOR v_entry IN SELECT jsonb_array_elements(v_data) LOOP
213+
RETURN NEXT v_entry;
214+
END LOOP;
215+
RETURN;
216+
END"
217+
`;
218+
219+
exports[`plpgsql-deparser bug fixes RETURN NEXT retvarno recovery should recover single declared variable in RETURN NEXT for SETOF functions 1`] = `
198220
"DECLARE
199221
v_val int := 42;
200222
BEGIN
201-
RETURN NEXT;
223+
RETURN NEXT v_val;
202224
RETURN;
203225
END"
204226
`;

‎packages/plpgsql-deparser/__tests__/deparser-fixes.test.ts‎

Lines changed: 47 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1022,8 +1022,8 @@ END$$`;
10221022
});
10231023
});
10241024

1025-
describe('RETURN NEXT retvarno loss (demonstrates bug)', () => {
1026-
it('should lose the variable in RETURN NEXT (libpg-query does not serialize retvarno)', async () => {
1025+
describe('RETURN NEXT retvarno recovery', () => {
1026+
it('should recover FOR loop variable in RETURN NEXT for SETOF functions', async () => {
10271027
const sql = `CREATE FUNCTION test_return_next_for_var(v_data jsonb) RETURNS SETOF jsonb
10281028
LANGUAGE plpgsql AS $$
10291029
DECLARE
@@ -1036,16 +1036,14 @@ BEGIN
10361036
RETURN;
10371037
END$$`;
10381038

1039+
await testUtils.expectAstMatch('RETURN NEXT FOR loop var', sql);
10391040
const parsed = parsePlPgSQLSync(sql) as unknown as PLpgSQLParseResult;
1040-
const deparsed = deparseSync(parsed);
1041+
const deparsed = deparseSync(parsed, undefined, { kind: 'setof' });
10411042
expect(deparsed).toMatchSnapshot();
1042-
1043-
// BUG: the variable v_entry is lost — RETURN NEXT v_entry becomes bare RETURN NEXT
1044-
expect(deparsed).toContain('RETURN NEXT;');
1045-
expect(deparsed).not.toContain('RETURN NEXT v_entry');
1043+
expect(deparsed).toContain('RETURN NEXT v_entry');
10461044
});
10471045

1048-
it('should lose the variable in RETURN NEXT for declared var (no FOR loop)', async () => {
1046+
it('should recover single declared variable in RETURN NEXT for SETOF functions', async () => {
10491047
const sql = `CREATE FUNCTION test_return_next_declared_var() RETURNS SETOF int
10501048
LANGUAGE plpgsql AS $$
10511049
DECLARE
@@ -1054,13 +1052,51 @@ BEGIN
10541052
RETURN NEXT v_val;
10551053
END$$`;
10561054

1055+
await testUtils.expectAstMatch('RETURN NEXT declared var', sql);
10571056
const parsed = parsePlPgSQLSync(sql) as unknown as PLpgSQLParseResult;
1058-
const deparsed = deparseSync(parsed);
1057+
const deparsed = deparseSync(parsed, undefined, { kind: 'setof' });
1058+
expect(deparsed).toMatchSnapshot();
1059+
expect(deparsed).toContain('RETURN NEXT v_val');
1060+
});
1061+
1062+
it('should leave RETURN NEXT bare for OUT-param functions', async () => {
1063+
const sql = `CREATE FUNCTION test_return_next_out(OUT x integer, OUT y text) RETURNS SETOF record
1064+
LANGUAGE plpgsql AS $$
1065+
BEGIN
1066+
FOR i IN 1..5 LOOP
1067+
x := i;
1068+
y := 'item_' || i::text;
1069+
RETURN NEXT;
1070+
END LOOP;
1071+
RETURN;
1072+
END$$`;
1073+
1074+
await testUtils.expectAstMatch('RETURN NEXT OUT params', sql);
1075+
const parsed = parsePlPgSQLSync(sql) as unknown as PLpgSQLParseResult;
1076+
const deparsed = deparseSync(parsed, undefined, { kind: 'out_params' });
10591077
expect(deparsed).toMatchSnapshot();
1078+
expect(deparsed).toContain('RETURN NEXT;');
1079+
expect(deparsed).not.toMatch(/RETURN NEXT\s+\w/);
1080+
});
1081+
1082+
it('should leave RETURN NEXT bare when returnInfo is not provided', async () => {
1083+
const sql = `CREATE FUNCTION test_return_next_for_var(v_data jsonb) RETURNS SETOF jsonb
1084+
LANGUAGE plpgsql AS $$
1085+
DECLARE
1086+
v_entry jsonb;
1087+
BEGIN
1088+
FOR v_entry IN SELECT jsonb_array_elements(v_data)
1089+
LOOP
1090+
RETURN NEXT v_entry;
1091+
END LOOP;
1092+
RETURN;
1093+
END$$`;
10601094

1061-
// BUG: the variable v_val is lost — RETURN NEXT v_val becomes bare RETURN NEXT
1095+
const parsed = parsePlPgSQLSync(sql) as unknown as PLpgSQLParseResult;
1096+
// Without returnInfo, inference should NOT happen (safety)
1097+
const deparsed = deparseSync(parsed);
1098+
expect(deparsed).toMatchSnapshot();
10621099
expect(deparsed).toContain('RETURN NEXT;');
1063-
expect(deparsed).not.toContain('RETURN NEXT v_val');
10641100
});
10651101
});
10661102
});

‎packages/plpgsql-deparser/__tests__/plpgsql-deparser.test.ts‎

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -40,10 +40,7 @@ describe('PLpgSQLDeparser', () => {
4040
// - Labeled blocks with EXIT statements
4141
// - Nested DECLARE inside FOR loops (lineno-based scope tracking)
4242
const KNOWN_FAILING_FIXTURES = new Set<string>([
43-
// libpg-query does not serialize retvarno for PLpgSQL_stmt_return_next,
44-
// so RETURN NEXT <variable> loses the variable during round-trip.
45-
'plpgsql_deparser_fixes-47.sql', // RETURN NEXT v_entry (FOR loop)
46-
'plpgsql_deparser_fixes-48.sql', // RETURN NEXT v_val (declared var)
43+
// No known failures - all fixtures pass!
4744
]);
4845

4946
it('should round-trip ALL generated fixtures (excluding known failures)', async () => {

‎packages/plpgsql-deparser/src/plpgsql-deparser.ts‎

Lines changed: 71 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -86,6 +86,8 @@ export interface PLpgSQLDeparserContext {
8686
loopVarLinenos?: Set<number>;
8787
/** Map of block lineno to the set of datum indices that belong to that block */
8888
blockDatumMap?: Map<number, Set<number>>;
89+
/** Name of the enclosing FOR loop variable (used to recover RETURN NEXT <var> when retvarno is missing) */
90+
enclosingForVarName?: string;
8991
}
9092

9193
/**
@@ -1179,7 +1181,7 @@ export class PLpgSQLDeparser {
11791181
}
11801182

11811183
if (fori.body) {
1182-
const bodyContext = { ...context, indentLevel: context.indentLevel + 1 };
1184+
const bodyContext = { ...context, indentLevel: context.indentLevel + 1, enclosingForVarName: varName };
11831185
for (const stmt of fori.body) {
11841186
const stmtStr = this.deparseStmt(stmt, bodyContext);
11851187
parts.push(this.indent(stmtStr + ';', bodyContext.indentLevel));
@@ -1210,7 +1212,7 @@ export class PLpgSQLDeparser {
12101212
parts.push(`${kw('FOR')} ${varName} ${kw('IN')} ${query} ${kw('LOOP')}`);
12111213

12121214
if (fors.body) {
1213-
const bodyContext = { ...context, indentLevel: context.indentLevel + 1 };
1215+
const bodyContext = { ...context, indentLevel: context.indentLevel + 1, enclosingForVarName: varName };
12141216
for (const stmt of fors.body) {
12151217
const stmtStr = this.deparseStmt(stmt, bodyContext);
12161218
parts.push(this.indent(stmtStr + ';', bodyContext.indentLevel));
@@ -1246,7 +1248,7 @@ export class PLpgSQLDeparser {
12461248
parts.push(`${forClause} ${kw('LOOP')}`);
12471249

12481250
if (forc.body) {
1249-
const bodyContext = { ...context, indentLevel: context.indentLevel + 1 };
1251+
const bodyContext = { ...context, indentLevel: context.indentLevel + 1, enclosingForVarName: varName };
12501252
for (const stmt of forc.body) {
12511253
const stmtStr = this.deparseStmt(stmt, bodyContext);
12521254
parts.push(this.indent(stmtStr + ';', bodyContext.indentLevel));
@@ -1283,7 +1285,7 @@ export class PLpgSQLDeparser {
12831285
parts.push(`${kw('FOREACH')} ${varName}${sliceClause} ${kw('IN ARRAY')} ${expr} ${kw('LOOP')}`);
12841286

12851287
if (foreach.body) {
1286-
const bodyContext = { ...context, indentLevel: context.indentLevel + 1 };
1288+
const bodyContext = { ...context, indentLevel: context.indentLevel + 1, enclosingForVarName: varName };
12871289
for (const stmt of foreach.body) {
12881290
const stmtStr = this.deparseStmt(stmt, bodyContext);
12891291
parts.push(this.indent(stmtStr + ';', bodyContext.indentLevel));
@@ -1363,22 +1365,83 @@ export class PLpgSQLDeparser {
13631365

13641366
/**
13651367
* Deparse a RETURN NEXT statement
1368+
*
1369+
* libpg-query does not serialize `retvarno` for PLpgSQL_stmt_return_next,
1370+
* so when a function has `RETURN NEXT variable`, the variable reference is
1371+
* lost. We recover it by:
1372+
* 1. Using the enclosing FOR loop variable (most common pattern)
1373+
* 2. Scanning datums for the sole user-declared variable (non-loop case)
1374+
* Bare `RETURN NEXT` remains valid for TABLE / OUT-param functions.
13661375
*/
13671376
private deparseReturnNext(ret: PLpgSQL_stmt_return_next, context: PLpgSQLDeparserContext): string {
13681377
const kw = this.keyword;
1369-
1378+
13701379
if (ret.expr) {
13711380
return `${kw('RETURN NEXT')} ${this.deparseExpr(ret.expr)}`;
13721381
}
1373-
1382+
13741383
if (ret.retvarno !== undefined && ret.retvarno >= 0) {
13751384
const varName = this.getVarName(ret.retvarno, context);
13761385
return `${kw('RETURN NEXT')} ${varName}`;
13771386
}
1378-
1387+
1388+
// retvarno is missing (libpg-query serialization gap). Try to recover.
1389+
const inferred = this.inferReturnNextVar(context);
1390+
if (inferred) {
1391+
return `${kw('RETURN NEXT')} ${inferred}`;
1392+
}
1393+
13791394
return kw('RETURN NEXT');
13801395
}
13811396

1397+
/**
1398+
* Attempt to infer the variable for a bare RETURN NEXT when retvarno is
1399+
* not available. Returns the variable name or undefined.
1400+
*
1401+
* We only infer when `returnInfo` tells us the function returns SETOF
1402+
* scalar. Without returnInfo we cannot distinguish a legitimate bare
1403+
* `RETURN NEXT` (OUT-param / TABLE function) from a dropped retvarno,
1404+
* so we leave the statement bare to avoid incorrect output.
1405+
*/
1406+
private inferReturnNextVar(context: PLpgSQLDeparserContext): string | undefined {
1407+
// Without return-type context we cannot safely infer
1408+
if (!context.returnInfo) {
1409+
return undefined;
1410+
}
1411+
1412+
// TABLE / OUT-param functions legitimately use bare RETURN NEXT
1413+
if (context.returnInfo.kind === 'out_params') {
1414+
return undefined;
1415+
}
1416+
1417+
// Only attempt recovery for SETOF scalar functions
1418+
if (context.returnInfo.kind !== 'setof') {
1419+
return undefined;
1420+
}
1421+
1422+
// 1. Inside a FOR loop → use the loop variable
1423+
if (context.enclosingForVarName) {
1424+
return context.enclosingForVarName;
1425+
}
1426+
1427+
// 2. Outside a FOR loop → look for a single user-declared var
1428+
if (context.datums) {
1429+
const candidates = context.datums.filter((d) => {
1430+
if ('PLpgSQL_var' in d) {
1431+
const v = d.PLpgSQL_var;
1432+
// Skip implicit 'found' and function parameters (no lineno)
1433+
return v.refname !== 'found' && v.lineno !== undefined;
1434+
}
1435+
return false;
1436+
});
1437+
if (candidates.length === 1 && 'PLpgSQL_var' in candidates[0]) {
1438+
return candidates[0].PLpgSQL_var.refname;
1439+
}
1440+
}
1441+
1442+
return undefined;
1443+
}
1444+
13821445
/**
13831446
* Deparse a RETURN QUERY statement
13841447
*/
@@ -1718,7 +1781,7 @@ export class PLpgSQLDeparser {
17181781
parts.push(`${forClause} ${kw('LOOP')}`);
17191782

17201783
if (fors.body) {
1721-
const bodyContext = { ...context, indentLevel: context.indentLevel + 1 };
1784+
const bodyContext = { ...context, indentLevel: context.indentLevel + 1, enclosingForVarName: varName };
17221785
for (const stmt of fors.body) {
17231786
const stmtStr = this.deparseStmt(stmt, bodyContext);
17241787
parts.push(this.indent(stmtStr + ';', bodyContext.indentLevel));

‎packages/plpgsql-deparser/test-utils/index.ts‎

Lines changed: 40 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import { parsePlPgSQL, parsePlPgSQLSync } from '@libpg-query/parser';
2-
import { deparseSync, PLpgSQLParseResult } from '../src';
2+
import { deparseSync, PLpgSQLParseResult, ReturnInfo } from '../src';
33
import { readFileSync, readdirSync, existsSync } from 'fs';
44
import * as path from 'path';
55
import { diff } from 'jest-diff';
@@ -328,6 +328,43 @@ function reconstructSql(originalSql: string, newBody: string): string {
328328
return parts.prefix + newBody + parts.suffix;
329329
}
330330

331+
/**
332+
* Extract ReturnInfo from a CREATE FUNCTION SQL string by inspecting the
333+
* RETURNS clause and OUT/INOUT parameters. Used by the round-trip test to
334+
* give the deparser enough context to recover dropped retvarno fields.
335+
*
336+
* We intentionally only extract for SETOF and OUT-param functions — these
337+
* are the cases where RETURN NEXT inference needs context. For other
338+
* return types (scalar, void, trigger) we return undefined to avoid
339+
* changing existing RETURN statement behaviour (e.g. bare RETURN in scalar
340+
* functions would become RETURN NULL if we provided returnInfo).
341+
*/
342+
export function extractReturnInfo(sql: string): ReturnInfo | undefined {
343+
// Normalise to a single line for simpler matching
344+
const norm = sql.replace(/\s+/g, ' ').trim();
345+
346+
// Strip the body ($$...$$) so we only inspect the signature
347+
const sig = norm.replace(/\$\$.*\$\$/s, '');
348+
349+
// OUT / INOUT parameters → out_params (prevents false RETURN NEXT inference)
350+
if (/\b(OUT|INOUT)\s+/i.test(sig)) {
351+
return { kind: 'out_params' };
352+
}
353+
354+
// RETURNS TABLE(...) → out_params
355+
if (/RETURNS\s+TABLE\s*\(/i.test(sig)) {
356+
return { kind: 'out_params' };
357+
}
358+
359+
// RETURNS SETOF <type> → needed for RETURN NEXT variable recovery
360+
if (/RETURNS\s+SETOF\b/i.test(sig)) {
361+
return { kind: 'setof' };
362+
}
363+
364+
// All other types: don't provide returnInfo to preserve existing behaviour
365+
return undefined;
366+
}
367+
331368
export class PLpgSQLTestUtils {
332369
protected printErrorMessage(sql: string, position: number) {
333370
const lineNumber = sql.slice(0, position).match(/\n/g)?.length || 0;
@@ -368,7 +405,8 @@ export class PLpgSQLTestUtils {
368405
throw createParseError('PARSE_FAILED', testName, sql);
369406
}
370407

371-
const deparsedBody = deparseSync(originalAst);
408+
const returnInfo = extractReturnInfo(sql);
409+
const deparsedBody = deparseSync(originalAst, undefined, returnInfo);
372410

373411
if (!deparsedBody || deparsedBody.trim().length === 0) {
374412
throw createParseError('DEPARSE_FAILED', testName, sql, deparsedBody);

0 commit comments

Comments
 (0)