feat(plpgsql-deparser): add context-based RETURN statement handling

Add ReturnInfo type and context to PLpgSQLDeparserContext for correct
RETURN statement syntax based on function return type:
- void/setof/trigger/out_params: bare RETURN is valid
- scalar: RETURN NULL is required for empty returns

Add getReturnInfo helper to plpgsql-parser to extract return type info
from CreateFunctionStmt AST node.

When returnInfo is provided in context, the deparser uses it to determine
the correct RETURN syntax. When not provided, it falls back to bare RETURN
for backward compatibility.

Includes tests for all return type scenarios (void, setof, trigger,
scalar, OUT params, RETURNS TABLE).

Author: pyramation

packages/plpgsql-deparser/__tests__/__snapshots__/hydrate-demo.test.ts.snap

@@ -69,7 +69,7 @@ BEGIN
             RAISE EXCEPTION 'p_round_to out of range: %', p_round_to;
       END IF;
       IF p_lock THEN
-            PERFORM SELECT pg_advisory_xact_lock(v_lock_key);
+            PERFORM pg_advisory_xact_lock(v_lock_key);
       END IF;
       IF p_debug THEN
             RAISE NOTICE 'big_kitchen_sink start=% org=% user=% from=% to=% min_total=%', v_now, p_org_id, p_user_id, p_from_ts, p_to_ts, v_min_total;
@@ -99,7 +99,7 @@ BEGIN
       SELECT
         t.orders_scanned,
         t.gross_total,
-        t.avg_total
+        t.avg_total INTO v_orders_scanned, v_gross, v_avg
       FROM totals AS t;
       IF p_apply_discount THEN
             v_rebate := round(v_gross * GREATEST(LEAST(v_discount_rate + v_jitter, 0.50), 0), p_round_to);
@@ -110,7 +110,7 @@ BEGIN
       v_net := round(((v_gross - v_discount) + v_tax) * power(10::numeric, 0), p_round_to);
       SELECT
         oi.sku,
-        CAST(sum(oi.quantity) AS bigint) AS qty
+        CAST(sum(oi.quantity) AS bigint) AS qty INTO v_top_sku, v_top_sku_qty
       FROM app_public.order_item AS oi
       JOIN app_public.app_order AS o ON o.id = oi.order_id
       WHERE

packages/plpgsql-deparser/__tests__/pretty/__snapshots__/plpgsql-pretty.test.ts.snap

@@ -36,7 +36,7 @@ begin
             raise exception 'p_round_to out of range: %', p_round_to;
       end if;
       if p_lock then
-            perform SELECT pg_advisory_xact_lock(v_lock_key);
+            perform pg_advisory_xact_lock(v_lock_key);
       end if;
       if p_debug then
             raise notice 'big_kitchen_sink start=% org=% user=% from=% to=% min_total=%', v_now, p_org_id, p_user_id, p_from_ts, p_to_ts, v_min_total;
@@ -67,7 +67,7 @@ begin
         SELECT
           t.orders_scanned,
           t.gross_total,
-          t.avg_total
+          t.avg_total into v_orders_scanned, v_gross, v_avg
                                                             FROM totals t;
       if p_apply_discount then
             v_discount := round(v_gross * GREATEST(LEAST(v_discount_rate + v_jitter, 0.50), 0), p_round_to);
@@ -78,7 +78,7 @@ begin
       v_net := round((v_gross - v_discount + v_tax) * power(10::numeric, 0), p_round_to);
       SELECT
           oi.sku,
-          sum(oi.quantity)::bigint AS qty
+          sum(oi.quantity)::bigint AS qty into v_top_sku, v_top_sku_qty
                                         FROM app_public.order_item oi
         JOIN app_public.app_order o ON o.id = oi.order_id
         WHERE o.org_id = p_org_id
@@ -237,7 +237,7 @@ BEGIN
             RAISE EXCEPTION 'p_round_to out of range: %', p_round_to;
       END IF;
       IF p_lock THEN
-            PERFORM SELECT pg_advisory_xact_lock(v_lock_key);
+            PERFORM pg_advisory_xact_lock(v_lock_key);
       END IF;
       IF p_debug THEN
             RAISE NOTICE 'big_kitchen_sink start=% org=% user=% from=% to=% min_total=%', v_now, p_org_id, p_user_id, p_from_ts, p_to_ts, v_min_total;
