diff --git a/__fixtures__/plpgsql-generated/generated.json b/__fixtures__/plpgsql-generated/generated.json index 13f6bf31..909e5936 100644 --- a/__fixtures__/plpgsql-generated/generated.json +++ b/__fixtures__/plpgsql-generated/generated.json @@ -136,6 +136,17 @@ "plpgsql_deparser_fixes-35.sql": "-- Test 35: CALL statement\nCREATE FUNCTION test_call_statement() RETURNS void\nLANGUAGE plpgsql AS $$\nBEGIN\n CALL my_procedure(1, 'hello');\n RETURN;\nEND$$", "plpgsql_deparser_fixes-36.sql": "-- =============================================================================\n-- Edge Case Tests: Real-World Patterns\n-- =============================================================================\n\n-- Test 36: Permission bitnum trigger pattern (the function that exposed the END; bug)\nCREATE FUNCTION test_permission_bitnum_trigger() RETURNS trigger\nLANGUAGE plpgsql AS $$\nDECLARE\n bitlen int;\n v_len int;\nBEGIN\n v_len := 32;\n BEGIN\n bitlen := bit_length(NEW.bitstr);\n EXCEPTION\n WHEN others THEN\n bitlen := 0;\n END;\n IF bitlen = 0 THEN\n NEW.bitstr := lpad('', v_len, '0');\n END IF;\n RETURN NEW;\nEND$$", "plpgsql_deparser_fixes-37.sql": "-- Test 37: Multi-step sign-in pattern (deeply nested IF chains)\nCREATE FUNCTION test_signin_pattern(v_email text) RETURNS record\nLANGUAGE plpgsql AS $$\nDECLARE\n v_user record;\n v_secret record;\nBEGIN\n SELECT * INTO v_user FROM users WHERE email = v_email;\n IF NOT FOUND THEN\n RAISE EXCEPTION 'USER_NOT_FOUND';\n END IF;\n SELECT * INTO v_secret FROM secrets WHERE user_id = v_user.id;\n IF NOT FOUND THEN\n RAISE EXCEPTION 'NO_CREDENTIALS';\n END IF;\n IF v_secret.locked_at IS NOT NULL THEN\n RAISE EXCEPTION 'ACCOUNT_LOCKED';\n END IF;\n RETURN v_user;\nEND$$", + "plpgsql_deparser_fixes-38.sql": "-- Test 38: COALESCE inside function call arguments\n-- Exercises CoalesceExpr as a FuncCall arg (e.g. jsonb_array_elements(COALESCE(v_config, '[]')))\nCREATE FUNCTION test_coalesce_in_func_args(v_config jsonb) RETURNS SETOF jsonb\nLANGUAGE plpgsql AS $$\nBEGIN\n RETURN QUERY SELECT jsonb_array_elements(COALESCE(v_config, '[]'::jsonb));\nEND$$", + "plpgsql_deparser_fixes-39.sql": "-- Test 39: Set-returning function call in FROM clause (RangeFunction)\n-- Exercises FuncCall inside RangeFunction for set-returning functions\nCREATE FUNCTION test_func_in_from_clause(v_data jsonb) RETURNS TABLE(key text, value text)\nLANGUAGE plpgsql AS $$\nBEGIN\n RETURN QUERY SELECT elem->>'key', elem->>'value'\n FROM jsonb_array_elements(v_data) AS elem;\nEND$$", + "plpgsql_deparser_fixes-40.sql": "-- Test 40: COALESCE in FROM clause with set-returning function (combined pattern)\nCREATE FUNCTION test_coalesce_in_from_srf(v_config jsonb) RETURNS TABLE(entry jsonb)\nLANGUAGE plpgsql AS $$\nBEGIN\n RETURN QUERY SELECT elem\n FROM jsonb_array_elements(COALESCE(v_config, '[]'::jsonb)) AS elem;\nEND$$", + "plpgsql_deparser_fixes-41.sql": "-- Test 41: FOR loop with set-returning function and COALESCE in SELECT INTO\nCREATE FUNCTION test_for_loop_srf_coalesce(v_items jsonb) RETURNS integer\nLANGUAGE plpgsql AS $$\nDECLARE\n v_count integer;\nBEGIN\n SELECT count(*) INTO v_count\n FROM jsonb_array_elements(COALESCE(v_items, '[]'::jsonb)) AS elem;\n RETURN v_count;\nEND$$", + "plpgsql_deparser_fixes-42.sql": "-- Test 42: Multiple set-returning functions in FROM clause\nCREATE FUNCTION test_multiple_srf_from(v_keys text[], v_values text[]) RETURNS TABLE(k text, v text)\nLANGUAGE plpgsql AS $$\nBEGIN\n RETURN QUERY SELECT unnest_k, unnest_v\n FROM unnest(v_keys) AS unnest_k, unnest(v_values) AS unnest_v;\nEND$$", + "plpgsql_deparser_fixes-43.sql": "-- Test 43: COALESCE in assignment with function call\nCREATE FUNCTION test_coalesce_assign_func() RETURNS trigger\nLANGUAGE plpgsql AS $$\nBEGIN\n NEW.updated_at := COALESCE(NEW.updated_at, now());\n NEW.name := COALESCE(NEW.name, 'default_' || NEW.id::text);\n RETURN NEW;\nEND$$", + "plpgsql_deparser_fixes-44.sql": "-- Test 44: Nested COALESCE expressions\nCREATE FUNCTION test_nested_coalesce(a text, b text, c text) RETURNS text\nLANGUAGE plpgsql AS $$\nBEGIN\n RETURN COALESCE(a, COALESCE(b, COALESCE(c, 'fallback')));\nEND$$", + "plpgsql_deparser_fixes-45.sql": "-- Test 45: SELECT INTO with COALESCE and function call in FROM\nCREATE FUNCTION test_select_into_from_srf(v_data jsonb) RETURNS jsonb\nLANGUAGE plpgsql AS $$\nDECLARE\n v_first jsonb;\nBEGIN\n SELECT elem INTO v_first\n FROM jsonb_array_elements(v_data) AS elem\n LIMIT 1;\n RETURN v_first;\nEND$$", + "plpgsql_deparser_fixes-46.sql": "-- Test 46: FOR loop with jsonb_array_elements and WHERE filter\nCREATE FUNCTION test_for_srf_with_filter(v_config jsonb, v_key text) RETURNS jsonb\nLANGUAGE plpgsql AS $$\nDECLARE\n v_entry jsonb;\nBEGIN\n FOR v_entry IN\n SELECT elem FROM jsonb_array_elements(v_config) AS elem\n WHERE elem->>'key' = v_key\n LOOP\n RETURN v_entry;\n END LOOP;\n RETURN NULL;\nEND$$", + "plpgsql_deparser_fixes-47.sql": "-- Test 47: RETURN NEXT with FOR loop variable (retvarno recovery)\n-- Exercises the case where libpg-query drops retvarno from PLpgSQL_stmt_return_next\nCREATE FUNCTION test_return_next_for_var(v_data jsonb) RETURNS SETOF jsonb\nLANGUAGE plpgsql AS $$\nDECLARE\n v_entry jsonb;\nBEGIN\n FOR v_entry IN SELECT jsonb_array_elements(v_data)\n LOOP\n RETURN NEXT v_entry;\n END LOOP;\n RETURN;\nEND$$", + "plpgsql_deparser_fixes-48.sql": "-- Test 48: RETURN NEXT with declared variable (no FOR loop, single candidate)\nCREATE FUNCTION test_return_next_declared_var() RETURNS SETOF int\nLANGUAGE plpgsql AS $$\nDECLARE\n v_val int := 42;\nBEGIN\n RETURN NEXT v_val;\nEND$$", "plpgsql_control-1.sql": "--\n-- Tests for PL/pgSQL control structures\n--\n\n-- integer FOR loop\n\ndo $$\nbegin\n -- basic case\n for i in 1..3 loop\n raise notice '1..3: i = %', i;\n end loop;\n -- with BY, end matches exactly\n for i in 1..10 by 3 loop\n raise notice '1..10 by 3: i = %', i;\n end loop;\n -- with BY, end does not match\n for i in 1..11 by 3 loop\n raise notice '1..11 by 3: i = %', i;\n end loop;\n -- zero iterations\n for i in 1..0 by 3 loop\n raise notice '1..0 by 3: i = %', i;\n end loop;\n -- REVERSE\n for i in reverse 10..0 by 3 loop\n raise notice 'reverse 10..0 by 3: i = %', i;\n end loop;\n -- potential overflow\n for i in 2147483620..2147483647 by 10 loop\n raise notice '2147483620..2147483647 by 10: i = %', i;\n end loop;\n -- potential overflow, reverse direction\n for i in reverse -2147483620..-2147483647 by 10 loop\n raise notice 'reverse -2147483620..-2147483647 by 10: i = %', i;\n end loop;\nend$$", "plpgsql_control-2.sql": "-- BY can't be zero or negative\ndo $$\nbegin\n for i in 1..3 by 0 loop\n raise notice '1..3 by 0: i = %', i;\n end loop;\nend$$", "plpgsql_control-3.sql": "do $$\nbegin\n for i in 1..3 by -1 loop\n raise notice '1..3 by -1: i = %', i;\n end loop;\nend$$", diff --git a/__fixtures__/plpgsql/plpgsql_deparser_fixes.sql b/__fixtures__/plpgsql/plpgsql_deparser_fixes.sql index 84bf64df..478656ae 100644 --- a/__fixtures__/plpgsql/plpgsql_deparser_fixes.sql +++ b/__fixtures__/plpgsql/plpgsql_deparser_fixes.sql @@ -505,3 +505,113 @@ BEGIN END IF; RETURN v_user; END$$; + +-- Test 38: COALESCE inside function call arguments +-- Exercises CoalesceExpr as a FuncCall arg (e.g. jsonb_array_elements(COALESCE(v_config, '[]'))) +CREATE FUNCTION test_coalesce_in_func_args(v_config jsonb) RETURNS SETOF jsonb +LANGUAGE plpgsql AS $$ +BEGIN + RETURN QUERY SELECT jsonb_array_elements(COALESCE(v_config, '[]'::jsonb)); +END$$; + +-- Test 39: Set-returning function call in FROM clause (RangeFunction) +-- Exercises FuncCall inside RangeFunction for set-returning functions +CREATE FUNCTION test_func_in_from_clause(v_data jsonb) RETURNS TABLE(key text, value text) +LANGUAGE plpgsql AS $$ +BEGIN + RETURN QUERY SELECT elem->>'key', elem->>'value' + FROM jsonb_array_elements(v_data) AS elem; +END$$; + +-- Test 40: COALESCE in FROM clause with set-returning function (combined pattern) +CREATE FUNCTION test_coalesce_in_from_srf(v_config jsonb) RETURNS TABLE(entry jsonb) +LANGUAGE plpgsql AS $$ +BEGIN + RETURN QUERY SELECT elem + FROM jsonb_array_elements(COALESCE(v_config, '[]'::jsonb)) AS elem; +END$$; + +-- Test 41: FOR loop with set-returning function and COALESCE in SELECT INTO +CREATE FUNCTION test_for_loop_srf_coalesce(v_items jsonb) RETURNS integer +LANGUAGE plpgsql AS $$ +DECLARE + v_count integer; +BEGIN + SELECT count(*) INTO v_count + FROM jsonb_array_elements(COALESCE(v_items, '[]'::jsonb)) AS elem; + RETURN v_count; +END$$; + +-- Test 42: Multiple set-returning functions in FROM clause +CREATE FUNCTION test_multiple_srf_from(v_keys text[], v_values text[]) RETURNS TABLE(k text, v text) +LANGUAGE plpgsql AS $$ +BEGIN + RETURN QUERY SELECT unnest_k, unnest_v + FROM unnest(v_keys) AS unnest_k, unnest(v_values) AS unnest_v; +END$$; + +-- Test 43: COALESCE in assignment with function call +CREATE FUNCTION test_coalesce_assign_func() RETURNS trigger +LANGUAGE plpgsql AS $$ +BEGIN + NEW.updated_at := COALESCE(NEW.updated_at, now()); + NEW.name := COALESCE(NEW.name, 'default_' || NEW.id::text); + RETURN NEW; +END$$; + +-- Test 44: Nested COALESCE expressions +CREATE FUNCTION test_nested_coalesce(a text, b text, c text) RETURNS text +LANGUAGE plpgsql AS $$ +BEGIN + RETURN COALESCE(a, COALESCE(b, COALESCE(c, 'fallback'))); +END$$; + +-- Test 45: SELECT INTO with COALESCE and function call in FROM +CREATE FUNCTION test_select_into_from_srf(v_data jsonb) RETURNS jsonb +LANGUAGE plpgsql AS $$ +DECLARE + v_first jsonb; +BEGIN + SELECT elem INTO v_first + FROM jsonb_array_elements(v_data) AS elem + LIMIT 1; + RETURN v_first; +END$$; + +-- Test 46: FOR loop with jsonb_array_elements and WHERE filter +CREATE FUNCTION test_for_srf_with_filter(v_config jsonb, v_key text) RETURNS jsonb +LANGUAGE plpgsql AS $$ +DECLARE + v_entry jsonb; +BEGIN + FOR v_entry IN + SELECT elem FROM jsonb_array_elements(v_config) AS elem + WHERE elem->>'key' = v_key + LOOP + RETURN v_entry; + END LOOP; + RETURN NULL; +END$$; + +-- Test 47: RETURN NEXT with FOR loop variable (retvarno recovery) +-- Exercises the case where libpg-query drops retvarno from PLpgSQL_stmt_return_next +CREATE FUNCTION test_return_next_for_var(v_data jsonb) RETURNS SETOF jsonb +LANGUAGE plpgsql AS $$ +DECLARE + v_entry jsonb; +BEGIN + FOR v_entry IN SELECT jsonb_array_elements(v_data) + LOOP + RETURN NEXT v_entry; + END LOOP; + RETURN; +END$$; + +-- Test 48: RETURN NEXT with declared variable (no FOR loop, single candidate) +CREATE FUNCTION test_return_next_declared_var() RETURNS SETOF int +LANGUAGE plpgsql AS $$ +DECLARE + v_val int := 42; +BEGIN + RETURN NEXT v_val; +END$$; diff --git a/packages/plpgsql-deparser/__tests__/__snapshots__/deparser-fixes.test.ts.snap b/packages/plpgsql-deparser/__tests__/__snapshots__/deparser-fixes.test.ts.snap index ac18db54..55c79ea0 100644 --- a/packages/plpgsql-deparser/__tests__/__snapshots__/deparser-fixes.test.ts.snap +++ b/packages/plpgsql-deparser/__tests__/__snapshots__/deparser-fixes.test.ts.snap @@ -1,5 +1,72 @@ // Jest Snapshot v1, https://jestjs.io/docs/snapshot-testing +exports[`plpgsql-deparser bug fixes CoalesceExpr and RangeFunction patterns should handle COALESCE in FROM clause with set-returning function 1`] = ` +"BEGIN + RETURN QUERY SELECT elem + FROM jsonb_array_elements(COALESCE(v_config, '[]'::jsonb)) AS elem; + RETURN; +END" +`; + +exports[`plpgsql-deparser bug fixes CoalesceExpr and RangeFunction patterns should handle COALESCE inside function call arguments 1`] = ` +"BEGIN + RETURN QUERY SELECT jsonb_array_elements(COALESCE(v_config, '[]'::jsonb)); + RETURN; +END" +`; + +exports[`plpgsql-deparser bug fixes CoalesceExpr and RangeFunction patterns should handle COALESCE with function call in assignment 1`] = ` +"BEGIN + NEW.updated_at := COALESCE(NEW.updated_at, now()); + NEW.name := COALESCE(NEW.name, 'default_' || NEW.id::text); + RETURN NEW; +END" +`; + +exports[`plpgsql-deparser bug fixes CoalesceExpr and RangeFunction patterns should handle FOR loop with SRF and WHERE filter 1`] = ` +"DECLARE + v_entry jsonb; +BEGIN + FOR v_entry IN SELECT elem FROM jsonb_array_elements(v_config) AS elem + WHERE elem->>'key' = v_key LOOP + RETURN v_entry; + END LOOP; + RETURN NULL; +END" +`; + +exports[`plpgsql-deparser bug fixes CoalesceExpr and RangeFunction patterns should handle SELECT INTO from set-returning function in FROM clause 1`] = ` +"DECLARE + v_first jsonb; +BEGIN + SELECT elem INTO v_first FROM jsonb_array_elements(v_data) AS elem + LIMIT 1; + RETURN v_first; +END" +`; + +exports[`plpgsql-deparser bug fixes CoalesceExpr and RangeFunction patterns should handle multiple SRFs in FROM clause 1`] = ` +"BEGIN + RETURN QUERY SELECT unnest_k, unnest_v + FROM unnest(v_keys) AS unnest_k, unnest(v_values) AS unnest_v; + RETURN; +END" +`; + +exports[`plpgsql-deparser bug fixes CoalesceExpr and RangeFunction patterns should handle nested COALESCE expressions 1`] = ` +"BEGIN + RETURN COALESCE(a, COALESCE(b, COALESCE(c, 'fallback'))); +END" +`; + +exports[`plpgsql-deparser bug fixes CoalesceExpr and RangeFunction patterns should handle set-returning function in FROM clause (RangeFunction) 1`] = ` +"BEGIN + RETURN QUERY SELECT elem->>'key', elem->>'value' + FROM jsonb_array_elements(v_data) AS elem; + RETURN; +END" +`; + exports[`plpgsql-deparser bug fixes INTO clause depth-aware scanner should handle INTO STRICT 1`] = ` "DECLARE v_id integer; @@ -116,6 +183,48 @@ exports[`plpgsql-deparser bug fixes PERFORM SELECT fix should strip SELECT keywo END" `; +exports[`plpgsql-deparser bug fixes RETURN NEXT retvarno recovery should leave RETURN NEXT bare for OUT-param functions 1`] = ` +"BEGIN + FOR i IN 1..5 LOOP + x := i; + y := 'item_' || i::text; + RETURN NEXT; + END LOOP; + RETURN; +END" +`; + +exports[`plpgsql-deparser bug fixes RETURN NEXT retvarno recovery should leave RETURN NEXT bare when returnInfo is not provided 1`] = ` +"DECLARE + v_entry jsonb; +BEGIN + FOR v_entry IN SELECT jsonb_array_elements(v_data) LOOP + RETURN NEXT; + END LOOP; + RETURN; +END" +`; + +exports[`plpgsql-deparser bug fixes RETURN NEXT retvarno recovery should recover FOR loop variable in RETURN NEXT for SETOF functions 1`] = ` +"DECLARE + v_entry jsonb; +BEGIN + FOR v_entry IN SELECT jsonb_array_elements(v_data) LOOP + RETURN NEXT v_entry; + END LOOP; + RETURN; +END" +`; + +exports[`plpgsql-deparser bug fixes RETURN NEXT retvarno recovery should recover single declared variable in RETURN NEXT for SETOF functions 1`] = ` +"DECLARE + v_val int := 42; +BEGIN + RETURN NEXT v_val; + RETURN; +END" +`; + exports[`plpgsql-deparser bug fixes Record field qualification (recfield) should handle OLD and NEW record references 1`] = ` "BEGIN IF OLD.status <> NEW.status THEN diff --git a/packages/plpgsql-deparser/__tests__/deparser-fixes.test.ts b/packages/plpgsql-deparser/__tests__/deparser-fixes.test.ts index 4b068ff5..b63dfc37 100644 --- a/packages/plpgsql-deparser/__tests__/deparser-fixes.test.ts +++ b/packages/plpgsql-deparser/__tests__/deparser-fixes.test.ts @@ -872,4 +872,231 @@ END$$`; expect(deparsed).toMatchSnapshot(); }); }); + + describe('CoalesceExpr and RangeFunction patterns', () => { + it('should handle COALESCE inside function call arguments', async () => { + const sql = `CREATE FUNCTION test_coalesce_in_func_args(v_config jsonb) RETURNS SETOF jsonb +LANGUAGE plpgsql AS $$ +BEGIN + RETURN QUERY SELECT jsonb_array_elements(COALESCE(v_config, '[]'::jsonb)); +END$$`; + + await testUtils.expectAstMatch('COALESCE in FuncCall args', sql); + + const parsed = parsePlPgSQLSync(sql) as unknown as PLpgSQLParseResult; + const deparsed = deparseSync(parsed); + expect(deparsed).toMatchSnapshot(); + expect(deparsed).toContain('COALESCE'); + expect(deparsed).toContain('jsonb_array_elements'); + }); + + it('should handle set-returning function in FROM clause (RangeFunction)', async () => { + const sql = `CREATE FUNCTION test_func_in_from(v_data jsonb) RETURNS TABLE(key text, value text) +LANGUAGE plpgsql AS $$ +BEGIN + RETURN QUERY SELECT elem->>'key', elem->>'value' + FROM jsonb_array_elements(v_data) AS elem; +END$$`; + + await testUtils.expectAstMatch('FuncCall in FROM clause', sql); + + const parsed = parsePlPgSQLSync(sql) as unknown as PLpgSQLParseResult; + const deparsed = deparseSync(parsed); + expect(deparsed).toMatchSnapshot(); + expect(deparsed).toContain('jsonb_array_elements'); + expect(deparsed).toContain('FROM'); + }); + + it('should handle COALESCE in FROM clause with set-returning function', async () => { + const sql = `CREATE FUNCTION test_coalesce_in_from_srf(v_config jsonb) RETURNS TABLE(entry jsonb) +LANGUAGE plpgsql AS $$ +BEGIN + RETURN QUERY SELECT elem + FROM jsonb_array_elements(COALESCE(v_config, '[]'::jsonb)) AS elem; +END$$`; + + await testUtils.expectAstMatch('COALESCE in FROM SRF', sql); + + const parsed = parsePlPgSQLSync(sql) as unknown as PLpgSQLParseResult; + const deparsed = deparseSync(parsed); + expect(deparsed).toMatchSnapshot(); + expect(deparsed).toContain('COALESCE'); + expect(deparsed).toContain('FROM'); + expect(deparsed).toContain('jsonb_array_elements'); + }); + + it('should handle SELECT INTO from set-returning function in FROM clause', async () => { + const sql = `CREATE FUNCTION test_select_into_from_srf(v_data jsonb) RETURNS jsonb +LANGUAGE plpgsql AS $$ +DECLARE + v_first jsonb; +BEGIN + SELECT elem INTO v_first + FROM jsonb_array_elements(v_data) AS elem + LIMIT 1; + RETURN v_first; +END$$`; + + await testUtils.expectAstMatch('SELECT INTO from SRF in FROM', sql); + + const parsed = parsePlPgSQLSync(sql) as unknown as PLpgSQLParseResult; + const deparsed = deparseSync(parsed); + expect(deparsed).toMatchSnapshot(); + expect(deparsed).toContain('INTO'); + expect(deparsed).toContain('jsonb_array_elements'); + expect(deparsed).toContain('FROM'); + }); + + it('should handle FOR loop with SRF and WHERE filter', async () => { + const sql = `CREATE FUNCTION test_for_srf_filter(v_config jsonb, v_key text) RETURNS jsonb +LANGUAGE plpgsql AS $$ +DECLARE + v_entry jsonb; +BEGIN + FOR v_entry IN + SELECT elem FROM jsonb_array_elements(v_config) AS elem + WHERE elem->>'key' = v_key + LOOP + RETURN v_entry; + END LOOP; + RETURN NULL; +END$$`; + + await testUtils.expectAstMatch('FOR loop SRF with WHERE', sql); + + const parsed = parsePlPgSQLSync(sql) as unknown as PLpgSQLParseResult; + const deparsed = deparseSync(parsed); + expect(deparsed).toMatchSnapshot(); + expect(deparsed).toContain('jsonb_array_elements'); + expect(deparsed).toContain('WHERE'); + }); + + it('should handle nested COALESCE expressions', async () => { + const sql = `CREATE FUNCTION test_nested_coalesce(a text, b text, c text) RETURNS text +LANGUAGE plpgsql AS $$ +BEGIN + RETURN COALESCE(a, COALESCE(b, COALESCE(c, 'fallback'))); +END$$`; + + await testUtils.expectAstMatch('nested COALESCE', sql); + + const parsed = parsePlPgSQLSync(sql) as unknown as PLpgSQLParseResult; + const deparsed = deparseSync(parsed); + expect(deparsed).toMatchSnapshot(); + expect(deparsed).toContain('COALESCE'); + }); + + it('should handle multiple SRFs in FROM clause', async () => { + const sql = `CREATE FUNCTION test_multiple_srf(v_keys text[], v_values text[]) RETURNS TABLE(k text, v text) +LANGUAGE plpgsql AS $$ +BEGIN + RETURN QUERY SELECT unnest_k, unnest_v + FROM unnest(v_keys) AS unnest_k, unnest(v_values) AS unnest_v; +END$$`; + + await testUtils.expectAstMatch('multiple SRFs in FROM', sql); + + const parsed = parsePlPgSQLSync(sql) as unknown as PLpgSQLParseResult; + const deparsed = deparseSync(parsed); + expect(deparsed).toMatchSnapshot(); + expect(deparsed).toContain('unnest'); + expect(deparsed).toContain('FROM'); + }); + + it('should handle COALESCE with function call in assignment', async () => { + const sql = `CREATE FUNCTION test_coalesce_assign() RETURNS trigger +LANGUAGE plpgsql AS $$ +BEGIN + NEW.updated_at := COALESCE(NEW.updated_at, now()); + NEW.name := COALESCE(NEW.name, 'default_' || NEW.id::text); + RETURN NEW; +END$$`; + + await testUtils.expectAstMatch('COALESCE assign with func', sql); + + const parsed = parsePlPgSQLSync(sql) as unknown as PLpgSQLParseResult; + const deparsed = deparseSync(parsed); + expect(deparsed).toMatchSnapshot(); + expect(deparsed).toContain('COALESCE'); + expect(deparsed).toContain('now()'); + }); + }); + + describe('RETURN NEXT retvarno recovery', () => { + it('should recover FOR loop variable in RETURN NEXT for SETOF functions', async () => { + const sql = `CREATE FUNCTION test_return_next_for_var(v_data jsonb) RETURNS SETOF jsonb +LANGUAGE plpgsql AS $$ +DECLARE + v_entry jsonb; +BEGIN + FOR v_entry IN SELECT jsonb_array_elements(v_data) + LOOP + RETURN NEXT v_entry; + END LOOP; + RETURN; +END$$`; + + await testUtils.expectAstMatch('RETURN NEXT FOR loop var', sql); + const parsed = parsePlPgSQLSync(sql) as unknown as PLpgSQLParseResult; + const deparsed = deparseSync(parsed, undefined, { kind: 'setof' }); + expect(deparsed).toMatchSnapshot(); + expect(deparsed).toContain('RETURN NEXT v_entry'); + }); + + it('should recover single declared variable in RETURN NEXT for SETOF functions', async () => { + const sql = `CREATE FUNCTION test_return_next_declared_var() RETURNS SETOF int +LANGUAGE plpgsql AS $$ +DECLARE + v_val int := 42; +BEGIN + RETURN NEXT v_val; +END$$`; + + await testUtils.expectAstMatch('RETURN NEXT declared var', sql); + const parsed = parsePlPgSQLSync(sql) as unknown as PLpgSQLParseResult; + const deparsed = deparseSync(parsed, undefined, { kind: 'setof' }); + expect(deparsed).toMatchSnapshot(); + expect(deparsed).toContain('RETURN NEXT v_val'); + }); + + it('should leave RETURN NEXT bare for OUT-param functions', async () => { + const sql = `CREATE FUNCTION test_return_next_out(OUT x integer, OUT y text) RETURNS SETOF record +LANGUAGE plpgsql AS $$ +BEGIN + FOR i IN 1..5 LOOP + x := i; + y := 'item_' || i::text; + RETURN NEXT; + END LOOP; + RETURN; +END$$`; + + await testUtils.expectAstMatch('RETURN NEXT OUT params', sql); + const parsed = parsePlPgSQLSync(sql) as unknown as PLpgSQLParseResult; + const deparsed = deparseSync(parsed, undefined, { kind: 'out_params' }); + expect(deparsed).toMatchSnapshot(); + expect(deparsed).toContain('RETURN NEXT;'); + expect(deparsed).not.toMatch(/RETURN NEXT\s+\w/); + }); + + it('should leave RETURN NEXT bare when returnInfo is not provided', async () => { + const sql = `CREATE FUNCTION test_return_next_for_var(v_data jsonb) RETURNS SETOF jsonb +LANGUAGE plpgsql AS $$ +DECLARE + v_entry jsonb; +BEGIN + FOR v_entry IN SELECT jsonb_array_elements(v_data) + LOOP + RETURN NEXT v_entry; + END LOOP; + RETURN; +END$$`; + + const parsed = parsePlPgSQLSync(sql) as unknown as PLpgSQLParseResult; + // Without returnInfo, inference should NOT happen (safety) + const deparsed = deparseSync(parsed); + expect(deparsed).toMatchSnapshot(); + expect(deparsed).toContain('RETURN NEXT;'); + }); + }); }); diff --git a/packages/plpgsql-deparser/src/plpgsql-deparser.ts b/packages/plpgsql-deparser/src/plpgsql-deparser.ts index 8a88fbbf..ada53fbc 100644 --- a/packages/plpgsql-deparser/src/plpgsql-deparser.ts +++ b/packages/plpgsql-deparser/src/plpgsql-deparser.ts @@ -86,6 +86,8 @@ export interface PLpgSQLDeparserContext { loopVarLinenos?: Set; /** Map of block lineno to the set of datum indices that belong to that block */ blockDatumMap?: Map>; + /** Name of the enclosing FOR loop variable (used to recover RETURN NEXT when retvarno is missing) */ + enclosingForVarName?: string; } /** @@ -1179,7 +1181,7 @@ export class PLpgSQLDeparser { } if (fori.body) { - const bodyContext = { ...context, indentLevel: context.indentLevel + 1 }; + const bodyContext = { ...context, indentLevel: context.indentLevel + 1, enclosingForVarName: varName }; for (const stmt of fori.body) { const stmtStr = this.deparseStmt(stmt, bodyContext); parts.push(this.indent(stmtStr + ';', bodyContext.indentLevel)); @@ -1210,7 +1212,7 @@ export class PLpgSQLDeparser { parts.push(`${kw('FOR')} ${varName} ${kw('IN')} ${query} ${kw('LOOP')}`); if (fors.body) { - const bodyContext = { ...context, indentLevel: context.indentLevel + 1 }; + const bodyContext = { ...context, indentLevel: context.indentLevel + 1, enclosingForVarName: varName }; for (const stmt of fors.body) { const stmtStr = this.deparseStmt(stmt, bodyContext); parts.push(this.indent(stmtStr + ';', bodyContext.indentLevel)); @@ -1246,7 +1248,7 @@ export class PLpgSQLDeparser { parts.push(`${forClause} ${kw('LOOP')}`); if (forc.body) { - const bodyContext = { ...context, indentLevel: context.indentLevel + 1 }; + const bodyContext = { ...context, indentLevel: context.indentLevel + 1, enclosingForVarName: varName }; for (const stmt of forc.body) { const stmtStr = this.deparseStmt(stmt, bodyContext); parts.push(this.indent(stmtStr + ';', bodyContext.indentLevel)); @@ -1283,7 +1285,7 @@ export class PLpgSQLDeparser { parts.push(`${kw('FOREACH')} ${varName}${sliceClause} ${kw('IN ARRAY')} ${expr} ${kw('LOOP')}`); if (foreach.body) { - const bodyContext = { ...context, indentLevel: context.indentLevel + 1 }; + const bodyContext = { ...context, indentLevel: context.indentLevel + 1, enclosingForVarName: varName }; for (const stmt of foreach.body) { const stmtStr = this.deparseStmt(stmt, bodyContext); parts.push(this.indent(stmtStr + ';', bodyContext.indentLevel)); @@ -1363,22 +1365,83 @@ export class PLpgSQLDeparser { /** * Deparse a RETURN NEXT statement + * + * libpg-query does not serialize `retvarno` for PLpgSQL_stmt_return_next, + * so when a function has `RETURN NEXT variable`, the variable reference is + * lost. We recover it by: + * 1. Using the enclosing FOR loop variable (most common pattern) + * 2. Scanning datums for the sole user-declared variable (non-loop case) + * Bare `RETURN NEXT` remains valid for TABLE / OUT-param functions. */ private deparseReturnNext(ret: PLpgSQL_stmt_return_next, context: PLpgSQLDeparserContext): string { const kw = this.keyword; - + if (ret.expr) { return `${kw('RETURN NEXT')} ${this.deparseExpr(ret.expr)}`; } - + if (ret.retvarno !== undefined && ret.retvarno >= 0) { const varName = this.getVarName(ret.retvarno, context); return `${kw('RETURN NEXT')} ${varName}`; } - + + // retvarno is missing (libpg-query serialization gap). Try to recover. + const inferred = this.inferReturnNextVar(context); + if (inferred) { + return `${kw('RETURN NEXT')} ${inferred}`; + } + return kw('RETURN NEXT'); } + /** + * Attempt to infer the variable for a bare RETURN NEXT when retvarno is + * not available. Returns the variable name or undefined. + * + * We only infer when `returnInfo` tells us the function returns SETOF + * scalar. Without returnInfo we cannot distinguish a legitimate bare + * `RETURN NEXT` (OUT-param / TABLE function) from a dropped retvarno, + * so we leave the statement bare to avoid incorrect output. + */ + private inferReturnNextVar(context: PLpgSQLDeparserContext): string | undefined { + // Without return-type context we cannot safely infer + if (!context.returnInfo) { + return undefined; + } + + // TABLE / OUT-param functions legitimately use bare RETURN NEXT + if (context.returnInfo.kind === 'out_params') { + return undefined; + } + + // Only attempt recovery for SETOF scalar functions + if (context.returnInfo.kind !== 'setof') { + return undefined; + } + + // 1. Inside a FOR loop → use the loop variable + if (context.enclosingForVarName) { + return context.enclosingForVarName; + } + + // 2. Outside a FOR loop → look for a single user-declared var + if (context.datums) { + const candidates = context.datums.filter((d) => { + if ('PLpgSQL_var' in d) { + const v = d.PLpgSQL_var; + // Skip implicit 'found' and function parameters (no lineno) + return v.refname !== 'found' && v.lineno !== undefined; + } + return false; + }); + if (candidates.length === 1 && 'PLpgSQL_var' in candidates[0]) { + return candidates[0].PLpgSQL_var.refname; + } + } + + return undefined; + } + /** * Deparse a RETURN QUERY statement */ @@ -1718,7 +1781,7 @@ export class PLpgSQLDeparser { parts.push(`${forClause} ${kw('LOOP')}`); if (fors.body) { - const bodyContext = { ...context, indentLevel: context.indentLevel + 1 }; + const bodyContext = { ...context, indentLevel: context.indentLevel + 1, enclosingForVarName: varName }; for (const stmt of fors.body) { const stmtStr = this.deparseStmt(stmt, bodyContext); parts.push(this.indent(stmtStr + ';', bodyContext.indentLevel)); diff --git a/packages/plpgsql-deparser/test-utils/index.ts b/packages/plpgsql-deparser/test-utils/index.ts index e6d856ee..360e655a 100644 --- a/packages/plpgsql-deparser/test-utils/index.ts +++ b/packages/plpgsql-deparser/test-utils/index.ts @@ -1,5 +1,5 @@ import { parsePlPgSQL, parsePlPgSQLSync } from '@libpg-query/parser'; -import { deparseSync, PLpgSQLParseResult } from '../src'; +import { deparseSync, PLpgSQLParseResult, ReturnInfo } from '../src'; import { readFileSync, readdirSync, existsSync } from 'fs'; import * as path from 'path'; import { diff } from 'jest-diff'; @@ -328,6 +328,43 @@ function reconstructSql(originalSql: string, newBody: string): string { return parts.prefix + newBody + parts.suffix; } +/** + * Extract ReturnInfo from a CREATE FUNCTION SQL string by inspecting the + * RETURNS clause and OUT/INOUT parameters. Used by the round-trip test to + * give the deparser enough context to recover dropped retvarno fields. + * + * We intentionally only extract for SETOF and OUT-param functions — these + * are the cases where RETURN NEXT inference needs context. For other + * return types (scalar, void, trigger) we return undefined to avoid + * changing existing RETURN statement behaviour (e.g. bare RETURN in scalar + * functions would become RETURN NULL if we provided returnInfo). + */ +export function extractReturnInfo(sql: string): ReturnInfo | undefined { + // Normalise to a single line for simpler matching + const norm = sql.replace(/\s+/g, ' ').trim(); + + // Strip the body ($$...$$) so we only inspect the signature + const sig = norm.replace(/\$\$.*\$\$/s, ''); + + // OUT / INOUT parameters → out_params (prevents false RETURN NEXT inference) + if (/\b(OUT|INOUT)\s+/i.test(sig)) { + return { kind: 'out_params' }; + } + + // RETURNS TABLE(...) → out_params + if (/RETURNS\s+TABLE\s*\(/i.test(sig)) { + return { kind: 'out_params' }; + } + + // RETURNS SETOF → needed for RETURN NEXT variable recovery + if (/RETURNS\s+SETOF\b/i.test(sig)) { + return { kind: 'setof' }; + } + + // All other types: don't provide returnInfo to preserve existing behaviour + return undefined; +} + export class PLpgSQLTestUtils { protected printErrorMessage(sql: string, position: number) { const lineNumber = sql.slice(0, position).match(/\n/g)?.length || 0; @@ -368,7 +405,8 @@ export class PLpgSQLTestUtils { throw createParseError('PARSE_FAILED', testName, sql); } - const deparsedBody = deparseSync(originalAst); + const returnInfo = extractReturnInfo(sql); + const deparsedBody = deparseSync(originalAst, undefined, returnInfo); if (!deparsedBody || deparsedBody.trim().length === 0) { throw createParseError('DEPARSE_FAILED', testName, sql, deparsedBody);