feat(plpgsql-deparser): enable heterogeneous deparse for AST-based transformations

pyramation · pyramation · commit e6383c90db38 · 2026-01-05T21:24:15.000Z
- Add deparseExprNode() helper to deparse expression AST nodes by wrapping in SELECT
- Update dehydrateQuery() for sql-expr kind to always prefer deparsed AST
- Update dehydrateQuery() for assign kind to prefer deparsed AST nodes
- Update hydrate-demo.test.ts to use AST-based modifications instead of string fields
- Add tests for schema renaming in hydrated expressions (sql-expr, assign, sql-stmt)

This enables AST-based transformations (e.g., schema renaming) to be properly
reflected in the deparsed PL/pgSQL function bodies. Previously, modifications
to AST nodes were ignored because dehydrateQuery() returned the original
string fields instead of deparsing the modified AST.
diff --git a/packages/plpgsql-deparser/__tests__/__snapshots__/hydrate-demo.test.ts.snap b/packages/plpgsql-deparser/__tests__/__snapshots__/hydrate-demo.test.ts.snap
@@ -48,21 +48,24 @@ exports[`hydrate demonstration with big-function.sql should parse, hydrate, modi
   v_min_total numeric := COALESCE(p_min_total, 0);
   v_sql text;
   v_rowcount int := 0;
-  v_lock_key bigint := ('x' || substr(md5(p_org_id::text), 1, 16))::bit(64)::bigint;
+  v_lock_key bigint := CAST(CAST('x' || substr(md5(p_org_id::text), 1, 16) AS pg_catalog.bit(64)) AS bigint);
   sqlstate CONSTANT text;
   sqlerrm CONSTANT text;
 BEGIN
   BEGIN
-      IF p_org_id IS NULL OR p_user_id IS NULL THEN
+      IF p_org_id IS NULL
+          OR p_user_id IS NULL THEN
             RAISE EXCEPTION 'p_org_id and p_user_id are required';
       END IF;
       IF p_from_ts > p_to_ts THEN
             RAISE EXCEPTION 'p_from_ts (%) must be <= p_to_ts (%)', p_from_ts, p_to_ts;
       END IF;
-      IF p_max_rows < 1 OR p_max_rows > 10000 THEN
+      IF p_max_rows < 1
+          OR p_max_rows > 10000 THEN
             RAISE EXCEPTION 'p_max_rows out of range: %', p_max_rows;
       END IF;
-      IF p_round_to < 0 OR p_round_to > 6 THEN
+      IF p_round_to < 0
+          OR p_round_to > 6 THEN
             RAISE EXCEPTION 'p_round_to out of range: %', p_round_to;
       END IF;
       IF p_lock THEN
@@ -104,7 +107,7 @@ BEGIN
             v_discount := 0;
       END IF;
       v_levy := round(GREATEST(v_gross - v_discount, 0) * v_tax_rate, p_round_to);
-      v_net := round((v_gross - v_discount + v_tax) * power(10::numeric, 0), p_round_to);
+      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
@@ -168,11 +171,7 @@ BEGIN
         updated_at = now();
       GET DIAGNOSTICS v_rowcount = ;
       v_orders_upserted := v_rowcount;
-      v_sql := format(
-          'SELECT count(*)::int FROM %I.%I WHERE org_id = $1 AND created_at >= $2 AND created_at < $3',
-          'app_public',
-          'app_order'
-        );
+      v_sql := format('SELECT count(*)::int FROM %I.%I WHERE org_id = $1 AND created_at >= $2 AND created_at < $3', 'app_public', 'app_order');
       EXECUTE v_sql INTO (unnamed row) USING p_org_id, p_from_ts, p_to_ts;
       IF p_debug THEN
             RAISE NOTICE 'dynamic count(app_order)=%', v_rowcount;
@@ -190,10 +189,7 @@ BEGIN
       avg_order_total := round(v_avg, p_round_to);
       top_sku := v_top_sku;
       top_sku_qty := v_top_sku_qty;
-      message := format(
-          'rollup ok: gross=%s discount=%s tax=%s net=%s (discount_rate=%s tax_rate=%s)',
-          v_gross, v_discount, v_tax, v_net, v_discount_rate, v_tax_rate
-        );
+      message := format('rollup ok: gross=%s discount=%s tax=%s net=%s (discount_rate=%s tax_rate=%s)', v_gross, v_discount, v_tax, v_net, v_discount_rate, v_tax_rate);
       RETURN NEXT;
       RETURN;
   END;
diff --git a/packages/plpgsql-deparser/__tests__/hydrate-demo.test.ts b/packages/plpgsql-deparser/__tests__/hydrate-demo.test.ts
@@ -90,6 +90,14 @@ function collectHydratedExprs(obj: any, limit: number): any[] {
   return results;
 }
 
+/**
+ * Modify AST nodes directly (not string fields).
+ * This demonstrates the proper way to transform hydrated PL/pgSQL ASTs.
+ * 
+ * For assign kind: modify targetExpr/valueExpr AST nodes
+ * For sql-expr kind: modify expr AST node
+ * For sql-stmt kind: modify parseResult AST
+ */
 function modifyAst(ast: any): any {
   let modCount = 0;
   let assignModCount = 0;
@@ -102,26 +110,46 @@ function modifyAst(ast: any): any {
         const query = node.PLpgSQL_expr.query;
         
         if (typeof query === 'object' && query.kind === 'assign') {
-          if (query.target === 'v_discount' && assignModCount === 0) {
-            query.target = 'v_rebate';
-            assignModCount++;
-            modCount++;
+          // Modify targetExpr AST node (not the string field)
+          // targetExpr is a ColumnRef with fields array containing String nodes
+          if (query.target === 'v_discount' && query.targetExpr && assignModCount === 0) {
+            // ColumnRef structure: { ColumnRef: { fields: [{ String: { sval: 'v_discount' } }] } }
+            if (query.targetExpr.ColumnRef?.fields?.[0]?.String) {
+              query.targetExpr.ColumnRef.fields[0].String.sval = 'v_rebate';
+              assignModCount++;
+              modCount++;
+            }
           }
-          if (query.target === 'v_tax' && assignModCount === 1) {
-            query.target = 'v_levy';
-            assignModCount++;
-            modCount++;
+          if (query.target === 'v_tax' && query.targetExpr && assignModCount === 1) {
+            if (query.targetExpr.ColumnRef?.fields?.[0]?.String) {
+              query.targetExpr.ColumnRef.fields[0].String.sval = 'v_levy';
+              assignModCount++;
+              modCount++;
+            }
           }
-          if (query.value === '0' && modCount < 5) {
-            query.value = '42';
-            modCount++;
+          // Modify valueExpr AST node for integer constants
+          // A_Const structure: { A_Const: { ival: { ival: 0 } } } or { A_Const: { sval: { sval: '0' } } }
+          if (query.value === '0' && query.valueExpr && modCount < 5) {
+            if (query.valueExpr.A_Const?.ival !== undefined) {
+              query.valueExpr.A_Const.ival.ival = 42;
+              modCount++;
+            } else if (query.valueExpr.A_Const?.sval !== undefined) {
+              query.valueExpr.A_Const.sval.sval = '42';
+              modCount++;
+            }
           }
         }
         
         if (typeof query === 'object' && query.kind === 'sql-expr') {
-          if (query.original === '0' && modCount < 8) {
-            query.original = '42';
-            modCount++;
+          // Modify expr AST node for integer constants
+          if (query.original === '0' && query.expr && modCount < 8) {
+            if (query.expr.A_Const?.ival !== undefined) {
+              query.expr.A_Const.ival.ival = 42;
+              modCount++;
+            } else if (query.expr.A_Const?.sval !== undefined) {
+              query.expr.A_Const.sval.sval = '42';
+              modCount++;
+            }
           }
         }
       }
diff --git a/packages/plpgsql-deparser/__tests__/hydrate.test.ts b/packages/plpgsql-deparser/__tests__/hydrate.test.ts
@@ -1,5 +1,5 @@
 import { loadModule, parsePlPgSQLSync } from '@libpg-query/parser';
-import { hydratePlpgsqlAst, isHydratedExpr, getOriginalQuery, PLpgSQLParseResult } from '../src';
+import { hydratePlpgsqlAst, dehydratePlpgsqlAst, deparseSync, isHydratedExpr, getOriginalQuery, PLpgSQLParseResult } from '../src';
 
 describe('hydratePlpgsqlAst', () => {
   beforeAll(async () => {
@@ -152,8 +152,145 @@ $$`;
         .toBe(result.stats.totalExpressions);
     });
   });