@@ -268,7 +268,7 @@ BEGIN
         SELECT
           t.orders_scanned,
           t.gross_total,
-          t.avg_total
+          t.avg_total INTO v_orders_scanned, v_gross, v_avg
                                                             FROM totals t;
       IF p_apply_discount THEN
             v_discount := round(v_gross * GREATEST(LEAST(v_discount_rate + v_jitter, 0.50), 0), p_round_to);
@@ -279,7 +279,7 @@ BEGIN
       v_net := round((v_gross - v_discount + v_tax) * power(10::numeric, 0), p_round_to);
       SELECT
           oi.sku,
-          sum(oi.quantity)::bigint AS qty
+          sum(oi.quantity)::bigint AS qty INTO v_top_sku, v_top_sku_qty
                                         FROM app_public.order_item oi
         JOIN app_public.app_order o ON o.id = oi.order_id
         WHERE o.org_id = p_org_id

packages/plpgsql-deparser/__tests__/return-context.test.ts

@@ -0,0 +1,134 @@
+import { loadModule, parsePlPgSQLSync, parseSync } from '@libpg-query/parser';
+import { deparseSync, ReturnInfo } from '../src';
+import { PLpgSQLParseResult } from '../src/types';
+
+describe('RETURN statement context handling', () => {
+  beforeAll(async () => {
+    await loadModule();
+  });
+
+  describe('deparseSync with returnInfo context', () => {
+    const parseBody = (sql: string): PLpgSQLParseResult => {
+      return parsePlPgSQLSync(sql) as unknown as PLpgSQLParseResult;
+    };
+
+    it('should output bare RETURN for void functions', () => {
+      const sql = `CREATE FUNCTION test_void() RETURNS void
+LANGUAGE plpgsql
+AS $$
+BEGIN
+    RETURN;
+END;
+$$`;
+      const parsed = parseBody(sql);
+      const returnInfo: ReturnInfo = { kind: 'void' };
+      const result = deparseSync(parsed, undefined, returnInfo);
+      
+      expect(result).toContain('RETURN;');
+      expect(result).not.toContain('RETURN NULL');
+    });
+
+    it('should output bare RETURN for setof functions', () => {
+      const sql = `CREATE FUNCTION test_setof() RETURNS SETOF integer
+LANGUAGE plpgsql
+AS $$
+BEGIN
+    RETURN NEXT 1;
+    RETURN NEXT 2;
+    RETURN;
+END;
+$$`;
+      const parsed = parseBody(sql);
+      const returnInfo: ReturnInfo = { kind: 'setof' };
+      const result = deparseSync(parsed, undefined, returnInfo);
+      
+      expect(result).toContain('RETURN;');
+      expect(result).not.toContain('RETURN NULL');
+    });
+
+    it('should output bare RETURN for trigger functions', () => {
+      const sql = `CREATE FUNCTION test_trigger() RETURNS trigger
+LANGUAGE plpgsql
+AS $$
+BEGIN
+    RETURN NEW;
+END;
+$$`;
+      const parsed = parseBody(sql);
+      const returnInfo: ReturnInfo = { kind: 'trigger' };
+      const result = deparseSync(parsed, undefined, returnInfo);
+      
+      expect(result).toContain('RETURN NEW');
+    });
+
+    it('should output bare RETURN for out_params functions', () => {
+      const sql = `CREATE FUNCTION test_out(OUT result integer) RETURNS integer
+LANGUAGE plpgsql
+AS $$
+BEGIN
+    result := 42;
+    RETURN;
+END;
+$$`;
+      const parsed = parseBody(sql);
+      const returnInfo: ReturnInfo = { kind: 'out_params' };
+      const result = deparseSync(parsed, undefined, returnInfo);
+      
+      expect(result).toContain('RETURN;');
+      expect(result).not.toContain('RETURN NULL');
+    });
+
+    it('should output RETURN NULL for scalar functions with empty return', () => {
+      const sql = `CREATE FUNCTION test_scalar(val integer) RETURNS text
+LANGUAGE plpgsql
+AS $$
+BEGIN
+    IF val > 0 THEN
+        RETURN 'positive';
+    END IF;
+    RETURN;
+END;
+$$`;
+      const parsed = parseBody(sql);
+      const returnInfo: ReturnInfo = { kind: 'scalar' };
+      const result = deparseSync(parsed, undefined, returnInfo);
+      
+      expect(result).toContain('RETURN NULL');
+    });
+
+    it('should preserve RETURN with expression regardless of context', () => {
+      const sql = `CREATE FUNCTION test_expr() RETURNS integer
+LANGUAGE plpgsql
+AS $$
+BEGIN
+    RETURN 42;
+END;
+$$`;
+      const parsed = parseBody(sql);
+      
+      // Test with scalar context
+      const scalarResult = deparseSync(parsed, undefined, { kind: 'scalar' });
+      expect(scalarResult).toContain('RETURN 42');
+      
+      // Test with void context (shouldn't change expression returns)
+      const voidResult = deparseSync(parsed, undefined, { kind: 'void' });
+      expect(voidResult).toContain('RETURN 42');
+    });
+
+    it('should default to bare RETURN when no context provided (backward compatibility)', () => {
+      const sql = `CREATE FUNCTION test_no_context() RETURNS void
+LANGUAGE plpgsql
+AS $$
+BEGIN
+    RETURN;
+END;
+$$`;
+      const parsed = parseBody(sql);
+      
+      // No returnInfo provided - should default to bare RETURN
+      const result = deparseSync(parsed);
+      expect(result).toContain('RETURN;');
+      expect(result).not.toContain('RETURN NULL');
+    });
+  });
+});

