fix: recover RETURN NEXT variable lost by libpg-query + CoalesceExpr/RangeFunction test fixtures

__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$$",

__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$$;

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