Merge pull request #256 from constructive-io/devin/1767232213-fix-dehydrate-ast

feat: deparse modified AST and improve INSERT pretty-printing

Author: pyramation

packages/deparser/src/deparser.ts

@@ -563,7 +563,9 @@ export class Deparser implements DeparserVisitor {
         output.push('VALUES');
         const lists = ListUtils.unwrapList(node.valuesLists).map(list => {
           const values = ListUtils.unwrapList(list).map(val => this.visit(val as Node, context));
-          return context.parens(values.join(', '));
+          // Put each value on its own line for pretty printing
+          const indentedValues = values.map(val => context.indent(val));
+          return '(\n' + indentedValues.join(',\n') + '\n)';
         });
         const indentedTuples = lists.map(tuple => {
           if (this.containsMultilineStringLiteral(tuple)) {
@@ -1116,7 +1118,13 @@ export class Deparser implements DeparserVisitor {
           } else {
             const updateContext = context.spawn('UpdateStmt', { update: true });
             const targets = targetList.map(target => this.visit(target as Node, updateContext));
-            output.push(targets.join(', '));
+            if (context.isPretty()) {
+              // Put each assignment on its own line for pretty printing
+              const indentedTargets = targets.map(target => context.indent(target));
+              output.push('\n' + indentedTargets.join(',\n'));
+            } else {
+              output.push(targets.join(', '));
+            }
           }
         }
 

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

@@ -71,34 +71,33 @@ 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;
 END IF;
-    WITH base AS (
-    SELECT
-      o.id,
-      o.total_amount::numeric AS total_amount,
-      o.currency,
-      o.created_at
-    FROM app_public.app_order o
-    WHERE o.org_id = p_org_id
-      AND o.user_id = p_user_id
-      AND o.created_at >= p_from_ts
-      AND o.created_at <  p_to_ts
-      AND o.total_amount::numeric >= v_min_total
-      AND o.currency = p_currency
-    ORDER BY o.created_at DESC
-    LIMIT p_max_rows
-  ),
-  totals AS (
-    SELECT
-      count(*)::int AS orders_scanned,
-      COALESCE(sum(total_amount), 0) AS gross_total,
-      COALESCE(avg(total_amount), 0) AS avg_total
-    FROM base
-  )
-  SELECT
-    t.orders_scanned,
-    t.gross_total,
-    t.avg_total
-                                                      FROM totals t;
+    WITH 
+  base AS (SELECT
+    o.id,
+    o.total_amount::numeric AS total_amount,
+    o.currency,
+    o.created_at
+  FROM app_public.app_order AS o
+  WHERE
+    o.org_id = p_org_id
+    AND o.user_id = p_user_id
+    AND o.created_at >= p_from_ts
+    AND o.created_at < p_to_ts
+    AND o.total_amount::numeric >= v_min_total
+    AND o.currency = p_currency
+  ORDER BY
+    o.created_at DESC
+  LIMIT p_max_rows),
+  totals AS (SELECT
+    (count(*))::int AS orders_scanned,
+    COALESCE(sum(total_amount), 0) AS gross_total,
+    COALESCE(avg(total_amount), 0) AS avg_total
+  FROM base)
+SELECT
+  t.orders_scanned,
+  t.gross_total,
+  t.avg_total
+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);
 ELSE
@@ -107,36 +106,40 @@ 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);
     SELECT
-    oi.sku,
-    sum(oi.quantity)::bigint AS 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
-    AND o.user_id = p_user_id
-    AND o.created_at >= p_from_ts
-    AND o.created_at <  p_to_ts
-    AND o.currency = p_currency
-  GROUP BY oi.sku
-  ORDER BY qty DESC, oi.sku ASC
-  LIMIT 1;
+  oi.sku,
+  CAST(sum(oi.quantity) AS bigint) AS qty
+FROM app_public.order_item AS oi
+JOIN app_public.app_order AS o ON o.id = oi.order_id
+WHERE
+  o.org_id = p_org_id
+  AND o.user_id = p_user_id
+  AND o.created_at >= p_from_ts
+  AND o.created_at < p_to_ts
+  AND o.currency = p_currency
+GROUP BY
+  oi.sku
+ORDER BY
+  qty DESC,
+  oi.sku ASC
+LIMIT 1;
     INSERT INTO app_public.order_rollup (
-    org_id,
-    user_id,
-    period_from,
-    period_to,
-    currency,
-    orders_scanned,
-    gross_total,
-    discount_total,
-    tax_total,
-    net_total,
-    avg_order_total,
-    top_sku,
-    top_sku_qty,
-    note,
-    updated_at
-  )
-  VALUES (
+  org_id,
+  user_id,
+  period_from,
+  period_to,
+  currency,
+  orders_scanned,
+  gross_total,
+  discount_total,
+  tax_total,
+  net_total,
+  avg_order_total,
+  top_sku,
+  top_sku_qty,
+  note,
+  updated_at
+) VALUES
+  (
     p_org_id,
     p_user_id,
     p_from_ts,
@@ -152,19 +155,17 @@ END IF;
     v_top_sku_qty,
     p_note,
     now()
-  )
-  ON CONFLICT (org_id, user_id, period_from, period_to, currency)
-  DO UPDATE SET
-    orders_scanned   = EXCLUDED.orders_scanned,
-    gross_total      = EXCLUDED.gross_total,
-    discount_total   = EXCLUDED.discount_total,
-    tax_total        = EXCLUDED.tax_total,
-    net_total        = EXCLUDED.net_total,
-    avg_order_total  = EXCLUDED.avg_order_total,
-    top_sku          = EXCLUDED.top_sku,
-    top_sku_qty      = EXCLUDED.top_sku_qty,
-    note             = COALESCE(EXCLUDED.note, app_public.order_rollup.note),
-    updated_at       = now();
+  ) ON CONFLICT (org_id, user_id, period_from, period_to, currency) DO UPDATE SET 
+  orders_scanned = excluded.orders_scanned,
+  gross_total = excluded.gross_total,
+  discount_total = excluded.discount_total,
+  tax_total = excluded.tax_total,
+  net_total = excluded.net_total,
+  avg_order_total = excluded.avg_order_total,
+  top_sku = excluded.top_sku,
+  top_sku_qty = excluded.top_sku_qty,
+  note = COALESCE(excluded.note, app_public.order_rollup.note),
+  updated_at = now();
     GET DIAGNOSTICS v_rowcount = ;
     v_orders_upserted := v_rowcount;
     v_sql := format(

packages/plpgsql-deparser/src/hydrate.ts

@@ -1,9 +1,11 @@
 import { parseSync, scanSync } from '@libpg-query/parser';
 import { ParseResult, Node } from '@pgsql/types';
+import { Deparser, DeparserOptions } from 'pgsql-deparser';
 import {
   HydratedExprQuery,
   HydratedExprRaw,
   HydratedExprSqlExpr,
+  HydratedExprSqlStmt,
   HydratedExprAssign,
   HydrationOptions,
   HydrationResult,
@@ -13,6 +15,18 @@ import {
 } from './hydrate-types';
 import { PLpgSQLParseResult } from './types';
 
+/**
+ * Options for dehydrating (converting back to strings) a hydrated PL/pgSQL AST
+ */
+export interface DehydrationOptions {
+  /**
+   * Options to pass to the SQL deparser when deparsing sql-stmt expressions.
+   * This allows callers to control formatting (pretty printing, etc.) of
+   * embedded SQL statements inside PL/pgSQL function bodies.
+   */
+  sqlDeparseOptions?: DeparserOptions;
+}
+
 function extractExprFromSelectWrapper(result: ParseResult): Node | undefined {
   const stmt = result.stmts?.[0]?.stmt as any;
   if (stmt?.SelectStmt?.targetList?.[0]?.ResTarget?.val) {
@@ -350,17 +364,17 @@ export function getOriginalQuery(query: string | HydratedExprQuery): string {
   return query.original;
 }
 
-export function dehydratePlpgsqlAst<T>(ast: T): T {
-  return dehydrateNode(ast) as T;
+export function dehydratePlpgsqlAst<T>(ast: T, options?: DehydrationOptions): T {
+  return dehydrateNode(ast, options) as T;
 }
 
-function dehydrateNode(node: any): any {
+function dehydrateNode(node: any, options?: DehydrationOptions): any {
   if (node === null || node === undefined) {
     return node;
   }
 
   if (Array.isArray(node)) {
-    return node.map(item => dehydrateNode(item));
+    return node.map(item => dehydrateNode(item, options));
   }
 
   if (typeof node !== 'object') {
@@ -375,7 +389,7 @@ function dehydrateNode(node: any): any {
     if (typeof query === 'string') {
       dehydratedQuery = query;
     } else if (isHydratedExpr(query)) {
-      dehydratedQuery = dehydrateQuery(query);
+      dehydratedQuery = dehydrateQuery(query, options?.sqlDeparseOptions);
     } else {
       dehydratedQuery = String(query);
     }
@@ -390,17 +404,39 @@ function dehydrateNode(node: any): any {
 
   const result: any = {};
   for (const [key, value] of Object.entries(node)) {
-    result[key] = dehydrateNode(value);
+    result[key] = dehydrateNode(value, options);
   }
   return result;
 }
 
-function dehydrateQuery(query: HydratedExprQuery): string {
+function dehydrateQuery(query: HydratedExprQuery, sqlDeparseOptions?: DeparserOptions): string {
   switch (query.kind) {
-    case 'assign':
-      return `${query.target} := ${query.value}`;
+    case 'assign': {
+      // For assignments, use the target and value strings directly
+      // These may have been modified by the caller
+      const assignQuery = query as HydratedExprAssign;
+      return `${assignQuery.target} := ${assignQuery.value}`;
+    }
+    case 'sql-stmt': {
+      // Deparse the modified parseResult back to SQL
+      // This enables AST-based transformations (e.g., schema renaming)
+      // Pass through sqlDeparseOptions to control formatting (pretty printing, etc.)
+      const stmtQuery = query as HydratedExprSqlStmt;
+      if (stmtQuery.parseResult?.stmts?.[0]?.stmt) {
+        try {
+          return Deparser.deparse(stmtQuery.parseResult.stmts[0].stmt, sqlDeparseOptions);
+        } catch {
+          // Fall back to original if deparse fails
+          return query.original;
+        }
+      }
+      return query.original;
+    }
     case 'sql-expr':
-    case 'sql-stmt':
+      // For sql-expr, return the original string
+      // Callers can modify query.original directly for simple transformations
+      // For AST-based transformations, use sql-stmt instead
+      return query.original;
     case 'raw':
     default:
       return query.original;

packages/plpgsql-deparser/src/index.ts

@@ -21,4 +21,4 @@ export const deparseFunction = async (
 export { PLpgSQLDeparser, PLpgSQLDeparserOptions };
 export * from './types';
 export * from './hydrate-types';
-export { hydratePlpgsqlAst, dehydratePlpgsqlAst, isHydratedExpr, getOriginalQuery } from './hydrate';
+export { hydratePlpgsqlAst, dehydratePlpgsqlAst, isHydratedExpr, getOriginalQuery, DehydrationOptions } from './hydrate';