packages/plpgsql-deparser/src/index.ts

@@ -1,4 +1,4 @@
-import { PLpgSQLDeparser, PLpgSQLDeparserOptions } from './plpgsql-deparser';
+import { PLpgSQLDeparser, PLpgSQLDeparserOptions, ReturnInfo, ReturnInfoKind } from './plpgsql-deparser';
 
 const deparseMethod = PLpgSQLDeparser.deparse;
 const deparseFunctionMethod = PLpgSQLDeparser.deparseFunction;
@@ -18,7 +18,7 @@ export const deparseFunction = async (
   return deparseFunctionMethod(...args);
 };
 
-export { PLpgSQLDeparser, PLpgSQLDeparserOptions };
+export { PLpgSQLDeparser, PLpgSQLDeparserOptions, ReturnInfo, ReturnInfoKind };
 export * from './types';
 export * from './hydrate-types';
 export { hydratePlpgsqlAst, dehydratePlpgsqlAst, isHydratedExpr, getOriginalQuery, DehydrationOptions } from './hydrate';

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

@@ -65,10 +65,23 @@ export interface PLpgSQLDeparserOptions {
   uppercase?: boolean;
 }
 
+/**
+ * Return type information for a PL/pgSQL function.
+ * Used to determine the correct RETURN statement syntax:
+ * - void/setof/trigger/out_params: bare RETURN is valid
+ * - scalar: RETURN NULL is required for empty returns
+ */
+export type ReturnInfoKind = 'void' | 'setof' | 'trigger' | 'scalar' | 'out_params';
+
+export interface ReturnInfo {
+  kind: ReturnInfoKind;
+}
+
 export interface PLpgSQLDeparserContext {
   indentLevel: number;
   options: PLpgSQLDeparserOptions;
   datums?: PLpgSQLDatum[];
+  returnInfo?: ReturnInfo;
 }
 
 /**
@@ -90,49 +103,62 @@ export class PLpgSQLDeparser {
 
   /**
    * Static method to deparse a PL/pgSQL parse result
+   * @param parseResult - The PL/pgSQL parse result
+   * @param options - Deparser options
+   * @param returnInfo - Optional return type info for correct RETURN statement handling
    */
