diff --git a/packages/plpgsql-parser/__tests__/plpgsql-parser.test.ts b/packages/plpgsql-parser/__tests__/plpgsql-parser.test.ts index 1d738444..9f1b280d 100644 --- a/packages/plpgsql-parser/__tests__/plpgsql-parser.test.ts +++ b/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', () => { diff --git a/packages/plpgsql-parser/src/parse.ts b/packages/plpgsql-parser/src/parse.ts index 22f4f67b..8142399c 100644 --- a/packages/plpgsql-parser/src/parse.ts +++ b/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,6 +109,7 @@ 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++) { @@ -106,7 +117,7 @@ export function parse(sql: string, options: ParseOptions = {}): ParsedScript { 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);