Merge pull request #302 from constructive-io/feat/plpgsql-parser-multifn-bodies

fix(plpgsql-parser): multi-function scripts no longer all get the first function's body

Author: pyramation

packages/plpgsql-parser/__tests__/plpgsql-parser.test.ts

@@ -91,6 +91,62 @@ describe('plpgsql-parser', () => {
 
       expect(result).toContain('visitor_renamed');
     });
+
+    it('should preserve each function body in multi-function scripts', () => {
+      const multiFunctionSql = `
+CREATE FUNCTION f1() RETURNS void LANGUAGE plpgsql AS $$
+BEGIN
+  RAISE NOTICE 'one';
+END;
+$$;
+
+CREATE FUNCTION f2() RETURNS void LANGUAGE plpgsql AS $$
+BEGIN
+  RAISE NOTICE 'two';
+END;
+$$;
+
+CREATE FUNCTION f3() RETURNS trigger LANGUAGE plpgsql AS $$
+BEGIN
+  IF TG_OP = 'INSERT' THEN
+    RETURN NEW;
+  END IF;
+END;
+$$;
+`;
+      const result = transformSync(multiFunctionSql, () => {});
+
+      expect(result).toContain(`'one'`);
+      expect(result).toContain(`'two'`);
+      expect(result).toContain('RETURN new');
+      expect((result.match(/'one'/g) || []).length).toBe(1);
+      expect((result.match(/'two'/g) || []).length).toBe(1);
+    });
+
+    it('should associate the correct plpgsql AST with each parsed function', () => {
+      const multiFunctionSql = `
+CREATE TABLE t (id int);
+
+CREATE FUNCTION fa() RETURNS void LANGUAGE plpgsql AS $$
+BEGIN
+  RAISE NOTICE 'alpha';
+END;
+$$;
+
+CREATE FUNCTION fb() RETURNS void LANGUAGE plpgsql AS $$
+BEGIN
+  RAISE NOTICE 'beta';
+END;
+$$;
+`;
+      const parsed = parse(multiFunctionSql);
+
+      expect(parsed.functions).toHaveLength(2);
+      expect(JSON.stringify(parsed.functions[0].plpgsql.raw)).toContain('alpha');
+      expect(JSON.stringify(parsed.functions[0].plpgsql.raw)).not.toContain('beta');
+      expect(JSON.stringify(parsed.functions[1].plpgsql.raw)).toContain('beta');
+      expect(JSON.stringify(parsed.functions[1].plpgsql.raw)).not.toContain('alpha');
+    });
   });
 
   describe('deparseSync', () => {

packages/plpgsql-parser/src/parse.ts

@@ -61,7 +61,14 @@ function isPlpgsqlFunction(stmt: any): boolean {
   return language === 'plpgsql';
 }
 
-function extractFunctionInfo(stmt: any, stmtIndex: number, fullSql: string): ParsedFunction | null {
+function getStatementSql(sqlBuffer: Buffer, rawStmt: any): string {
+  const start = rawStmt?.stmt_location ?? 0;
+  const len = rawStmt?.stmt_len;
+  const end = len !== undefined ? start + len : sqlBuffer.length;
+  return sqlBuffer.slice(start, end).toString('utf8');
+}
+
+function extractFunctionInfo(stmt: any, stmtIndex: number, stmtSql: string): ParsedFunction | null {
   const createFunctionStmt = stmt?.CreateFunctionStmt;
   if (!createFunctionStmt) return null;
 
@@ -72,7 +79,10 @@ function extractFunctionInfo(stmt: any, stmtIndex: number, fullSql: string): Par
   if (!body) return null;
 
   try {
-    const plpgsqlRaw = parsePlPgSQLSync(fullSql) as unknown as PLpgSQLParseResult;
+    // Parse only this statement's SQL. Parsing the full script would return
+    // every function's PL/pgSQL AST in plpgsql_funcs, and downstream deparse
+    // pairs each statement with plpgsql_funcs[0].
+    const plpgsqlRaw = parsePlPgSQLSync(stmtSql) as unknown as PLpgSQLParseResult;
     const { ast: hydrated, stats, errors } = hydratePlpgsqlAst(plpgsqlRaw);
 
     return {
@@ -99,14 +109,15 @@ export function parse(sql: string, options: ParseOptions = {}): ParsedScript {
   const sqlResult: ParseResult = parseSqlSync(sql);
   const items: ParsedItem[] = [];
   const functions: ParsedFunction[] = [];
+  const sqlBuffer = Buffer.from(sql, 'utf8');
 
   if (sqlResult.stmts) {
     for (let i = 0; i < sqlResult.stmts.length; i++) {
       const rawStmt = sqlResult.stmts[i];
       const stmt = rawStmt?.stmt;
 
       if (stmt && isPlpgsqlFunction(stmt) && hydrate) {
-        const fnInfo = extractFunctionInfo(stmt, i, sql);
+        const fnInfo = extractFunctionInfo(stmt, i, getStatementSql(sqlBuffer, rawStmt));
         if (fnInfo) {
           items.push(fnInfo);
           functions.push(fnInfo);