From 688c95ef7b7461635a7ed3c182c1d64f86b5e194 Mon Sep 17 00:00:00 2001 From: Dan Lynch Date: Tue, 6 Jan 2026 16:44:43 +0000 Subject: [PATCH 1/2] test(plpgsql-deparser): add schema transform demo test Demonstrates the heterogeneous AST transformation pipeline: - Parse SQL containing PL/pgSQL functions - Hydrate embedded SQL expressions into AST nodes - Traverse and transform schema names in both outer SQL AST and embedded SQL - Dehydrate back to strings - Deparse to final SQL output Includes 4 test cases: - Simple function with schema-qualified table reference - Trigger function with INSERT into schema-qualified table - RETURN QUERY function with schema-qualified table - Function calls inside PL/pgSQL expressions --- .../__tests__/schema-transform.demo.test.ts | 376 ++++++++++++++++++ 1 file changed, 376 insertions(+) create mode 100644 packages/plpgsql-deparser/__tests__/schema-transform.demo.test.ts diff --git a/packages/plpgsql-deparser/__tests__/schema-transform.demo.test.ts b/packages/plpgsql-deparser/__tests__/schema-transform.demo.test.ts new file mode 100644 index 00000000..34194300 --- /dev/null +++ b/packages/plpgsql-deparser/__tests__/schema-transform.demo.test.ts @@ -0,0 +1,376 @@ +/** + * Schema Transform Demo + * + * This test demonstrates the heterogeneous AST transformation pipeline: + * 1. Parse SQL containing PL/pgSQL functions + * 2. Hydrate embedded SQL expressions into AST nodes + * 3. Traverse and transform schema names in both: + * - Outer SQL AST (CreateFunctionStmt, return types, etc.) + * - Embedded SQL inside PL/pgSQL function bodies + * 4. Dehydrate back to strings + * 5. Deparse to final SQL output + * + * This pattern is useful for: + * - Schema renaming (e.g., old_schema -> new_schema) + * - Identifier rewriting + * - Cross-cutting AST transformations + */ + +import { loadModule, parsePlPgSQLSync, parseSync } from '@libpg-query/parser'; +import { Deparser } from 'pgsql-deparser'; +import { hydratePlpgsqlAst, dehydratePlpgsqlAst, PLpgSQLParseResult, deparseSync } from '../src'; + +describe('schema transform demo', () => { + beforeAll(async () => { + await loadModule(); + }); + + /** + * Transform schema names in SQL AST nodes. + * Handles RangeVar, TypeName, FuncCall, and other schema-qualified references. + */ + function transformSchemaInSqlAst( + node: any, + oldSchema: string, + newSchema: string + ): void { + if (node === null || node === undefined || typeof node !== 'object') { + return; + } + + if (Array.isArray(node)) { + for (const item of node) { + transformSchemaInSqlAst(item, oldSchema, newSchema); + } + return; + } + + // Handle RangeVar nodes (table references like old_schema.users) + if ('RangeVar' in node) { + const rangeVar = node.RangeVar; + if (rangeVar.schemaname === oldSchema) { + rangeVar.schemaname = newSchema; + } + } + + // Handle direct relation references (INSERT/UPDATE/DELETE statements) + // These have schemaname directly on the relation object, not wrapped in RangeVar + if ('relation' in node && node.relation && typeof node.relation === 'object') { + const relation = node.relation; + if (relation.schemaname === oldSchema) { + relation.schemaname = newSchema; + } + } + + // Handle TypeName nodes (type references like old_schema.my_type) + if ('TypeName' in node) { + const typeName = node.TypeName; + if (Array.isArray(typeName.names) && typeName.names.length >= 2) { + const firstNameNode = typeName.names[0]; + if (firstNameNode?.String?.sval === oldSchema) { + firstNameNode.String.sval = newSchema; + } + } + } + + // Handle FuncCall nodes (function calls like old_schema.my_func()) + if ('FuncCall' in node) { + const funcCall = node.FuncCall; + if (Array.isArray(funcCall.funcname) && funcCall.funcname.length >= 2) { + const firstNameNode = funcCall.funcname[0]; + if (firstNameNode?.String?.sval === oldSchema) { + firstNameNode.String.sval = newSchema; + } + } + } + + // Handle CreateFunctionStmt funcname (CREATE FUNCTION old_schema.my_func) + if ('CreateFunctionStmt' in node) { + const createFunc = node.CreateFunctionStmt; + if (Array.isArray(createFunc.funcname) && createFunc.funcname.length >= 2) { + const firstNameNode = createFunc.funcname[0]; + if (firstNameNode?.String?.sval === oldSchema) { + firstNameNode.String.sval = newSchema; + } + } + } + + // Handle direct type references (returnType in CreateFunctionStmt) + if ('names' in node && 'typemod' in node && Array.isArray(node.names) && node.names.length >= 2) { + const firstNameNode = node.names[0]; + if (firstNameNode?.String?.sval === oldSchema) { + firstNameNode.String.sval = newSchema; + } + } + + // Recurse into all object properties + for (const value of Object.values(node)) { + transformSchemaInSqlAst(value, oldSchema, newSchema); + } + } + + /** + * Transform schema names in hydrated PL/pgSQL AST. + * Walks through PLpgSQL_expr nodes and transforms embedded SQL ASTs. + */ + function transformSchemaInPlpgsqlAst( + node: any, + oldSchema: string, + newSchema: string + ): void { + if (node === null || node === undefined || typeof node !== 'object') { + return; + } + + if (Array.isArray(node)) { + for (const item of node) { + transformSchemaInPlpgsqlAst(item, oldSchema, newSchema); + } + return; + } + + // Handle PLpgSQL_expr nodes with hydrated queries + if ('PLpgSQL_expr' in node) { + const expr = node.PLpgSQL_expr; + const query = expr.query; + + if (query && typeof query === 'object' && 'kind' in query) { + // Handle sql-stmt kind (full SQL statements like SELECT, INSERT) + if (query.kind === 'sql-stmt' && query.parseResult) { + transformSchemaInSqlAst(query.parseResult, oldSchema, newSchema); + } + + // Handle sql-expr kind (SQL expressions like function calls) + if (query.kind === 'sql-expr' && query.expr) { + transformSchemaInSqlAst(query.expr, oldSchema, newSchema); + } + + // Handle assign kind (assignments like var := expr) + if (query.kind === 'assign') { + if (query.targetExpr) { + transformSchemaInSqlAst(query.targetExpr, oldSchema, newSchema); + } + if (query.valueExpr) { + transformSchemaInSqlAst(query.valueExpr, oldSchema, newSchema); + } + } + } + } + + // Handle PLpgSQL_type nodes (variable type declarations) + if ('PLpgSQL_type' in node) { + const plType = node.PLpgSQL_type; + if (plType.typname && plType.typname.startsWith(oldSchema + '.')) { + plType.typname = plType.typname.replace(oldSchema + '.', newSchema + '.'); + } + } + + // Recurse into all object properties + for (const value of Object.values(node)) { + transformSchemaInPlpgsqlAst(value, oldSchema, newSchema); + } + } + + it('should transform schema names in a simple PL/pgSQL function', async () => { + // Simple function with schema-qualified table reference in the body + const sql = ` + CREATE FUNCTION old_schema.get_user_count() + RETURNS int + LANGUAGE plpgsql + AS $$ + DECLARE + user_count int; + BEGIN + SELECT count(*) INTO user_count FROM old_schema.users; + RETURN user_count; + END$$; + `; + + // Step 1: Parse the SQL (includes PL/pgSQL parsing) + const sqlParsed = parseSync(sql) as any; + const plpgsqlParsed = parsePlPgSQLSync(sql) as unknown as PLpgSQLParseResult; + + // Step 2: Hydrate the PL/pgSQL AST (parses embedded SQL into AST nodes) + const { ast: hydratedAst, stats } = hydratePlpgsqlAst(plpgsqlParsed); + + // Verify hydration worked + expect(stats.parsedExpressions).toBeGreaterThan(0); + expect(stats.failedExpressions).toBe(0); + + // Step 3: Transform schema names in both ASTs + const oldSchema = 'old_schema'; + const newSchema = 'new_schema'; + + // Transform outer SQL AST (CreateFunctionStmt) + transformSchemaInSqlAst(sqlParsed, oldSchema, newSchema); + + // Transform PL/pgSQL AST (embedded SQL in function body) + transformSchemaInPlpgsqlAst(hydratedAst, oldSchema, newSchema); + + // Step 4: Dehydrate the PL/pgSQL AST (converts AST back to strings) + const dehydratedAst = dehydratePlpgsqlAst(hydratedAst); + + // Step 5: Deparse the PL/pgSQL body + const newBody = deparseSync(dehydratedAst); + + // Step 6: Stitch the new body back into the SQL AST + const createFunctionStmt = sqlParsed.stmts[0].stmt.CreateFunctionStmt; + const asOption = createFunctionStmt.options.find( + (opt: any) => opt.DefElem?.defname === 'as' + ); + if (asOption?.DefElem?.arg?.List?.items?.[0]?.String) { + asOption.DefElem.arg.List.items[0].String.sval = newBody; + } + + // Step 7: Deparse the full SQL statement + const output = Deparser.deparse(sqlParsed.stmts[0].stmt); + + // Verify transformations + // Function name should be transformed + expect(output).toContain('new_schema.get_user_count'); + expect(output).not.toContain('old_schema.get_user_count'); + + // Table reference in SELECT should be transformed + expect(output).toContain('new_schema.users'); + expect(output).not.toContain('old_schema.users'); + + // Verify the output is valid SQL by re-parsing + const reparsed = parseSync(output); + expect(reparsed.stmts).toHaveLength(1); + expect(reparsed.stmts[0].stmt).toHaveProperty('CreateFunctionStmt'); + }); + + it('should transform schema names in trigger functions', async () => { + const sql = ` + CREATE FUNCTION old_schema.audit_trigger() + RETURNS trigger + LANGUAGE plpgsql + AS $$ + BEGIN + INSERT INTO old_schema.audit_log (table_name, action) + VALUES (TG_TABLE_NAME, TG_OP); + RETURN NEW; + END$$; + `; + + const sqlParsed = parseSync(sql) as any; + const plpgsqlParsed = parsePlPgSQLSync(sql) as unknown as PLpgSQLParseResult; + const { ast: hydratedAst } = hydratePlpgsqlAst(plpgsqlParsed); + + const oldSchema = 'old_schema'; + const newSchema = 'audit_schema'; + + transformSchemaInSqlAst(sqlParsed, oldSchema, newSchema); + transformSchemaInPlpgsqlAst(hydratedAst, oldSchema, newSchema); + + const dehydratedAst = dehydratePlpgsqlAst(hydratedAst); + const newBody = deparseSync(dehydratedAst); + + const createFunctionStmt = sqlParsed.stmts[0].stmt.CreateFunctionStmt; + const asOption = createFunctionStmt.options.find( + (opt: any) => opt.DefElem?.defname === 'as' + ); + if (asOption?.DefElem?.arg?.List?.items?.[0]?.String) { + asOption.DefElem.arg.List.items[0].String.sval = newBody; + } + + const output = Deparser.deparse(sqlParsed.stmts[0].stmt); + + expect(output).toContain('audit_schema.audit_trigger'); + expect(output).toContain('audit_schema.audit_log'); + expect(output).not.toContain('old_schema'); + + // Verify valid SQL + const reparsed = parseSync(output); + expect(reparsed.stmts).toHaveLength(1); + }); + + it('should transform schema names in RETURN QUERY functions', async () => { + const sql = ` + CREATE FUNCTION app_public.get_active_users() + RETURNS SETOF int + LANGUAGE plpgsql + AS $$ + BEGIN + RETURN QUERY SELECT id FROM app_public.users WHERE is_active = true; + RETURN; + END$$; + `; + + const sqlParsed = parseSync(sql) as any; + const plpgsqlParsed = parsePlPgSQLSync(sql) as unknown as PLpgSQLParseResult; + const { ast: hydratedAst } = hydratePlpgsqlAst(plpgsqlParsed); + + transformSchemaInSqlAst(sqlParsed, 'app_public', 'myapp_public'); + transformSchemaInPlpgsqlAst(hydratedAst, 'app_public', 'myapp_public'); + + const dehydratedAst = dehydratePlpgsqlAst(hydratedAst); + const newBody = deparseSync(dehydratedAst); + + const createFunctionStmt = sqlParsed.stmts[0].stmt.CreateFunctionStmt; + const asOption = createFunctionStmt.options.find( + (opt: any) => opt.DefElem?.defname === 'as' + ); + if (asOption?.DefElem?.arg?.List?.items?.[0]?.String) { + asOption.DefElem.arg.List.items[0].String.sval = newBody; + } + + const output = Deparser.deparse(sqlParsed.stmts[0].stmt); + + // All app_public references should be transformed + expect(output).toContain('myapp_public.get_active_users'); + expect(output).toContain('myapp_public.users'); + // Use regex with word boundary to avoid matching 'app_public' inside 'myapp_public' + expect(output).not.toMatch(/\bapp_public\./) + + // Verify valid SQL + const reparsed = parseSync(output); + expect(reparsed.stmts).toHaveLength(1); + }); + + it('should transform function calls inside PL/pgSQL expressions', async () => { + const sql = ` + CREATE FUNCTION old_schema.calculate_total(p_amount numeric) + RETURNS numeric + LANGUAGE plpgsql + AS $$ + DECLARE + tax_rate numeric; + BEGIN + tax_rate := old_schema.get_tax_rate(); + RETURN p_amount * (1 + tax_rate); + END$$; + `; + + const sqlParsed = parseSync(sql) as any; + const plpgsqlParsed = parsePlPgSQLSync(sql) as unknown as PLpgSQLParseResult; + const { ast: hydratedAst } = hydratePlpgsqlAst(plpgsqlParsed); + + transformSchemaInSqlAst(sqlParsed, 'old_schema', 'billing_schema'); + transformSchemaInPlpgsqlAst(hydratedAst, 'old_schema', 'billing_schema'); + + const dehydratedAst = dehydratePlpgsqlAst(hydratedAst); + const newBody = deparseSync(dehydratedAst); + + const createFunctionStmt = sqlParsed.stmts[0].stmt.CreateFunctionStmt; + const asOption = createFunctionStmt.options.find( + (opt: any) => opt.DefElem?.defname === 'as' + ); + if (asOption?.DefElem?.arg?.List?.items?.[0]?.String) { + asOption.DefElem.arg.List.items[0].String.sval = newBody; + } + + const output = Deparser.deparse(sqlParsed.stmts[0].stmt); + + // Function name should be transformed + expect(output).toContain('billing_schema.calculate_total'); + + // Function call in assignment should be transformed + expect(output).toContain('billing_schema.get_tax_rate'); + expect(output).not.toContain('old_schema'); + + // Verify valid SQL + const reparsed = parseSync(output); + expect(reparsed.stmts).toHaveLength(1); + }); +}); From 08d5a57d1215af222891f2d846aecadab332b849 Mon Sep 17 00:00:00 2001 From: Dan Lynch Date: Tue, 6 Jan 2026 17:09:05 +0000 Subject: [PATCH 2/2] test(plpgsql-deparser): add schema rename mapped snapshot test Adds comprehensive schema rename testing with: - SQL fixture file with 12 complex procedures covering various patterns (SELECT, INSERT, UPDATE, DELETE, RETURN QUERY, CTEs, subqueries, etc.) - Snapshot test that tracks all schema references found during traversal - Snapshot of final transformed SQL output The rename map snapshot shows 41 total references across 3 schemas: - app_public: 39 references (functions, tables, relations, func_calls) - app_private: 1 reference - app_internal: 1 reference All references are correctly transformed in the output SQL. --- .../plpgsql/plpgsql_schema_rename.sql | 213 +++++++++ .../schema-rename-mapped.test.ts.snap | 424 +++++++++++++++++ .../__tests__/schema-rename-mapped.test.ts | 437 ++++++++++++++++++ 3 files changed, 1074 insertions(+) create mode 100644 __fixtures__/plpgsql/plpgsql_schema_rename.sql create mode 100644 packages/plpgsql-deparser/__tests__/__snapshots__/schema-rename-mapped.test.ts.snap create mode 100644 packages/plpgsql-deparser/__tests__/schema-rename-mapped.test.ts diff --git a/__fixtures__/plpgsql/plpgsql_schema_rename.sql b/__fixtures__/plpgsql/plpgsql_schema_rename.sql new file mode 100644 index 00000000..784f00b2 --- /dev/null +++ b/__fixtures__/plpgsql/plpgsql_schema_rename.sql @@ -0,0 +1,213 @@ +-- Fixtures to test schema rename traversal +-- These exercise complex scenarios with multiple schema references across different contexts + +-- Test 1: Function with schema-qualified table references in SELECT +CREATE FUNCTION app_public.get_user_stats(p_user_id int) +RETURNS int +LANGUAGE plpgsql AS $$ +DECLARE + total_count int; +BEGIN + SELECT count(*) INTO total_count + FROM app_public.users u + JOIN app_public.orders o ON o.user_id = u.id + WHERE u.id = p_user_id; + RETURN total_count; +END$$; + +-- Test 2: Trigger function with INSERT into schema-qualified table +CREATE FUNCTION app_public.audit_changes() +RETURNS trigger +LANGUAGE plpgsql AS $$ +BEGIN + INSERT INTO app_public.audit_log (table_name, operation, old_data, new_data, changed_at) + VALUES (TG_TABLE_NAME, TG_OP, to_json(OLD), to_json(NEW), now()); + + IF TG_OP = 'DELETE' THEN + RETURN OLD; + END IF; + RETURN NEW; +END$$; + +-- Test 3: Function with UPDATE to schema-qualified table +CREATE FUNCTION app_public.update_user_status(p_user_id int, p_status text) +RETURNS void +LANGUAGE plpgsql AS $$ +BEGIN + UPDATE app_public.users + SET status = p_status, updated_at = now() + WHERE id = p_user_id; + + INSERT INTO app_public.status_history (user_id, status, changed_at) + VALUES (p_user_id, p_status, now()); +END$$; + +-- Test 4: Function with DELETE from schema-qualified table +CREATE FUNCTION app_public.cleanup_old_sessions(p_days int) +RETURNS int +LANGUAGE plpgsql AS $$ +DECLARE + deleted_count int; +BEGIN + DELETE FROM app_public.sessions + WHERE created_at < now() - (p_days || ' days')::interval; + + GET DIAGNOSTICS deleted_count = ROW_COUNT; + RETURN deleted_count; +END$$; + +-- Test 5: SETOF function with RETURN QUERY and schema-qualified tables +CREATE FUNCTION app_public.get_active_orders(p_status text) +RETURNS SETOF int +LANGUAGE plpgsql AS $$ +BEGIN + RETURN QUERY + SELECT o.id + FROM app_public.orders o + JOIN app_public.users u ON u.id = o.user_id + WHERE o.status = p_status + AND u.is_active = true; + RETURN; +END$$; + +-- Test 6: Function with schema-qualified function calls in expressions +CREATE FUNCTION app_public.calculate_order_total(p_order_id int) +RETURNS numeric +LANGUAGE plpgsql AS $$ +DECLARE + subtotal numeric; + tax_amount numeric; + discount numeric; +BEGIN + SELECT sum(quantity * price) INTO subtotal + FROM app_public.order_items + WHERE order_id = p_order_id; + + tax_amount := app_public.get_tax_rate() * subtotal; + discount := app_public.get_discount(p_order_id); + + RETURN subtotal + tax_amount - discount; +END$$; + +-- Test 7: Function with multiple schema references in complex query +CREATE FUNCTION app_public.get_user_dashboard(p_user_id int) +RETURNS TABLE(metric_name text, metric_value numeric) +LANGUAGE plpgsql AS $$ +BEGIN + RETURN QUERY + SELECT 'total_orders'::text, count(*)::numeric + FROM app_public.orders + WHERE user_id = p_user_id + UNION ALL + SELECT 'total_spent'::text, coalesce(sum(total), 0)::numeric + FROM app_public.orders + WHERE user_id = p_user_id + UNION ALL + SELECT 'active_subscriptions'::text, count(*)::numeric + FROM app_public.subscriptions + WHERE user_id = p_user_id AND status = 'active'; + RETURN; +END$$; + +-- Test 8: Trigger function with conditional logic and multiple tables +CREATE FUNCTION app_public.sync_user_profile() +RETURNS trigger +LANGUAGE plpgsql AS $$ +DECLARE + profile_exists boolean; +BEGIN + SELECT EXISTS( + SELECT 1 FROM app_public.profiles WHERE user_id = NEW.id + ) INTO profile_exists; + + IF NOT profile_exists THEN + INSERT INTO app_public.profiles (user_id, created_at) + VALUES (NEW.id, now()); + ELSE + UPDATE app_public.profiles + SET updated_at = now() + WHERE user_id = NEW.id; + END IF; + + PERFORM app_public.notify_profile_change(NEW.id); + RETURN NEW; +END$$; + +-- Test 9: Function with CTE and schema-qualified references +CREATE FUNCTION app_public.get_top_customers(p_limit int) +RETURNS SETOF int +LANGUAGE plpgsql AS $$ +BEGIN + RETURN QUERY + WITH customer_totals AS ( + SELECT user_id, sum(total) as total_spent + FROM app_public.orders + WHERE status = 'completed' + GROUP BY user_id + ) + SELECT ct.user_id + FROM customer_totals ct + JOIN app_public.users u ON u.id = ct.user_id + WHERE u.is_active = true + ORDER BY ct.total_spent DESC + LIMIT p_limit; + RETURN; +END$$; + +-- Test 10: Function with subquery in WHERE clause +CREATE FUNCTION app_public.get_users_with_orders() +RETURNS SETOF int +LANGUAGE plpgsql AS $$ +BEGIN + RETURN QUERY + SELECT u.id + FROM app_public.users u + WHERE EXISTS ( + SELECT 1 FROM app_public.orders o + WHERE o.user_id = u.id + ); + RETURN; +END$$; + +-- Test 11: Function referencing multiple schemas +CREATE FUNCTION app_public.cross_schema_report(p_date date) +RETURNS TABLE(source text, count bigint) +LANGUAGE plpgsql AS $$ +BEGIN + RETURN QUERY + SELECT 'public_users'::text, count(*) + FROM app_public.users + WHERE created_at::date = p_date + UNION ALL + SELECT 'private_logs'::text, count(*) + FROM app_private.activity_logs + WHERE logged_at::date = p_date + UNION ALL + SELECT 'internal_metrics'::text, count(*) + FROM app_internal.metrics + WHERE recorded_at::date = p_date; + RETURN; +END$$; + +-- Test 12: Procedure with schema-qualified references +CREATE PROCEDURE app_public.process_batch(p_batch_id int) +LANGUAGE plpgsql AS $$ +DECLARE + item record; +BEGIN + FOR item IN + SELECT * FROM app_public.batch_items + WHERE batch_id = p_batch_id + LOOP + INSERT INTO app_public.processed_items (item_id, processed_at) + VALUES (item.id, now()); + + UPDATE app_public.batch_items + SET status = 'processed' + WHERE id = item.id; + END LOOP; + + UPDATE app_public.batches + SET status = 'completed', completed_at = now() + WHERE id = p_batch_id; +END$$; diff --git a/packages/plpgsql-deparser/__tests__/__snapshots__/schema-rename-mapped.test.ts.snap b/packages/plpgsql-deparser/__tests__/__snapshots__/schema-rename-mapped.test.ts.snap new file mode 100644 index 00000000..a1cdfed9 --- /dev/null +++ b/packages/plpgsql-deparser/__tests__/__snapshots__/schema-rename-mapped.test.ts.snap @@ -0,0 +1,424 @@ +// Jest Snapshot v1, https://jestjs.io/docs/snapshot-testing + +exports[`schema rename mapped should transform schema names and snapshot rename map and output: rename-map 1`] = ` +{ + "app_internal": { + "newSchema": "myapp_internal_v2", + "referenceCount": 1, + "references": [ + { + "name": "metrics", + "type": "table_ref", + }, + ], + }, + "app_private": { + "newSchema": "myapp_private_v2", + "referenceCount": 1, + "references": [ + { + "name": "activity_logs", + "type": "table_ref", + }, + ], + }, + "app_public": { + "newSchema": "myapp_v2", + "referenceCount": 39, + "references": [ + { + "name": "get_user_stats", + "type": "function_name", + }, + { + "name": "users", + "type": "table_ref", + }, + { + "name": "orders", + "type": "table_ref", + }, + { + "name": "audit_changes", + "type": "function_name", + }, + { + "name": "audit_log", + "type": "relation", + }, + { + "name": "update_user_status", + "type": "function_name", + }, + { + "name": "users", + "type": "relation", + }, + { + "name": "status_history", + "type": "relation", + }, + { + "name": "cleanup_old_sessions", + "type": "function_name", + }, + { + "name": "sessions", + "type": "relation", + }, + { + "name": "get_active_orders", + "type": "function_name", + }, + { + "name": "orders", + "type": "table_ref", + }, + { + "name": "users", + "type": "table_ref", + }, + { + "name": "calculate_order_total", + "type": "function_name", + }, + { + "name": "order_items", + "type": "table_ref", + }, + { + "name": "get_tax_rate", + "type": "func_call", + }, + { + "name": "get_discount", + "type": "func_call", + }, + { + "name": "get_user_dashboard", + "type": "function_name", + }, + { + "name": "orders", + "type": "table_ref", + }, + { + "name": "orders", + "type": "table_ref", + }, + { + "name": "subscriptions", + "type": "table_ref", + }, + { + "name": "sync_user_profile", + "type": "function_name", + }, + { + "name": "profiles", + "type": "table_ref", + }, + { + "name": "profiles", + "type": "relation", + }, + { + "name": "profiles", + "type": "relation", + }, + { + "name": "notify_profile_change", + "type": "func_call", + }, + { + "name": "get_top_customers", + "type": "function_name", + }, + { + "name": "users", + "type": "table_ref", + }, + { + "name": "orders", + "type": "table_ref", + }, + { + "name": "get_users_with_orders", + "type": "function_name", + }, + { + "name": "users", + "type": "table_ref", + }, + { + "name": "orders", + "type": "table_ref", + }, + { + "name": "cross_schema_report", + "type": "function_name", + }, + { + "name": "users", + "type": "table_ref", + }, + { + "name": "process_batch", + "type": "function_name", + }, + { + "name": "processed_items", + "type": "relation", + }, + { + "name": "batch_items", + "type": "relation", + }, + { + "name": "batch_items", + "type": "table_ref", + }, + { + "name": "batches", + "type": "relation", + }, + ], + }, +} +`; + +exports[`schema rename mapped should transform schema names and snapshot rename map and output: transformed-sql 1`] = ` +"CREATE FUNCTION myapp_v2.get_user_stats( + p_user_id int +) RETURNS int LANGUAGE plpgsql AS $$DECLARE + total_count int; +BEGIN + SELECT count(*) INTO total_count + FROM myapp_v2.users AS u + JOIN myapp_v2.orders AS o ON o.user_id = u.id + WHERE + u.id = p_user_id; + RETURN total_count; +END$$; + +CREATE FUNCTION myapp_v2.audit_changes() RETURNS trigger LANGUAGE plpgsql AS $$BEGIN + INSERT INTO myapp_v2.audit_log ( + table_name, + operation, + old_data, + new_data, + changed_at + ) VALUES + ( + tg_table_name, + tg_op, + to_json(old), + to_json(new), + now() + ); + IF tg_op = 'DELETE' THEN + RETURN old; + END IF; + RETURN new; +END$$; + +CREATE FUNCTION myapp_v2.update_user_status( + p_user_id int, + p_status text +) RETURNS void LANGUAGE plpgsql AS $$BEGIN + UPDATE myapp_v2.users SET status = p_status,updated_at = now() WHERE id = p_user_id; + INSERT INTO myapp_v2.status_history ( + user_id, + status, + changed_at + ) VALUES + ( + p_user_id, + p_status, + now() + ); + RETURN; +END$$; + +CREATE FUNCTION myapp_v2.cleanup_old_sessions( + p_days int +) RETURNS int LANGUAGE plpgsql AS $$DECLARE + deleted_count int; +BEGIN + DELETE FROM myapp_v2.sessions WHERE created_at < (now() - CAST(p_days || ' days' AS interval)); + GET DIAGNOSTICS deleted_count = ; + RETURN deleted_count; +END$$; + +CREATE FUNCTION myapp_v2.get_active_orders( + p_status text +) RETURNS SETOF int LANGUAGE plpgsql AS $$BEGIN + RETURN QUERY SELECT o.id + FROM myapp_v2.orders AS o + JOIN myapp_v2.users AS u ON u.id = o.user_id + WHERE + o.status = p_status + AND u.is_active = true; + RETURN; +END$$; + +CREATE FUNCTION myapp_v2.calculate_order_total( + p_order_id int +) RETURNS numeric LANGUAGE plpgsql AS $$DECLARE + subtotal numeric; + tax_amount numeric; + discount numeric; +BEGIN + SELECT sum(quantity * price) INTO subtotal + FROM myapp_v2.order_items + WHERE + order_id = p_order_id; + tax_amount := myapp_v2.get_tax_rate() * subtotal; + discount := myapp_v2.get_discount(p_order_id); + RETURN (subtotal + tax_amount) - discount; +END$$; + +CREATE FUNCTION myapp_v2.get_user_dashboard( + p_user_id int +) RETURNS TABLE ( + metric_name text, + metric_value numeric +) LANGUAGE plpgsql AS $$BEGIN + RETURN QUERY (SELECT + 'total_orders'::text, + (count(*))::numeric + FROM myapp_v2.orders + WHERE + user_id = p_user_id + UNION + ALL + SELECT + 'total_spent'::text, + CAST(COALESCE(sum(total), 0) AS numeric) + FROM myapp_v2.orders + WHERE + user_id = p_user_id) + UNION + ALL + SELECT + 'active_subscriptions'::text, + (count(*))::numeric + FROM myapp_v2.subscriptions + WHERE + user_id = p_user_id + AND status = 'active'; + RETURN; +END$$; + +CREATE FUNCTION myapp_v2.sync_user_profile() RETURNS trigger LANGUAGE plpgsql AS $$DECLARE + profile_exists boolean; +BEGIN + SELECT + EXISTS (SELECT 1 + FROM myapp_v2.profiles + WHERE + user_id = new.id) INTO profile_exists; + IF NOT (profile_exists) THEN + INSERT INTO myapp_v2.profiles ( + user_id, + created_at + ) VALUES + ( + new.id, + now() + ); + ELSE + UPDATE myapp_v2.profiles SET updated_at = now() WHERE user_id = new.id; + END IF; + PERFORM myapp_v2.notify_profile_change(new.id); + RETURN new; +END$$; + +CREATE FUNCTION myapp_v2.get_top_customers( + p_limit int +) RETURNS SETOF int LANGUAGE plpgsql AS $$BEGIN + RETURN QUERY WITH + customer_totals AS (SELECT + user_id, + sum(total) AS total_spent + FROM myapp_v2.orders + WHERE + status = 'completed' + GROUP BY + user_id) + SELECT ct.user_id + FROM customer_totals AS ct + JOIN myapp_v2.users AS u ON u.id = ct.user_id + WHERE + u.is_active = true + ORDER BY + ct.total_spent DESC + LIMIT p_limit; + RETURN; +END$$; + +CREATE FUNCTION myapp_v2.get_users_with_orders() RETURNS SETOF int LANGUAGE plpgsql AS $$BEGIN + RETURN QUERY SELECT u.id + FROM myapp_v2.users AS u + WHERE + EXISTS (SELECT 1 + FROM myapp_v2.orders AS o + WHERE + o.user_id = u.id); + RETURN; +END$$; + +CREATE FUNCTION myapp_v2.cross_schema_report( + p_date date +) RETURNS TABLE ( + source text, + count bigint +) LANGUAGE plpgsql AS $$BEGIN + RETURN QUERY (SELECT + 'public_users'::text, + count(*) + FROM myapp_v2.users + WHERE + created_at::date = p_date + UNION + ALL + SELECT + 'private_logs'::text, + count(*) + FROM myapp_private_v2.activity_logs + WHERE + logged_at::date = p_date) + UNION + ALL + SELECT + 'internal_metrics'::text, + count(*) + FROM myapp_internal_v2.metrics + WHERE + recorded_at::date = p_date; + RETURN; +END$$; + +CREATE PROCEDURE myapp_v2.process_batch( + p_batch_id int +) LANGUAGE plpgsql AS $$DECLARE + item RECORD; +BEGIN + FOR item IN SELECT * + FROM myapp_v2.batch_items + WHERE + batch_id = p_batch_id LOOP + INSERT INTO myapp_v2.processed_items ( + item_id, + processed_at + ) VALUES + ( + item.id, + now() + ); + UPDATE myapp_v2.batch_items SET status = 'processed' WHERE id = item.id; + END LOOP; + UPDATE myapp_v2.batches SET status = 'completed',completed_at = now() WHERE id = p_batch_id; + RETURN; +END$$;" +`; diff --git a/packages/plpgsql-deparser/__tests__/schema-rename-mapped.test.ts b/packages/plpgsql-deparser/__tests__/schema-rename-mapped.test.ts new file mode 100644 index 00000000..9ba6dddc --- /dev/null +++ b/packages/plpgsql-deparser/__tests__/schema-rename-mapped.test.ts @@ -0,0 +1,437 @@ +/** + * Schema Rename Mapped Test + * + * This test demonstrates schema renaming with a rename map that tracks all + * schema references found during traversal. It reads a complex SQL fixture + * file and snapshots both: + * 1. The rename map (all schema references found) + * 2. The final SQL output after transformation + */ + +import { loadModule, parsePlPgSQLSync, parseSync } from '@libpg-query/parser'; +import { Deparser } from 'pgsql-deparser'; +import { hydratePlpgsqlAst, dehydratePlpgsqlAst, PLpgSQLParseResult, deparseSync } from '../src'; +import { readFileSync } from 'fs'; +import * as path from 'path'; + +const FIXTURE_PATH = path.join(__dirname, '../../../__fixtures__/plpgsql/plpgsql_schema_rename.sql'); + +interface SchemaReference { + type: 'function_name' | 'return_type' | 'table_ref' | 'func_call' | 'relation' | 'type_name'; + schema: string; + name: string; + location: string; +} + +interface RenameMap { + [oldSchema: string]: { + newSchema: string; + references: SchemaReference[]; + }; +} + +describe('schema rename mapped', () => { + let fixtureSQL: string; + + beforeAll(async () => { + await loadModule(); + fixtureSQL = readFileSync(FIXTURE_PATH, 'utf-8'); + }); + + /** + * Collect schema references from SQL AST and optionally transform them. + */ + function collectAndTransformSqlAst( + node: any, + renameMap: RenameMap, + location: string + ): void { + if (node === null || node === undefined || typeof node !== 'object') { + return; + } + + if (Array.isArray(node)) { + for (let i = 0; i < node.length; i++) { + collectAndTransformSqlAst(node[i], renameMap, `${location}[${i}]`); + } + return; + } + + // Handle RangeVar nodes (table references like app_public.users) + if ('RangeVar' in node) { + const rangeVar = node.RangeVar; + if (rangeVar.schemaname && renameMap[rangeVar.schemaname]) { + const ref: SchemaReference = { + type: 'table_ref', + schema: rangeVar.schemaname, + name: rangeVar.relname || 'unknown', + location: `${location}.RangeVar`, + }; + renameMap[rangeVar.schemaname].references.push(ref); + rangeVar.schemaname = renameMap[rangeVar.schemaname].newSchema; + } + } + + // Handle direct relation references (INSERT/UPDATE/DELETE statements) + if ('relation' in node && node.relation && typeof node.relation === 'object') { + const relation = node.relation; + if (relation.schemaname && renameMap[relation.schemaname]) { + const ref: SchemaReference = { + type: 'relation', + schema: relation.schemaname, + name: relation.relname || 'unknown', + location: `${location}.relation`, + }; + renameMap[relation.schemaname].references.push(ref); + relation.schemaname = renameMap[relation.schemaname].newSchema; + } + } + + // Handle TypeName nodes (type references) + if ('TypeName' in node) { + const typeName = node.TypeName; + if (Array.isArray(typeName.names) && typeName.names.length >= 2) { + const firstNameNode = typeName.names[0]; + const schemaName = firstNameNode?.String?.sval; + if (schemaName && renameMap[schemaName]) { + const secondNameNode = typeName.names[1]; + const ref: SchemaReference = { + type: 'type_name', + schema: schemaName, + name: secondNameNode?.String?.sval || 'unknown', + location: `${location}.TypeName`, + }; + renameMap[schemaName].references.push(ref); + firstNameNode.String.sval = renameMap[schemaName].newSchema; + } + } + } + + // Handle FuncCall nodes (function calls like app_public.get_tax_rate()) + if ('FuncCall' in node) { + const funcCall = node.FuncCall; + if (Array.isArray(funcCall.funcname) && funcCall.funcname.length >= 2) { + const firstNameNode = funcCall.funcname[0]; + const schemaName = firstNameNode?.String?.sval; + if (schemaName && renameMap[schemaName]) { + const secondNameNode = funcCall.funcname[1]; + const ref: SchemaReference = { + type: 'func_call', + schema: schemaName, + name: secondNameNode?.String?.sval || 'unknown', + location: `${location}.FuncCall`, + }; + renameMap[schemaName].references.push(ref); + firstNameNode.String.sval = renameMap[schemaName].newSchema; + } + } + } + + // Handle CreateFunctionStmt funcname + if ('CreateFunctionStmt' in node) { + const createFunc = node.CreateFunctionStmt; + if (Array.isArray(createFunc.funcname) && createFunc.funcname.length >= 2) { + const firstNameNode = createFunc.funcname[0]; + const schemaName = firstNameNode?.String?.sval; + if (schemaName && renameMap[schemaName]) { + const secondNameNode = createFunc.funcname[1]; + const ref: SchemaReference = { + type: 'function_name', + schema: schemaName, + name: secondNameNode?.String?.sval || 'unknown', + location: `${location}.CreateFunctionStmt.funcname`, + }; + renameMap[schemaName].references.push(ref); + firstNameNode.String.sval = renameMap[schemaName].newSchema; + } + } + } + + // Handle direct type references (returnType in CreateFunctionStmt) + if ('names' in node && 'typemod' in node && Array.isArray(node.names) && node.names.length >= 2) { + const firstNameNode = node.names[0]; + const schemaName = firstNameNode?.String?.sval; + if (schemaName && renameMap[schemaName]) { + const secondNameNode = node.names[1]; + const ref: SchemaReference = { + type: 'return_type', + schema: schemaName, + name: secondNameNode?.String?.sval || 'unknown', + location: `${location}.returnType`, + }; + renameMap[schemaName].references.push(ref); + firstNameNode.String.sval = renameMap[schemaName].newSchema; + } + } + + // Recurse into all object properties + for (const [key, value] of Object.entries(node)) { + collectAndTransformSqlAst(value, renameMap, `${location}.${key}`); + } + } + + /** + * Collect schema references from hydrated PL/pgSQL AST and transform them. + */ + function collectAndTransformPlpgsqlAst( + node: any, + renameMap: RenameMap, + location: string + ): void { + if (node === null || node === undefined || typeof node !== 'object') { + return; + } + + if (Array.isArray(node)) { + for (let i = 0; i < node.length; i++) { + collectAndTransformPlpgsqlAst(node[i], renameMap, `${location}[${i}]`); + } + return; + } + + // Handle PLpgSQL_expr nodes with hydrated queries + if ('PLpgSQL_expr' in node) { + const expr = node.PLpgSQL_expr; + const query = expr.query; + + if (query && typeof query === 'object' && 'kind' in query) { + if (query.kind === 'sql-stmt' && query.parseResult) { + collectAndTransformSqlAst(query.parseResult, renameMap, `${location}.PLpgSQL_expr.query.parseResult`); + } + if (query.kind === 'sql-expr' && query.expr) { + collectAndTransformSqlAst(query.expr, renameMap, `${location}.PLpgSQL_expr.query.expr`); + } + if (query.kind === 'assign') { + if (query.targetExpr) { + collectAndTransformSqlAst(query.targetExpr, renameMap, `${location}.PLpgSQL_expr.query.targetExpr`); + } + if (query.valueExpr) { + collectAndTransformSqlAst(query.valueExpr, renameMap, `${location}.PLpgSQL_expr.query.valueExpr`); + } + } + } + } + + // Handle PLpgSQL_type nodes (variable type declarations) + if ('PLpgSQL_type' in node) { + const plType = node.PLpgSQL_type; + if (plType.typname) { + for (const oldSchema of Object.keys(renameMap)) { + if (plType.typname.startsWith(oldSchema + '.')) { + const typeName = plType.typname.substring(oldSchema.length + 1); + const ref: SchemaReference = { + type: 'type_name', + schema: oldSchema, + name: typeName, + location: `${location}.PLpgSQL_type.typname`, + }; + renameMap[oldSchema].references.push(ref); + plType.typname = renameMap[oldSchema].newSchema + '.' + typeName; + break; + } + } + } + } + + // Recurse into all object properties + for (const [key, value] of Object.entries(node)) { + collectAndTransformPlpgsqlAst(value, renameMap, `${location}.${key}`); + } + } + + /** + * Transform a single SQL statement with schema renaming. + */ + function transformStatement( + sql: string, + renameMap: RenameMap, + stmtIndex: number + ): string { + const sqlParsed = parseSync(sql) as any; + + // Check if this is a PL/pgSQL function/procedure + const stmt = sqlParsed.stmts[0]?.stmt; + const isPlpgsql = stmt?.CreateFunctionStmt?.options?.some( + (opt: any) => opt.DefElem?.defname === 'language' && + opt.DefElem?.arg?.String?.sval?.toLowerCase() === 'plpgsql' + ); + + // Transform outer SQL AST + collectAndTransformSqlAst(sqlParsed, renameMap, `stmt[${stmtIndex}]`); + + if (isPlpgsql) { + try { + const plpgsqlParsed = parsePlPgSQLSync(sql) as unknown as PLpgSQLParseResult; + const { ast: hydratedAst } = hydratePlpgsqlAst(plpgsqlParsed); + + // Transform PL/pgSQL AST + collectAndTransformPlpgsqlAst(hydratedAst, renameMap, `stmt[${stmtIndex}].plpgsql`); + + // Dehydrate and deparse + const dehydratedAst = dehydratePlpgsqlAst(hydratedAst); + const newBody = deparseSync(dehydratedAst); + + // Stitch body back into SQL AST + const createFunctionStmt = sqlParsed.stmts[0].stmt.CreateFunctionStmt; + const asOption = createFunctionStmt?.options?.find( + (opt: any) => opt.DefElem?.defname === 'as' + ); + if (asOption?.DefElem?.arg?.List?.items?.[0]?.String) { + asOption.DefElem.arg.List.items[0].String.sval = newBody; + } + } catch (err) { + // If PL/pgSQL parsing fails, just use the SQL AST transformation + console.warn(`PL/pgSQL parsing failed for statement ${stmtIndex}:`, err); + } + } + + return Deparser.deparse(sqlParsed.stmts[0].stmt); + } + + /** + * Split SQL file into individual statements. + * Handles dollar-quoted strings and skips comment-only lines. + */ + function splitStatements(sql: string): string[] { + const statements: string[] = []; + let current = ''; + let inDollarQuote = false; + let dollarTag = ''; + let inLineComment = false; + + for (let i = 0; i < sql.length; i++) { + const char = sql[i]; + + // Handle line comments + if (!inDollarQuote && char === '-' && sql[i + 1] === '-') { + inLineComment = true; + current += char; + continue; + } + if (inLineComment && char === '\n') { + inLineComment = false; + current += char; + continue; + } + if (inLineComment) { + current += char; + continue; + } + + current += char; + + // Handle dollar quotes + if (char === '$' && !inDollarQuote) { + let tag = '$'; + let j = i + 1; + while (j < sql.length && (sql[j].match(/[a-zA-Z0-9_]/) || sql[j] === '$')) { + tag += sql[j]; + if (sql[j] === '$') { + j++; + break; + } + j++; + } + if (tag.endsWith('$') && tag.length >= 2) { + inDollarQuote = true; + dollarTag = tag; + current += sql.slice(i + 1, j); + i = j - 1; + } + } else if (inDollarQuote && char === '$') { + // Check for closing tag + const remaining = sql.slice(i); + if (remaining.startsWith(dollarTag)) { + current += dollarTag.slice(1); + i += dollarTag.length - 1; + inDollarQuote = false; + dollarTag = ''; + } + } else if (!inDollarQuote && char === ';') { + const trimmed = current.trim(); + // Remove leading comment lines and check if there's actual SQL + const withoutComments = trimmed.replace(/^(--[^\n]*\n\s*)+/, '').trim(); + if (withoutComments.length > 0) { + statements.push(trimmed); + } + current = ''; + } + } + + // Handle last statement without semicolon + const trimmed = current.trim(); + const withoutComments = trimmed.replace(/^(--[^\n]*\n\s*)+/, '').trim(); + if (withoutComments.length > 0) { + statements.push(trimmed); + } + + return statements; + } + + it('should transform schema names and snapshot rename map and output', () => { + // Define the rename map with schemas to transform + const renameMap: RenameMap = { + 'app_public': { + newSchema: 'myapp_v2', + references: [], + }, + 'app_private': { + newSchema: 'myapp_private_v2', + references: [], + }, + 'app_internal': { + newSchema: 'myapp_internal_v2', + references: [], + }, + }; + + // Split fixture into individual statements + const statements = splitStatements(fixtureSQL); + expect(statements.length).toBeGreaterThan(0); + + // Transform each statement + const transformedStatements: string[] = []; + for (let i = 0; i < statements.length; i++) { + const stmt = statements[i]; + // Skip comment-only statements (after removing leading comments) + const withoutLeadingComments = stmt.replace(/^(--[^\n]*\n\s*)+/, '').trim(); + if (!withoutLeadingComments || withoutLeadingComments.startsWith('--')) { + continue; + } + try { + const transformed = transformStatement(stmt, renameMap, i); + transformedStatements.push(transformed); + } catch (err) { + // Log the error for debugging + const errMsg = err instanceof Error ? err.message : String(err); + console.warn(`Failed to transform statement ${i}: ${errMsg}`); + transformedStatements.push(`-- TRANSFORM FAILED: ${stmt.substring(0, 100)}...`); + } + } + + // Create a summary of the rename map (without location details for cleaner snapshot) + const renameMapSummary: Record }> = {}; + for (const [oldSchema, data] of Object.entries(renameMap)) { + renameMapSummary[oldSchema] = { + newSchema: data.newSchema, + referenceCount: data.references.length, + references: data.references.map(r => ({ type: r.type, name: r.name })), + }; + } + + // Snapshot the rename map + expect(renameMapSummary).toMatchSnapshot('rename-map'); + + // Snapshot the transformed SQL + const finalSQL = transformedStatements.join(';\n\n') + ';'; + expect(finalSQL).toMatchSnapshot('transformed-sql'); + + // Verify no old schema references remain in output + expect(finalSQL).not.toMatch(/\bapp_public\./); + expect(finalSQL).not.toMatch(/\bapp_private\./); + expect(finalSQL).not.toMatch(/\bapp_internal\./); + + // Verify new schema references are present + expect(finalSQL).toContain('myapp_v2.'); + }); +});