+
+  describe('heterogeneous deparse (AST-based transformations)', () => {
+    it('should deparse modified sql-expr AST nodes (schema renaming)', () => {
+      // Note: This test only checks RangeVar nodes in SQL expressions.
+      // Type references in DECLARE (PLpgSQL_type.typname) are strings, not AST nodes,
+      // and require separate string-based transformation.
+      const sql = `CREATE FUNCTION test_func() RETURNS void
+LANGUAGE plpgsql
+AS $$
+BEGIN
+    IF EXISTS (SELECT 1 FROM "old-schema".users WHERE id = 1) THEN
+        RAISE NOTICE 'found';
+    END IF;
+END;
+$$`;
+
+      const parsed = parsePlPgSQLSync(sql) as unknown as PLpgSQLParseResult;
+      const { ast: hydratedAst } = hydratePlpgsqlAst(parsed);
+
+      // Modify schema names in the hydrated AST
+      transformSchemaNames(hydratedAst, 'old-schema', 'new_schema');
+
+      // Dehydrate and deparse
+      const dehydratedAst = dehydratePlpgsqlAst(hydratedAst);
+      const deparsedBody = deparseSync(dehydratedAst);
+
+      // The deparsed body should contain the new schema name
+      expect(deparsedBody).toContain('new_schema');
+      // And should NOT contain the old schema name
+      expect(deparsedBody).not.toContain('old-schema');
+    });
+
+    it('should deparse modified assign AST nodes (schema renaming in assignments)', () => {
+      const sql = `CREATE FUNCTION test_func() RETURNS void
+LANGUAGE plpgsql
+AS $$
+DECLARE
+    v_count integer;
+BEGIN
+    v_count := (SELECT count(*) FROM "old-schema".users);
+END;
+$$`;
+
+      const parsed = parsePlPgSQLSync(sql) as unknown as PLpgSQLParseResult;
+      const { ast: hydratedAst } = hydratePlpgsqlAst(parsed);
+
+      // Modify schema names in the hydrated AST
+      transformSchemaNames(hydratedAst, 'old-schema', 'new_schema');
+
+      // Dehydrate and deparse
+      const dehydratedAst = dehydratePlpgsqlAst(hydratedAst);
+      const deparsedBody = deparseSync(dehydratedAst);
+
+      // The deparsed body should contain the new schema name
+      expect(deparsedBody).toContain('new_schema');
+      // And should NOT contain the old schema name
+      expect(deparsedBody).not.toContain('old-schema');
+    });
+
+    it('should deparse modified sql-stmt AST nodes (schema renaming in SQL statements)', () => {
+      const sql = `CREATE FUNCTION test_func() RETURNS void
+LANGUAGE plpgsql
+AS $$
+BEGIN
+    INSERT INTO "old-schema".logs (message) VALUES ('test');
+END;
+$$`;
+
+      const parsed = parsePlPgSQLSync(sql) as unknown as PLpgSQLParseResult;
+      const { ast: hydratedAst } = hydratePlpgsqlAst(parsed);
+
+      // Modify schema names in the hydrated AST
+      transformSchemaNames(hydratedAst, 'old-schema', 'new_schema');
+
+      // Dehydrate and deparse
+      const dehydratedAst = dehydratePlpgsqlAst(hydratedAst);
+      const deparsedBody = deparseSync(dehydratedAst);
+
+      // The deparsed body should contain the new schema name
+      expect(deparsedBody).toContain('new_schema');
+      // And should NOT contain the old schema name
+      expect(deparsedBody).not.toContain('old-schema');
+    });
+  });
 });
 