-  static deparse(parseResult: PLpgSQLParseResult, options?: PLpgSQLDeparserOptions): string {
-    return new PLpgSQLDeparser(options).deparseResult(parseResult);
+  static deparse(parseResult: PLpgSQLParseResult, options?: PLpgSQLDeparserOptions, returnInfo?: ReturnInfo): string {
+    return new PLpgSQLDeparser(options).deparseResult(parseResult, returnInfo);
   }
 
   /**
    * Static method to deparse a single PL/pgSQL function body
+   * @param func - The PL/pgSQL function AST
+   * @param options - Deparser options
+   * @param returnInfo - Optional return type info for correct RETURN statement handling
    */
-  static deparseFunction(func: PLpgSQL_function, options?: PLpgSQLDeparserOptions): string {
-    return new PLpgSQLDeparser(options).deparseFunction(func);
+  static deparseFunction(func: PLpgSQL_function, options?: PLpgSQLDeparserOptions, returnInfo?: ReturnInfo): string {
+    return new PLpgSQLDeparser(options).deparseFunction(func, returnInfo);
   }
 
   /**
    * Deparse a complete PL/pgSQL parse result
+   * @param parseResult - The PL/pgSQL parse result
+   * @param returnInfo - Optional return type info for correct RETURN statement handling
    */
-  deparseResult(parseResult: PLpgSQLParseResult): string {
+  deparseResult(parseResult: PLpgSQLParseResult, returnInfo?: ReturnInfo): string {
     if (!parseResult.plpgsql_funcs || parseResult.plpgsql_funcs.length === 0) {
       return '';
     }
 
     return parseResult.plpgsql_funcs
-      .map(func => this.deparseFunctionNode(func))
+      .map(func => this.deparseFunctionNode(func, returnInfo))
       .join(this.options.newline + this.options.newline);
   }
 
   /**
    * Deparse a PLpgSQL_function node wrapper
+   * @param node - The PLpgSQL_function node wrapper
+   * @param returnInfo - Optional return type info for correct RETURN statement handling
    */
-  deparseFunctionNode(node: PLpgSQLFunctionNode): string {
+  deparseFunctionNode(node: PLpgSQLFunctionNode, returnInfo?: ReturnInfo): string {
     if ('PLpgSQL_function' in node) {
-      return this.deparseFunction(node.PLpgSQL_function);
+      return this.deparseFunction(node.PLpgSQL_function, returnInfo);
     }
     throw new Error('Unknown function node type');
   }
 
   /**
    * Deparse a PL/pgSQL function body
+   * @param func - The PL/pgSQL function AST
+   * @param returnInfo - Optional return type info for correct RETURN statement handling
    */
-  deparseFunction(func: PLpgSQL_function): string {
+  deparseFunction(func: PLpgSQL_function, returnInfo?: ReturnInfo): string {
     const context: PLpgSQLDeparserContext = {
       indentLevel: 0,
       options: this.options,
       datums: func.datums,
+      returnInfo,
     };
 
     const parts: string[] = [];
@@ -930,6 +956,13 @@ export class PLpgSQLDeparser {
 
   /**
    * Deparse a RETURN statement
+   * 
+   * PostgreSQL requires different RETURN syntax based on function type:
+   * - void/setof/trigger/out_params: bare RETURN is valid
+   * - scalar: RETURN NULL is required for empty returns
+   * 
+   * When returnInfo is provided in context, we use it to determine the correct syntax.
+   * When not provided, we fall back to heuristics that scan the function body.
    */
   private deparseReturn(ret: PLpgSQL_stmt_return, context: PLpgSQLDeparserContext): string {
     const kw = this.keyword;
@@ -943,6 +976,19 @@ export class PLpgSQLDeparser {
       return `${kw('RETURN')} ${varName}`;
     }
 
+    // Empty RETURN - need to determine if we should output bare RETURN or RETURN NULL
+    // Use context.returnInfo if available, otherwise use heuristics
+    if (context.returnInfo) {
+      // Context-based: use the provided return type info
+      if (context.returnInfo.kind === 'scalar') {
+        return `${kw('RETURN')} ${kw('NULL')}`;
+      }
+      // void, setof, trigger, out_params all use bare RETURN
+      return kw('RETURN');
+    }
+    
+    // Heuristic fallback: bare RETURN is the safest default
+    // This maintains backward compatibility for callers that don't provide returnInfo
     return kw('RETURN');
   }