From e6383c90db389fdb80e456cb54b30121e7ef436d Mon Sep 17 00:00:00 2001 From: Dan Lynch Date: Mon, 5 Jan 2026 21:24:15 +0000 Subject: [PATCH] feat(plpgsql-deparser): enable heterogeneous deparse for AST-based transformations - 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. --- .../__snapshots__/hydrate-demo.test.ts.snap | 24 ++- .../__tests__/hydrate-demo.test.ts | 56 +++++-- .../__tests__/hydrate.test.ts | 139 +++++++++++++++++- packages/plpgsql-deparser/src/hydrate.ts | 79 +++++++++- 4 files changed, 262 insertions(+), 36 deletions(-) diff --git a/packages/plpgsql-deparser/__tests__/__snapshots__/hydrate-demo.test.ts.snap b/packages/plpgsql-deparser/__tests__/__snapshots__/hydrate-demo.test.ts.snap index a2e36417..ac5cd8bb 100644 --- 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 index 29c1bb4b..ab4bbb14 100644 --- 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 index 8a2eecdc..8357f217 100644 --- 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 index d9e7a3f7..44e8e18f 100644 --- 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;