+/**
+ * Helper function to transform schema names in a hydrated PL/pgSQL AST.
+ * This walks the AST and modifies schemaname properties wherever they appear.
+ */
+function transformSchemaNames(obj: any, oldSchema: string, newSchema: string): void {
+  if (obj === null || obj === undefined) return;
+
+  if (typeof obj === 'object') {
+    // Check for RangeVar nodes (table references) - wrapped in RangeVar key
+    if ('RangeVar' in obj && obj.RangeVar?.schemaname === oldSchema) {
+      obj.RangeVar.schemaname = newSchema;
+    }
+    
+    // Check for direct schemaname property (e.g., InsertStmt.relation, UpdateStmt.relation)
+    // These are RangeVar-like objects without the RangeVar wrapper
+    if ('schemaname' in obj && obj.schemaname === oldSchema && 'relname' in obj) {
+      obj.schemaname = newSchema;
+    }
+
+    // Check for hydrated expressions with sql-expr kind
+    if ('PLpgSQL_expr' in obj) {
+      const query = obj.PLpgSQL_expr.query;
+      if (query && typeof query === 'object') {
+        // Transform the embedded SQL AST
+        if (query.kind === 'sql-expr' && query.expr) {
+          transformSchemaNames(query.expr, oldSchema, newSchema);
+        }
+        if (query.kind === 'sql-stmt' && query.parseResult) {
+          transformSchemaNames(query.parseResult, oldSchema, newSchema);
+        }
+        if (query.kind === 'assign') {
+          if (query.valueExpr) {
+            transformSchemaNames(query.valueExpr, oldSchema, newSchema);
+          }
+          if (query.targetExpr) {
+            transformSchemaNames(query.targetExpr, oldSchema, newSchema);
+          }
+        }
+      }
+    }
+
+    for (const value of Object.values(obj)) {
+      transformSchemaNames(value, oldSchema, newSchema);
+    }
+  }
+
+  if (Array.isArray(obj)) {
+    for (const item of obj) {
+      transformSchemaNames(item, oldSchema, newSchema);
+    }
+  }
+}
+
 function findExprByKind(obj: any, kind: string): any {
   if (obj === null || obj === undefined) return null;
   
diff --git a/packages/plpgsql-deparser/src/hydrate.ts b/packages/plpgsql-deparser/src/hydrate.ts
@@ -409,13 +409,69 @@ function dehydrateNode(node: any, options?: DehydrationOptions): any {
   return result;
 }
 
+/**
+ * Deparse a single expression AST node by wrapping it in a SELECT statement,
+ * deparsing, and stripping the SELECT prefix.
+ */
+function deparseExprNode(expr: Node, sqlDeparseOptions?: DeparserOptions): string | null {
+  try {
+    // Wrap the expression in a minimal SELECT statement
+    const wrappedStmt = {
+      SelectStmt: {
+        targetList: [
+          {
+            ResTarget: {
+              val: expr
+            }
+          }
+        ]
+      }
+    };
+    const deparsed = Deparser.deparse(wrappedStmt, sqlDeparseOptions);
+    // Strip the "SELECT " prefix (case-insensitive, handles whitespace/newlines)
+    const stripped = deparsed.replace(/^SELECT\s+/i, '').replace(/;?\s*$/, '');
+    return stripped;
+  } catch {
+    return null;
+  }
+}
+
+/**
+ * Normalize whitespace for comparison purposes.
+ * This helps detect if a string field was modified vs just having different formatting.
+ */
+function normalizeForComparison(str: string): string {
+  return str.replace(/\s+/g, ' ').trim().toLowerCase();
+}
+
 function dehydrateQuery(query: HydratedExprQuery, sqlDeparseOptions?: DeparserOptions): string {
   switch (query.kind) {
     case 'assign': {
-      // For assignments, use the target and value strings directly
-      // These may have been modified by the caller
+      // For assignments, always prefer deparsing the AST nodes if they exist.
+      // This enables AST-based transformations (e.g., schema renaming).
+      // Fall back to string fields if AST nodes are missing or deparse fails.
       const assignQuery = query as HydratedExprAssign;
-      return `${assignQuery.target} := ${assignQuery.value}`;
+      
+      let target = assignQuery.target;
+      let value = assignQuery.value;
+      
+      // For target: prefer deparsed AST if available
+      if (assignQuery.targetExpr) {
+        const deparsedTarget = deparseExprNode(assignQuery.targetExpr, sqlDeparseOptions);
+        if (deparsedTarget !== null) {
+          target = deparsedTarget;
+        }
+      }
+      
+      // For value: prefer deparsed AST if available
+      if (assignQuery.valueExpr) {
+        const deparsedValue = deparseExprNode(assignQuery.valueExpr, sqlDeparseOptions);
+        if (deparsedValue !== null) {
+          value = deparsedValue;
+        }
+      }
+      
+      return `${target} := ${value}`;
     }
     case 'sql-stmt': {
       // Deparse the modified parseResult back to SQL
@@ -432,11 +488,20 @@ function dehydrateQuery(query: HydratedExprQuery, sqlDeparseOptions?: DeparserOp
       }
       return query.original;
     }
-    case 'sql-expr':
-      // For sql-expr, return the original string
-      // Callers can modify query.original directly for simple transformations
-      // For AST-based transformations, use sql-stmt instead
+    case 'sql-expr': {
+      // For sql-expr, always prefer deparsing the AST.
+      // This enables AST-based transformations (e.g., schema renaming).
+      // Fall back to original only if deparse fails.
+      const exprQuery = query as HydratedExprSqlExpr;
+      if (exprQuery.expr) {
+        const deparsed = deparseExprNode(exprQuery.expr, sqlDeparseOptions);
+        if (deparsed !== null) {
+          return deparsed;
+        }
+      }
+      // Fall back to original if deparse fails
       return query.original;
+    }
     case 'raw':
     default:
       return query.original;
