Skip to content

Commit 6c0c2fc

Browse files
pyramationpyramation
committed
feat(plpgsql-parser): add traverse functionality for PL/pgSQL ASTs
- Add walk() function for traversing PL/pgSQL AST nodes with visitor pattern - Add walkParsedScript() convenience function for walking both SQL and PL/pgSQL - Support automatic recursion into hydrated SQL expressions via @pgsql/traverse - Add PLpgSQLNodePath class for path tracking during traversal - Add comprehensive tests for traverse functionality
1 parent 5676bd8 commit 6c0c2fc

5 files changed

Lines changed: 3586 additions & 5565 deletions

File tree

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

Lines changed: 247 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,247 @@
1+
import { parse, walk, walkParsedScript, PLpgSQLNodePath, loadModule } from '../src';
2+
import type { PLpgSQLVisitor } from '../src';
3+
4+
describe('plpgsql-parser traverse', () => {
5+
beforeAll(async () => {
6+
await loadModule();
7+
});
8+
9+
const simpleFunctionSql = `
10+
CREATE FUNCTION test_func(p_id int)
11+
RETURNS void
12+
LANGUAGE plpgsql
13+
AS $$
14+
DECLARE
15+
v_name text;
16+
BEGIN
17+
SELECT name INTO v_name FROM users WHERE id = p_id;
18+
RAISE NOTICE 'Hello %', v_name;
19+
END;
20+
$$;
21+
`;
22+
23+
describe('walk', () => {
24+
it('should visit PL/pgSQL nodes with a visitor object', () => {
25+
const parsed = parse(simpleFunctionSql);
26+
expect(parsed.functions.length).toBe(1);
27+
28+
const visitedTags: string[] = [];
29+
const visitor: PLpgSQLVisitor = {
30+
PLpgSQL_function: (path) => {
31+
visitedTags.push(path.tag);
32+
},
33+
PLpgSQL_stmt_block: (path) => {
34+
visitedTags.push(path.tag);
35+
},
36+
PLpgSQL_var: (path) => {
37+
visitedTags.push(path.tag);
38+
},
39+
PLpgSQL_stmt_execsql: (path) => {
40+
visitedTags.push(path.tag);
41+
},
42+
PLpgSQL_stmt_raise: (path) => {
43+
visitedTags.push(path.tag);
44+
},
45+
};
46+
47+
walk(parsed.functions[0].plpgsql.hydrated, visitor);
48+
49+
expect(visitedTags).toContain('PLpgSQL_function');
50+
expect(visitedTags).toContain('PLpgSQL_stmt_block');
51+
expect(visitedTags).toContain('PLpgSQL_var');
52+
expect(visitedTags).toContain('PLpgSQL_stmt_execsql');
53+
expect(visitedTags).toContain('PLpgSQL_stmt_raise');
54+
});
55+
56+
it('should visit PL/pgSQL nodes with a walker function', () => {
57+
const parsed = parse(simpleFunctionSql);
58+
59+
const visitedTags: string[] = [];
60+
walk(parsed.functions[0].plpgsql.hydrated, (path: PLpgSQLNodePath) => {
61+
visitedTags.push(path.tag);
62+
});
63+
64+
expect(visitedTags.length).toBeGreaterThan(0);
65+
expect(visitedTags).toContain('PLpgSQL_function');
66+
});
67+
68+
it('should provide correct path information', () => {
69+
const parsed = parse(simpleFunctionSql);
70+
71+
let blockPath: (string | number)[] = [];
72+
const visitor: PLpgSQLVisitor = {
73+
PLpgSQL_stmt_block: (path) => {
74+
blockPath = path.path;
75+
},
76+
};
77+
78+
walk(parsed.functions[0].plpgsql.hydrated, visitor);
79+
80+
expect(blockPath).toContain('action');
81+
});
82+
83+
it('should allow skipping children by returning false', () => {
84+
const parsed = parse(simpleFunctionSql);
85+
86+
const visitedTags: string[] = [];
87+
const visitor: PLpgSQLVisitor = {
88+
PLpgSQL_function: (path) => {
89+
visitedTags.push(path.tag);
90+
return false; // Skip children
91+
},
92+
PLpgSQL_stmt_block: (path) => {
93+
visitedTags.push(path.tag);
94+
},
95+
};
96+
97+
walk(parsed.functions[0].plpgsql.hydrated, visitor);
98+
99+
expect(visitedTags).toContain('PLpgSQL_function');
100+
expect(visitedTags).not.toContain('PLpgSQL_stmt_block');
101+
});
102+
});
103+
104+
describe('walkParsedScript', () => {
105+
it('should walk both SQL and PL/pgSQL nodes', () => {
106+
const parsed = parse(simpleFunctionSql);
107+
108+
const plpgsqlTags: string[] = [];
109+
const sqlTags: string[] = [];
110+
111+
walkParsedScript(
112+
parsed,
113+
{
114+
PLpgSQL_function: (path) => {
115+
plpgsqlTags.push(path.tag);
116+
},
117+
PLpgSQL_stmt_block: (path) => {
118+
plpgsqlTags.push(path.tag);
119+
},
120+
},
121+
{
122+
CreateFunctionStmt: (path) => {
123+
sqlTags.push(path.tag);
124+
},
125+
}
126+
);
127+
128+
expect(plpgsqlTags).toContain('PLpgSQL_function');
129+
expect(sqlTags).toContain('CreateFunctionStmt');
130+
});
131+
});
132+
133+
describe('SQL expression traversal', () => {
134+
it('should traverse into hydrated SQL expressions when sqlVisitor is provided', () => {
135+
const parsed = parse(simpleFunctionSql);
136+
137+
const sqlTags: string[] = [];
138+
const visitor: PLpgSQLVisitor = {
139+
PLpgSQL_expr: () => {
140+
// Just visit the expression node
141+
},
142+
};
143+
144+
walk(parsed.functions[0].plpgsql.hydrated, visitor, {
145+
walkSqlExpressions: true,
146+
sqlVisitor: {
147+
SelectStmt: (path) => {
148+
sqlTags.push(path.tag);
149+
},
150+
RangeVar: (path) => {
151+
sqlTags.push(path.tag);
152+
},
153+
},
154+
});
155+
156+
// The SELECT statement inside the function should be visited
157+
expect(sqlTags).toContain('SelectStmt');
158+
expect(sqlTags).toContain('RangeVar');
159+
});
160+
});
161+
162+
describe('control flow statements', () => {
163+
it('should traverse IF statements', () => {
164+
const ifFunctionSql = `
165+
CREATE FUNCTION test_if(p_val int)
166+
RETURNS text
167+
LANGUAGE plpgsql
168+
AS $$
169+
BEGIN
170+
IF p_val > 10 THEN
171+
RETURN 'big';
172+
ELSIF p_val > 5 THEN
173+
RETURN 'medium';
174+
ELSE
175+
RETURN 'small';
176+
END IF;
177+
END;
178+
$$;
179+
`;
180+
181+
const parsed = parse(ifFunctionSql);
182+
const visitedTags: string[] = [];
183+
184+
walk(parsed.functions[0].plpgsql.hydrated, (path: PLpgSQLNodePath) => {
185+
visitedTags.push(path.tag);
186+
});
187+
188+
expect(visitedTags).toContain('PLpgSQL_stmt_if');
189+
expect(visitedTags).toContain('PLpgSQL_if_elsif');
190+
expect(visitedTags).toContain('PLpgSQL_stmt_return');
191+
});
192+
193+
it('should traverse LOOP statements', () => {
194+
const loopFunctionSql = `
195+
CREATE FUNCTION test_loop()
196+
RETURNS void
197+
LANGUAGE plpgsql
198+
AS $$
199+
DECLARE
200+
i int := 0;
201+
BEGIN
202+
WHILE i < 10 LOOP
203+
i := i + 1;
204+
END LOOP;
205+
END;
206+
$$;
207+
`;
208+
209+
const parsed = parse(loopFunctionSql);
210+
const visitedTags: string[] = [];
211+
212+
walk(parsed.functions[0].plpgsql.hydrated, (path: PLpgSQLNodePath) => {
213+
visitedTags.push(path.tag);
214+
});
215+
216+
expect(visitedTags).toContain('PLpgSQL_stmt_while');
217+
expect(visitedTags).toContain('PLpgSQL_stmt_assign');
218+
});
219+
220+
it('should traverse FOR loops', () => {
221+
const forFunctionSql = `
222+
CREATE FUNCTION test_for()
223+
RETURNS void
224+
LANGUAGE plpgsql
225+
AS $$
226+
DECLARE
227+
rec record;
228+
BEGIN
229+
FOR rec IN SELECT * FROM users LOOP
230+
RAISE NOTICE '%', rec.name;
231+
END LOOP;
232+
END;
233+
$$;
234+
`;
235+
236+
const parsed = parse(forFunctionSql);
237+
const visitedTags: string[] = [];
238+
239+
walk(parsed.functions[0].plpgsql.hydrated, (path: PLpgSQLNodePath) => {
240+
visitedTags.push(path.tag);
241+
});
242+
243+
expect(visitedTags).toContain('PLpgSQL_stmt_fors');
244+
expect(visitedTags).toContain('PLpgSQL_stmt_raise');
245+
});
246+
});
247+
});

‎packages/plpgsql-parser/package.json‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,7 @@
4747
},
4848
"dependencies": {
4949
"@libpg-query/parser": "^17.6.3",
50+
"@pgsql/traverse": "workspace:*",
5051
"@pgsql/types": "^17.6.2",
5152
"pgsql-deparser": "workspace:*",
5253
"plpgsql-deparser": "workspace:*"

‎packages/plpgsql-parser/src/index.ts‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,15 @@ export * from './types';
22
export { parse, parseSync, loadModule } from './parse';
33
export { deparse, deparseSync } from './deparse';
44
export { transform, transformSync } from './transform';
5+
export {
6+
walk,
7+
walkParsedScript,
8+
PLpgSQLNodePath,
9+
type PLpgSQLWalker,
10+
type PLpgSQLVisitor,
11+
type PLpgSQLNodeTag,
12+
type WalkOptions
13+
} from './traverse';
514

615
export {
716
hydratePlpgsqlAst,

0 commit comments

Comments
 (0)