diff --git a/.agents/skills/ast-traversal/SKILL.md b/.agents/skills/ast-traversal/SKILL.md new file mode 100644 index 00000000..0eca0939 --- /dev/null +++ b/.agents/skills/ast-traversal/SKILL.md @@ -0,0 +1,196 @@ +--- +name: ast-traversal +description: How to walk PostgreSQL SQL and PL/pgSQL ASTs in this monorepo — choosing between walk, walkSql, walkSqlAst, walkPlpgsqlAst, and traverse; statement context, visitor composition, abort, and mutation. Use when reading, validating, or rewriting SQL/PL/pgSQL ASTs. +--- + +# AST Traversal + +Everything in this repo that inspects or rewrites SQL goes through one of five +functions. Pick the right one first; the rest of the work follows from it. + +## Choosing a function + +| Function | Package | Walks | Mutates | +|---|---|---|---| +| `walk(ast, visitors, opts?)` | `@pgsql/traverse` | any AST — parsed script, `ParseResult`, SQL node, PL/pgSQL node | no | +| `walkSql(text, visitors, opts?)` | `plpgsql-parser` | SQL **text**: parses, hydrates PL/pgSQL bodies, then `walk`s | no | +| `walkSqlAst(ast, visitor)` | `@pgsql/traverse` | SQL AST only | no | +| `walkPlpgsqlAst(ast, visitor, opts?)` | `@pgsql/traverse` | PL/pgSQL AST only | no | +| `traverse(ast, mutableVisitor)` | `@pgsql/traverse` | SQL AST, with insert/remove/replace | yes | + +Decision rules: + +- **Have a string?** `walkSql`. It is the only entry point that parses. +- **Have an AST and don't care which universe it is from?** `walk`. +- **Want exactly one node universe and no statement context?** `walkSqlAst` or + `walkPlpgsqlAst`. These are the primitives `walk` is built on; reach for them + when writing a reusable visitor that some other walker will drive (this is why + `@pgsql/transform` exports visitor *factories* rather than walkers). +- **Need to change the tree structurally?** `traverse`. Note that read-only + walkers hand you the real node objects, so field-level edits (renaming a + schema, rewriting a name list) work under `walk` too — `traverse` is for + inserting, removing, and replacing nodes. + +Never hand-roll `transformSync(sql, ..., { hydrate: true })` plus a +per-statement loop plus a PL/pgSQL walk. That harness *is* `walk`. + +## Visitors + +A visitor is either a callback (fires on every node) or an object keyed by node +tag. SQL tags and `PLpgSQL_*` tags may be mixed in one object — `walk` routes +each node to the right walker. + +```ts +import { walk } from '@pgsql/traverse'; + +walk(ast, { + RangeVar: (path) => console.log(path.node.relname), + PLpgSQL_stmt_dynexecute: (path) => console.log('dynamic EXECUTE', path.node) +}); +``` + +Pass an **array** of visitors to run independent concerns in a single parse: + +```ts +walkSql(sql, [blockedSchemas, readOnlySchemas, blockedFunctions]); +``` + +Each callback receives `(path, ctx)`: + +- `path` — a `NodePath` (SQL) or `PlpgsqlNodePath` (PL/pgSQL): `tag`, `node`, + `parent`, `keyPath`. This is *structure*: where the node sits. +- `ctx` — a `WalkContext`: `stmtTag`, `stmtIndex`, `isWrite`, `isRead`, + `insideFunction`, `functionName`, `abort()`. This is *situation*: what the + node is part of. + +The reserved `statement` key fires once per top-level statement, before its +children — the hook for per-statement setup or classification. + +## Control flow + +Two distinct mechanisms, do not confuse them: + +- `return false` — skip this node's children, keep walking siblings. +- `ctx.abort(reason?)` — end the whole walk. `walk`/`walkSql` return + `{ aborted, reason, reasons }`. This is what a validator wants: the first + rejection ends the work. + +```ts +const result = walkSql(sql, { + RangeVar: (path, ctx) => { + if (ctx.isWrite && path.node.schemaname === 'audit') { + ctx.abort(`cannot write to ${path.node.schemaname}`); + } + } +}); +if (result.aborted) reject(result.reason); +``` + +Unparseable input from `walkSql` comes back as `{ aborted: true, reason }` +rather than a thrown error, so "rejected" and "not understood" are one code path. + +## Worked examples + +### Collect every table a script touches, including inside function bodies + +```ts +import { loadModule, walkSql } from 'plpgsql-parser'; + +await loadModule(); // once per process: libpg-query is WASM + +const tables = new Set(); +walkSql(sql, { + RangeVar: (path, ctx) => { + const name = path.node.schemaname + ? `${path.node.schemaname}.${path.node.relname}` + : path.node.relname; + tables.add(ctx.insideFunction ? `${name} (via ${ctx.functionName})` : name); + } +}); +``` + +### Classify statements without a second parse + +```ts +walkSql(sql, { + statement: (path, ctx) => { + console.log(ctx.stmtIndex, path.tag, ctx.isWrite ? 'write' : 'read'); + } +}); +``` + +### Rewrite schema names on a parsed script + +Field-level rewrites work through the read-only walker because `path.node` is the +live node. This is exactly how `@pgsql/transform` renames schemas: + +```ts +import { transformSync } from 'plpgsql-parser'; +import { walk } from '@pgsql/traverse'; + +const out = transformSync(sql, (ctx) => { + walk(ctx, { + RangeVar: (path) => { + const to = mapping.get(path.node.schemaname); + if (to) path.node.schemaname = to; + } + }); +}, { hydrate: true, pretty: true }); +``` + +`transformSync` gives the callback a parsed script (`{ sql, functions }`), which +`walk` dispatches over directly — statements first, then every hydrated body. + +### Restructure the tree + +```ts +import { traverse } from '@pgsql/traverse'; + +traverse(ast, { + RawStmt: { + enter: (path) => { + if (isRedundant(path.node)) path.remove(); + } + } +}); +``` + +### Drive a reusable visitor from a primitive + +When a caller already owns the traversal loop (per-statement state, custom +ordering), build the visitor separately and hand it to the primitive: + +```ts +import { walkSqlAst } from '@pgsql/traverse'; + +for (const stmt of parseResult.stmts) { + walkSqlAst(stmt.stmt, createFactsVisitor(factsFor(stmt))); +} +``` + +## Options + +`walk` and `walkSql` share: + +- `walkFunctionBodies` (default `true`) — walk hydrated PL/pgSQL bodies. Under + `walkSql`, `false` also skips the PL/pgSQL parse, so turn it off when the + visitors only care about top-level SQL. +- `walkSqlExpressions` (default `true`) — recurse into the SQL expressions inside + those bodies. +- `sqlVisitor` — override the visitor used for those SQL expressions. + +## Gotchas + +- **Call `loadModule()` before any parse.** `libpg-query` is WASM; `parseSync` / + `walkSql` throw `WASM module not initialized` otherwise. +- **PL/pgSQL bodies are opaque until hydrated.** Without `{ hydrate: true }` a + function body is a query string, and no SQL visitor will ever fire inside it. +- **Untagged typed fields exist.** `CreatePolicyStmt.table` is a bare `RangeVar` + with no `{ RangeVar: ... }` wrapper. `walkSqlAst` handles this via the runtime + schema, which is why hand-written recursion over `Object.keys` misses nodes. +- **`@pgsql/traverse` must stay parser-free.** It depends on types only; adding a + parser dependency would cycle and pull WASM into the leaf package. Anything + that needs to parse belongs in `plpgsql-parser` or above. +- **Fixtures are the regression gate for transform work.** After changing any + traversal in `@pgsql/transform`, `pnpm --filter @pgsql/transform test` output + must be byte-identical; see the `testing-fixtures` skill. diff --git a/.agents/skills/code-generation/SKILL.md b/.agents/skills/code-generation/SKILL.md index 6e30b8b4..36481973 100644 --- a/.agents/skills/code-generation/SKILL.md +++ b/.agents/skills/code-generation/SKILL.md @@ -13,7 +13,7 @@ Three packages generate TypeScript from the PostgreSQL protobuf definition at `_ | Package | Script | What it generates | |---------|--------|-------------------| | `@pgsql/utils` | `npm run build:proto` | AST helper functions (`src/`), wrapped helpers (`wrapped.ts`), runtime schema (`runtime-schema.ts`) | -| `@pgsql/traverse` | `npm run build:proto` | Visitor-pattern traversal utilities | +| `@pgsql/traverse` | `npm run build:proto` | Runtime schema driving the SQL walker (`walk` / `walkSqlAst`); see the `ast-traversal` skill | | `@pgsql/transform-ast` | `npm run build:proto` | Multi-version AST transformer utilities | Each package has a `scripts/pg-proto-parser.ts` that configures `PgProtoParser` with package-specific options (which features to enable, output paths, type sources). diff --git a/.agents/skills/testing-fixtures/SKILL.md b/.agents/skills/testing-fixtures/SKILL.md index f268e9f4..1fa0fbe9 100644 --- a/.agents/skills/testing-fixtures/SKILL.md +++ b/.agents/skills/testing-fixtures/SKILL.md @@ -168,6 +168,30 @@ npm run kitchen-sink # generate transform-specific kitchen-sink tests npm run test:ast # run AST round-trip validation ``` +## Changing a Walker + +`@pgsql/transform` drives its whole pipeline (schema mapping, routing, qualify, +role and extension transforms, round-trip validation) through the walkers in +`@pgsql/traverse` — see the `ast-traversal` skill. Its fixtures are therefore the +regression gate for any traversal change, including ones made in another package: + +```bash +pnpm --filter @pgsql/traverse test +pnpm --filter plpgsql-parser test +pnpm --filter @pgsql/transform test # fixtures + snapshots must be unchanged +``` + +A traversal change that alters fixture output is a behavior change, not a +refactor. Two failure modes to look for specifically: + +- **Nodes no longer reached** — untagged typed fields (e.g. `CreatePolicyStmt.table`) + are only found via the runtime schema, so a hand-rolled recursion silently drops + them and a fixture loses a rename. +- **Nodes reached twice** — a visitor fired both by an outer walk and by a nested + one double-applies edits. + +Never edit a fixture or snapshot to make a walker change pass. + ## Package Scripts Reference ### `packages/deparser` (primary fixture pipeline) diff --git a/.github/workflows/run-tests.yaml b/.github/workflows/run-tests.yaml index f1932455..9ccb7439 100644 --- a/.github/workflows/run-tests.yaml +++ b/.github/workflows/run-tests.yaml @@ -23,6 +23,7 @@ jobs: - pg-proto-parser - '@pgsql/quotes' - '@pgsql/transform-ast' + - '@pgsql/traverse' - '@pgsql/transform' - '@pgsql/scripts' steps: diff --git a/AGENTS.md b/AGENTS.md index c71b42b7..de23ca42 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -10,12 +10,12 @@ A pnpm monorepo for PostgreSQL AST parsing, deparsing, and code generation. All |---------|-----------|---------| | `pgsql-parser` | `packages/parser` | Parse SQL to AST (wraps `libpg-query` WASM) | | `pgsql-deparser` | `packages/deparser` | Convert AST back to SQL (pure TypeScript) | -| `plpgsql-parser` | `packages/plpgsql-parser` | Parse PL/pgSQL to AST | +| `plpgsql-parser` | `packages/plpgsql-parser` | Parse PL/pgSQL to AST; `walkSql(text, ...)` for text-in traversal | | `plpgsql-deparser` | `packages/plpgsql-deparser` | Convert PL/pgSQL AST back to SQL | | `pgsql-types` | `packages/pgsql-types` | Narrowed TypeScript types inferred from SQL fixtures | | `@pgsql/types` | (published from proto-parser codegen) | Core TypeScript type definitions for PostgreSQL AST nodes | | `@pgsql/utils` | `packages/utils` | Type-safe AST node creation utilities | -| `@pgsql/traverse` | `packages/traverse` | Visitor-pattern AST traversal | +| `@pgsql/traverse` | `packages/traverse` | Visitor-pattern traversal of SQL and PL/pgSQL ASTs: `walk`, `walkSqlAst`, `walkPlpgsqlAst`, `traverse` | | `@pgsql/transform-ast` | `packages/transform-ast` | Multi-version AST transformer (PG 13-17) | | `@pgsql/transform` | `packages/transform` | SQL schema transformation, statement classification (AST facts), qualification, round-trip validation | | `@pgsql/quotes` | `packages/quotes` | SQL identifier/string quoting and keyword classification | @@ -37,6 +37,7 @@ Detailed workflow documentation lives in `.agents/skills/`: | Skill | Path | Covers | |-------|------|--------| +| **AST Traversal** | `.agents/skills/ast-traversal/SKILL.md` | Walking SQL and PL/pgSQL ASTs: choosing `walk` / `walkSql` / `walkSqlAst` / `walkPlpgsqlAst` / `traverse`, statement context, visitor composition, abort, mutation | | **Testing & Fixtures** | `.agents/skills/testing-fixtures/SKILL.md` | Fixture-based testing pipeline, adding new test fixtures, kitchen-sink workflow, PL/pgSQL fixtures, transform tests | | **Code Generation** | `.agents/skills/code-generation/SKILL.md` | Protobuf codegen (`build:proto`), type inference/generation (`pgsql-types`), keyword generation (`@pgsql/quotes`), version-specific deparsers | @@ -115,6 +116,12 @@ Version configuration lives in `config/versions.json` — maps PG versions (13-1 - TypeScript throughout, compiled to both CJS and ESM - `@pgsql/types` provides all AST node types — use them for type safety +- Traversal: reach for `walk` from `@pgsql/traverse` (any AST: SQL, PL/pgSQL, or a + parsed script) or `walkSql` from `plpgsql-parser` (SQL text). Use the + `walkSqlAst` / `walkPlpgsqlAst` primitives only when you deliberately want a + single node universe with no statement context, and `traverse` when you need to + mutate. Never hand-roll a `transformSync(..., { hydrate: true })` + + per-statement loop harness — that is what `walk` is for - `@pgsql/quotes` handles SQL identifier quoting — use `QuoteUtils` methods - Test files go in `__tests__/` within each package - Fixture SQL files go in `__fixtures__/kitchen-sink/` (see testing-fixtures skill) diff --git a/README.md b/README.md index 21603b84..99f6eb6a 100644 --- a/README.md +++ b/README.md @@ -149,19 +149,30 @@ const walker: Walker = (path: NodePath) => { walk(ast, walker); -// Using a visitor object (recommended for multiple node types) -const visitor: Visitor = { +// Using a visitor object (recommended for multiple node types). +// SQL and PL/pgSQL node tags may be mixed; the second argument to each +// callback describes the statement the node belongs to. +walk(ast, { SelectStmt: (path) => { console.log('SELECT statement:', path.node); }, - RangeVar: (path) => { + RangeVar: (path, ctx) => { console.log('Table:', path.node.relname); - console.log('Path to table:', path.path); - console.log('Parent node:', path.parent?.tag); - } -}; + console.log('Is a write target:', ctx.isWrite); + console.log('Inside function:', ctx.functionName); + }, + PLpgSQL_stmt_dynexecute: (_path, ctx) => ctx.abort('dynamic EXECUTE') +}); +``` + +Starting from SQL **text** instead of an AST? Use `walkSql` from +`plpgsql-parser`, which parses, hydrates PL/pgSQL bodies, and then walks: + +```typescript +import { loadModule, walkSql } from 'plpgsql-parser'; -walk(ast, visitor); +await loadModule(); +walkSql('SELECT * FROM users', { RangeVar: (path) => console.log(path.node.relname) }); ``` ## 📦 Packages @@ -175,7 +186,7 @@ walk(ast, visitor); | [**pg-proto-parser**](./packages/proto-parser) | PostgreSQL protobuf parser and code generator | • Generate TypeScript interfaces from protobuf
• Create enum mappings and utilities
• AST helper generation | | [**@pgsql/transform-ast**](./packages/transform-ast) | Multi-version PostgreSQL AST transformer | • Transform ASTs between PostgreSQL versions (13→17)
• Single source of truth deparser pipeline
• Backward compatibility for legacy SQL | | [**@pgsql/transform**](./packages/transform) | SQL transformation & classification | • Schema-name rewriting (incl. PL/pgSQL bodies)
• Per-statement AST facts (`classifyStatements`)
• Qualification & round-trip validation | -| [**@pgsql/traverse**](./packages/traverse) | PostgreSQL AST traversal utilities | • Visitor pattern for traversing PostgreSQL AST nodes
• NodePath context with parent/path information
• Runtime schema-based precise traversal | +| [**@pgsql/traverse**](./packages/traverse) | SQL + PL/pgSQL AST traversal | • One `walk()` for SQL ASTs, PL/pgSQL ASTs, and parsed scripts
• `NodePath` structure plus statement context (`isWrite`, `insideFunction`, ...)
• Visitor composition, skip-children, and whole-walk abort | ## 🛠️ Development diff --git a/packages/plpgsql-parser/README.md b/packages/plpgsql-parser/README.md index 67fe4535..92bf9177 100644 --- a/packages/plpgsql-parser/README.md +++ b/packages/plpgsql-parser/README.md @@ -103,99 +103,82 @@ Options: ## Traverse API -The package provides a visitor pattern for traversing PL/pgSQL ASTs, similar to `@pgsql/traverse` but designed for PL/pgSQL node types. +The walkers themselves live in [`@pgsql/traverse`](../traverse) and are +re-exported here, so one import covers parsing and traversal. This package owns +the one entry point that genuinely needs a parser: **SQL text in**. -### `walk(root, callback, options?)` +### `walkSql(sql, visitors, options?)` -Walks the tree of PL/pgSQL AST nodes using a visitor pattern. +Parses a SQL string, hydrates its PL/pgSQL function bodies, and walks both with +the given visitors — SQL statements and PL/pgSQL bodies in a single pass. ```typescript -import { parse, walk, loadModule } from 'plpgsql-parser'; -import type { PLpgSQLVisitor } from 'plpgsql-parser'; +import { loadModule, walkSql } from 'plpgsql-parser'; await loadModule(); -const parsed = parse(` - CREATE FUNCTION get_user(p_id int) - RETURNS text - LANGUAGE plpgsql - AS $$ - BEGIN - RETURN (SELECT name FROM users WHERE id = p_id); - END; - $$; -`); - -// Visit PL/pgSQL nodes -const visitor: PLpgSQLVisitor = { - PLpgSQL_stmt_block: (path) => { - console.log('Found block at path:', path.path); - }, - PLpgSQL_stmt_return: (path) => { - console.log('Found return statement'); +const result = walkSql(sql, { + // SQL nodes, at the top level and inside function bodies + RangeVar: (path, ctx) => { + if (ctx.isWrite && path.node.schemaname === 'audit') { + ctx.abort('the audit schema is read-only'); + } + if (ctx.insideFunction) { + console.log(`${path.node.relname} referenced by ${ctx.functionName}`); + } }, -}; + // PL/pgSQL-only nodes, in the same visitor + PLpgSQL_stmt_dynexecute: (_path, ctx) => ctx.abort('dynamic EXECUTE is not allowed') +}); -walk(parsed.functions[0].plpgsql.hydrated, visitor); +result.aborted; // true when a visitor called ctx.abort() +result.reason; // 'the audit schema is read-only' ``` +Pass an array of visitors to compose independent policies in one parse. Every +callback receives a `WalkContext` (`stmtTag`, `stmtIndex`, `isWrite`, `isRead`, +`insideFunction`, `functionName`, `abort`) — see the +[`@pgsql/traverse` README](../traverse) for the full traversal reference. + Options: -- `walkSqlExpressions` (default: `true`) - Whether to recurse into hydrated SQL expressions -- `sqlVisitor` - SQL visitor to use when walking hydrated SQL expressions (from `@pgsql/traverse`) +- `walkFunctionBodies` (default: `true`) - Hydrate and walk PL/pgSQL function bodies. `false` skips the PL/pgSQL parse entirely +- `walkSqlExpressions` (default: `true`) - Recurse into hydrated SQL expressions inside bodies +- `sqlVisitor` - Override the visitor used for those SQL expressions + +Unparseable input is reported as `{ aborted: true, reason }` rather than +throwing, so a validator can treat "rejected" and "could not be understood" +uniformly. -### `walkParsedScript(parsed, plpgsqlVisitor, sqlVisitor?)` +### `walk(ast, visitors, options?)` -Convenience function that walks both SQL statements and PL/pgSQL function bodies. +Re-exported from `@pgsql/traverse`. Same behavior as `walkSql`, but takes an AST +you already have — a `ParsedScript` from `parse()`, a `ParseResult`, a SQL node, +or a PL/pgSQL node: ```typescript -import { parse, walkParsedScript, loadModule } from 'plpgsql-parser'; +import { loadModule, parse, walk } from 'plpgsql-parser'; await loadModule(); const parsed = parse(` CREATE TABLE users (id int); - CREATE FUNCTION get_user(p_id int) RETURNS text LANGUAGE plpgsql AS $$ + CREATE FUNCTION get_user(id int) RETURNS text LANGUAGE plpgsql AS $$ BEGIN - RETURN (SELECT name FROM users WHERE id = p_id); + RETURN (SELECT name FROM users WHERE users.id = id); END; $$; `); -walkParsedScript( - parsed, - // PL/pgSQL visitor - { - PLpgSQL_stmt_return: (path) => { - console.log('PL/pgSQL return statement'); - }, - }, - // SQL visitor (optional) - visits both top-level SQL and embedded SQL in functions - { - CreateStmt: (path) => { - console.log('CREATE TABLE statement'); - }, - RangeVar: (path) => { - console.log('Table reference:', path.node.relname); - }, - } -); +walk(parsed, { + CreateStmt: () => console.log('CREATE TABLE statement'), + RangeVar: (path) => console.log('Table reference:', path.node.relname), + PLpgSQL_stmt_return: () => console.log('PL/pgSQL return statement') +}); ``` -### `PLpgSQLNodePath` - -The path object passed to visitor functions: - -```typescript -class PLpgSQLNodePath { - tag: TTag; // Node type (e.g., 'PLpgSQL_stmt_block') - node: any; // The actual node data - parent: PLpgSQLNodePath | null; // Parent path - keyPath: readonly (string | number)[]; // Full path array - - get path(): (string | number)[]; // Copy of keyPath - get key(): string | number; // Last element of path -} -``` +Also re-exported: `walkSqlAst` (SQL-only primitive), `walkPlpgsqlAst` +(PL/pgSQL-only primitive), `PlpgsqlNodePath`, and the `WalkContext` / +`UnifiedVisitor` / `WalkResult` types. ## Re-exports @@ -207,6 +190,7 @@ For power users, the package re-exports underlying primitives: - `deparsePlpgsqlBody` - PL/pgSQL deparser from `plpgsql-deparser` - `hydratePlpgsqlAst` - Hydration utility from `plpgsql-deparser` - `dehydratePlpgsqlAst` - Dehydration utility from `plpgsql-deparser` +- `walk`, `walkSqlAst`, `walkPlpgsqlAst` - Walkers from `@pgsql/traverse` ## License diff --git a/packages/plpgsql-parser/__tests__/traverse.test.ts b/packages/plpgsql-parser/__tests__/traverse.test.ts index bffa4649..272d3a1b 100644 --- a/packages/plpgsql-parser/__tests__/traverse.test.ts +++ b/packages/plpgsql-parser/__tests__/traverse.test.ts @@ -1,5 +1,5 @@ -import { parse, walk, walkParsedScript, PLpgSQLNodePath, loadModule } from '../src'; -import type { PLpgSQLVisitor } from '../src'; +import { loadModule,parse, PlpgsqlNodePath, walk } from '../src'; +import type { PlpgsqlVisitor } from '../src'; describe('plpgsql-parser traverse', () => { beforeAll(async () => { @@ -26,7 +26,7 @@ describe('plpgsql-parser traverse', () => { expect(parsed.functions.length).toBe(1); const visitedTags: string[] = []; - const visitor: PLpgSQLVisitor = { + const visitor: PlpgsqlVisitor = { PLpgSQL_function: (path) => { visitedTags.push(path.tag); }, @@ -57,7 +57,7 @@ describe('plpgsql-parser traverse', () => { const parsed = parse(simpleFunctionSql); const visitedTags: string[] = []; - walk(parsed.functions[0].plpgsql.hydrated, (path: PLpgSQLNodePath) => { + walk(parsed.functions[0].plpgsql.hydrated, (path: PlpgsqlNodePath) => { visitedTags.push(path.tag); }); @@ -69,7 +69,7 @@ describe('plpgsql-parser traverse', () => { const parsed = parse(simpleFunctionSql); let blockPath: (string | number)[] = []; - const visitor: PLpgSQLVisitor = { + const visitor: PlpgsqlVisitor = { PLpgSQL_stmt_block: (path) => { blockPath = path.path; }, @@ -84,7 +84,7 @@ describe('plpgsql-parser traverse', () => { const parsed = parse(simpleFunctionSql); const visitedTags: string[] = []; - const visitor: PLpgSQLVisitor = { + const visitor: PlpgsqlVisitor = { PLpgSQL_function: (path) => { visitedTags.push(path.tag); return false; // Skip children @@ -101,30 +101,25 @@ describe('plpgsql-parser traverse', () => { }); }); - describe('walkParsedScript', () => { - it('should walk both SQL and PL/pgSQL nodes', () => { + describe('walking a parsed script', () => { + it('should walk both SQL and PL/pgSQL nodes from one visitor', () => { const parsed = parse(simpleFunctionSql); - + const plpgsqlTags: string[] = []; const sqlTags: string[] = []; - - walkParsedScript( - parsed, - { - PLpgSQL_function: (path) => { - plpgsqlTags.push(path.tag); - }, - PLpgSQL_stmt_block: (path) => { - plpgsqlTags.push(path.tag); - }, + + walk(parsed, { + PLpgSQL_function: (path) => { + plpgsqlTags.push(path.tag); }, - { - CreateFunctionStmt: (path) => { - sqlTags.push(path.tag); - }, - } - ); - + PLpgSQL_stmt_block: (path) => { + plpgsqlTags.push(path.tag); + }, + CreateFunctionStmt: (path) => { + sqlTags.push(path.tag); + }, + }); + expect(plpgsqlTags).toContain('PLpgSQL_function'); expect(sqlTags).toContain('CreateFunctionStmt'); }); @@ -135,7 +130,7 @@ describe('plpgsql-parser traverse', () => { const parsed = parse(simpleFunctionSql); const sqlTags: string[] = []; - const visitor: PLpgSQLVisitor = { + const visitor: PlpgsqlVisitor = { PLpgSQL_expr: () => { // Just visit the expression node }, @@ -181,7 +176,7 @@ describe('plpgsql-parser traverse', () => { const parsed = parse(ifFunctionSql); const visitedTags: string[] = []; - walk(parsed.functions[0].plpgsql.hydrated, (path: PLpgSQLNodePath) => { + walk(parsed.functions[0].plpgsql.hydrated, (path: PlpgsqlNodePath) => { visitedTags.push(path.tag); }); @@ -209,7 +204,7 @@ describe('plpgsql-parser traverse', () => { const parsed = parse(loopFunctionSql); const visitedTags: string[] = []; - walk(parsed.functions[0].plpgsql.hydrated, (path: PLpgSQLNodePath) => { + walk(parsed.functions[0].plpgsql.hydrated, (path: PlpgsqlNodePath) => { visitedTags.push(path.tag); }); @@ -236,7 +231,7 @@ describe('plpgsql-parser traverse', () => { const parsed = parse(forFunctionSql); const visitedTags: string[] = []; - walk(parsed.functions[0].plpgsql.hydrated, (path: PLpgSQLNodePath) => { + walk(parsed.functions[0].plpgsql.hydrated, (path: PlpgsqlNodePath) => { visitedTags.push(path.tag); }); diff --git a/packages/plpgsql-parser/__tests__/walk-sql.test.ts b/packages/plpgsql-parser/__tests__/walk-sql.test.ts new file mode 100644 index 00000000..0c73ada9 --- /dev/null +++ b/packages/plpgsql-parser/__tests__/walk-sql.test.ts @@ -0,0 +1,287 @@ +import type { WalkContext } from '../src'; +import { loadModule, walkSql } from '../src'; + +beforeAll(async () => { + await loadModule(); +}); + +const CREATE_FUNCTION_SQL = ` + CREATE FUNCTION my_func() RETURNS void LANGUAGE plpgsql AS $$ + DECLARE + v_row metaschema_public.users%ROWTYPE; + BEGIN + SELECT * INTO v_row FROM metaschema_public.users WHERE id = 1; + INSERT INTO public.audit_log (msg) VALUES ('accessed'); + END; + $$; +`; + +describe('walkSql', () => { + describe('basic traversal', () => { + it('parses and walks a simple SELECT', () => { + const visited: string[] = []; + const result = walkSql('SELECT * FROM public.users', { + RangeVar: (path) => { + visited.push(`${path.node.schemaname ?? ''}.${path.node.relname}`); + }, + }); + expect(result.aborted).toBe(false); + expect(visited).toContain('public.users'); + }); + + it('visits function calls', () => { + const funcs: string[] = []; + walkSql("SELECT pg_catalog.set_config('a', 'b', false)", { + FuncCall: (path) => { + funcs.push( + (path.node.funcname ?? []) + .map((part: any) => part?.String?.sval ?? '') + .filter(Boolean) + .join('.'), + ); + }, + }); + expect(funcs).toContain('pg_catalog.set_config'); + }); + + it('handles empty SQL', () => { + expect(walkSql('', {}).aborted).toBe(false); + expect(walkSql(' ', {}).aborted).toBe(false); + }); + + it('reports unparseable SQL as an abort rather than throwing', () => { + const result = walkSql('NOT VALID SQL !!!', {}); + expect(result.aborted).toBe(true); + expect(result.reason).toBeDefined(); + }); + }); + + describe('statement context', () => { + it('fires the statement hook once per statement with its tag', () => { + const tags: string[] = []; + walkSql('SELECT 1; INSERT INTO t(a) VALUES(1)', { + statement: (path) => { + tags.push(path.tag); + }, + }); + expect(tags).toEqual(['SelectStmt', 'InsertStmt']); + }); + + it('marks writes and reads', () => { + const contexts: WalkContext[] = []; + const collect = { + statement: (_path: unknown, ctx: WalkContext): void => { + contexts.push(ctx); + }, + }; + + walkSql('INSERT INTO t(a) VALUES(1)', collect); + expect(contexts[0].isWrite).toBe(true); + expect(contexts[0].isRead).toBe(false); + + contexts.length = 0; + walkSql('SELECT * FROM t', collect); + expect(contexts[0].isWrite).toBe(false); + expect(contexts[0].isRead).toBe(true); + }); + + it('gives nested nodes the context of their own statement', () => { + const seen: Array<{ table: string; stmtTag: string | null; isWrite: boolean }> = []; + walkSql('SELECT * FROM a.reads; UPDATE b.writes SET x = 1', { + RangeVar: (path, ctx) => { + seen.push({ table: path.node.relname, stmtTag: ctx.stmtTag, isWrite: ctx.isWrite }); + }, + }); + expect(seen).toEqual([ + { table: 'reads', stmtTag: 'SelectStmt', isWrite: false }, + { table: 'writes', stmtTag: 'UpdateStmt', isWrite: true }, + ]); + }); + + it('refines the context to the nearest enclosing statement', () => { + const seen: Array<{ table: string; stmtTag: string | null; isWrite: boolean }> = []; + walkSql( + 'WITH w AS (INSERT INTO a.written VALUES (1) RETURNING *) SELECT * FROM b.read', + { + RangeVar: (path, ctx) => { + seen.push({ table: path.node.relname, stmtTag: ctx.stmtTag, isWrite: ctx.isWrite }); + }, + }, + ); + expect(seen).toEqual( + expect.arrayContaining([ + { table: 'written', stmtTag: 'InsertStmt', isWrite: true }, + { table: 'read', stmtTag: 'SelectStmt', isWrite: false }, + ]), + ); + }); + + it('marks a write inside a function body as a write', () => { + const seen: Array<{ table: string; isWrite: boolean; functionName: string | null }> = []; + walkSql( + `CREATE FUNCTION w() RETURNS void LANGUAGE plpgsql AS $$ + BEGIN + INSERT INTO infra.servers (name) VALUES ('x'); + END; + $$;`, + { + RangeVar: (path, ctx) => { + seen.push({ + table: path.node.relname, + isWrite: ctx.isWrite, + functionName: ctx.functionName, + }); + }, + }, + ); + expect(seen).toEqual([{ table: 'servers', isWrite: true, functionName: 'w' }]); + }); + + it('numbers statements', () => { + const indexes: number[] = []; + walkSql('SELECT 1; SELECT 2; SELECT 3', { + statement: (_path, ctx) => void indexes.push(ctx.stmtIndex), + }); + expect(indexes).toEqual([0, 1, 2]); + }); + }); + + describe('visitor composition', () => { + it('fires every visitor on each node in a single pass', () => { + const first: string[] = []; + const second: string[] = []; + walkSql('SELECT * FROM users', [ + { RangeVar: (path) => void first.push(path.node.relname) }, + { RangeVar: (path) => void second.push(path.node.relname) }, + ]); + expect(first).toEqual(['users']); + expect(second).toEqual(['users']); + }); + + it('accepts a bare walker function', () => { + const tags: string[] = []; + walkSql('SELECT 1', (path) => void tags.push(path.tag)); + expect(tags).toContain('SelectStmt'); + }); + }); + + describe('control flow', () => { + it('skips a node\u2019s children when a visitor returns false', () => { + const tags: string[] = []; + walkSql('SELECT * FROM users', { + SelectStmt: () => false, + RangeVar: (path) => void tags.push(path.tag), + }); + expect(tags).toEqual([]); + }); + + it('ends the whole walk on abort, skipping later statements', () => { + const tables: string[] = []; + const result = walkSql('SELECT * FROM a.one; SELECT * FROM b.two', { + RangeVar: (path, ctx) => { + tables.push(path.node.relname); + ctx.abort('seen enough'); + }, + }); + expect(tables).toEqual(['one']); + expect(result.aborted).toBe(true); + expect(result.reason).toBe('seen enough'); + expect(result.reasons).toEqual(['seen enough']); + }); + + it('aborts from inside a function body', () => { + const result = walkSql(CREATE_FUNCTION_SQL, { + PLpgSQL_stmt_execsql: (_path, ctx) => ctx.abort('no SQL in bodies'), + }); + expect(result.aborted).toBe(true); + expect(result.reason).toBe('no SQL in bodies'); + }); + }); + + describe('function bodies', () => { + it('visits SQL nodes inside hydrated PL/pgSQL bodies', () => { + const tables: string[] = []; + const result = walkSql(CREATE_FUNCTION_SQL, { + RangeVar: (path) => { + if (path.node.schemaname) tables.push(`${path.node.schemaname}.${path.node.relname}`); + }, + }); + expect(result.aborted).toBe(false); + expect(tables).toContain('metaschema_public.users'); + expect(tables).toContain('public.audit_log'); + }); + + it('reports insideFunction and the function name', () => { + const contexts: WalkContext[] = []; + walkSql(CREATE_FUNCTION_SQL, { + RangeVar: (_path, ctx) => void contexts.push(ctx), + }); + const insideFn = contexts.filter((ctx) => ctx.insideFunction); + expect(insideFn.length).toBeGreaterThan(0); + expect(insideFn[0].functionName).toBe('my_func'); + }); + + it('visits every statement of a multi-statement body', () => { + const schemas: string[] = []; + walkSql( + `CREATE FUNCTION multi_stmt() RETURNS void LANGUAGE plpgsql AS $$ + BEGIN + INSERT INTO schema_a.table1 (x) VALUES (1); + UPDATE schema_b.table2 SET x = 2; + DELETE FROM schema_c.table3 WHERE id = 3; + END; + $$;`, + { RangeVar: (path) => void (path.node.schemaname && schemas.push(path.node.schemaname)) }, + ); + expect(schemas).toEqual(expect.arrayContaining(['schema_a', 'schema_b', 'schema_c'])); + }); + + it('visits PL/pgSQL-only nodes such as dynamic EXECUTE', () => { + const dynamic = walkSql( + `CREATE FUNCTION dangerous() RETURNS void LANGUAGE plpgsql AS $$ + BEGIN + EXECUTE 'DROP TABLE users'; + END; + $$;`, + { PLpgSQL_stmt_dynexecute: (_path, ctx) => ctx.abort('dynamic EXECUTE') }, + ); + expect(dynamic.aborted).toBe(true); + expect(dynamic.reason).toBe('dynamic EXECUTE'); + }); + + it('leaves bodies unparsed when walkFunctionBodies is false', () => { + const result = walkSql( + `CREATE FUNCTION dangerous() RETURNS void LANGUAGE plpgsql AS $$ + BEGIN + EXECUTE 'DROP TABLE users'; + END; + $$;`, + { PLpgSQL_stmt_dynexecute: (_path, ctx) => ctx.abort('dynamic EXECUTE') }, + { walkFunctionBodies: false }, + ); + expect(result.aborted).toBe(false); + }); + + it('visits function calls made with PERFORM', () => { + const funcs: string[] = []; + walkSql( + `CREATE FUNCTION caller() RETURNS void LANGUAGE plpgsql AS $$ + BEGIN + PERFORM pg_catalog.set_config('role', 'admin', false); + END; + $$;`, + { + FuncCall: (path) => { + funcs.push( + (path.node.funcname ?? []) + .map((part: any) => part?.String?.sval ?? '') + .filter(Boolean) + .join('.'), + ); + }, + }, + ); + expect(funcs).toContain('pg_catalog.set_config'); + }); + }); +}); diff --git a/packages/plpgsql-parser/package.json b/packages/plpgsql-parser/package.json index 6fecf2d9..8810e7fa 100644 --- a/packages/plpgsql-parser/package.json +++ b/packages/plpgsql-parser/package.json @@ -1,6 +1,6 @@ { "name": "plpgsql-parser", - "version": "18.2.3", + "version": "18.3.0", "author": "Constructive ", "description": "Combined SQL + PL/pgSQL parser with hydrated ASTs and transform API", "main": "index.js", diff --git a/packages/plpgsql-parser/src/index.ts b/packages/plpgsql-parser/src/index.ts index 775ef231..00c34221 100644 --- a/packages/plpgsql-parser/src/index.ts +++ b/packages/plpgsql-parser/src/index.ts @@ -2,15 +2,27 @@ export * from './types'; export { parse, parseSync, loadModule } from './parse'; export { deparse, deparseSync } from './deparse'; export { transform, transformSync } from './transform'; -export { - walk, - walkParsedScript, - PLpgSQLNodePath, - type PLpgSQLWalker, - type PLpgSQLVisitor, - type PLpgSQLNodeTag, - type WalkOptions -} from './traverse'; +export { walkSql, type WalkSqlOptions } from './traverse'; + +// The walkers live in @pgsql/traverse; re-exported here so a single import +// covers parsing and traversal. +export { + PlpgsqlNodePath, + type PlpgsqlNodeTag, + type PlpgsqlVisitor, + type PlpgsqlWalker, + type PlpgsqlWalkOptions, + READ_STATEMENTS, + type UnifiedVisitor, + type UnifiedWalker, + walk, + type WalkContext, + type WalkOptions, + walkPlpgsqlAst, + type WalkResult, + walkSqlAst, + WRITE_STATEMENTS +} from '@pgsql/traverse'; export { getReturnInfo, getReturnInfoFromParsedFunction } from './return-info'; export { diff --git a/packages/plpgsql-parser/src/traverse.ts b/packages/plpgsql-parser/src/traverse.ts index a37ab818..8d6165af 100644 --- a/packages/plpgsql-parser/src/traverse.ts +++ b/packages/plpgsql-parser/src/traverse.ts @@ -1,673 +1,66 @@ /** - * PL/pgSQL AST Traversal - * - * Provides a visitor pattern for traversing PL/pgSQL ASTs, similar to @pgsql/traverse - * but designed for PL/pgSQL node types. Automatically recurses into hydrated SQL - * expressions using @pgsql/traverse. + * SQL text traversal. + * + * The walkers themselves live in `@pgsql/traverse`, which knows nothing about + * parsing. This module owns the one thing that genuinely needs a parser: going + * from a **SQL string** to a walk over both its statements and its hydrated + * PL/pgSQL function bodies. */ -import { walk as walkSql } from '@pgsql/traverse'; -import type { Walker as SqlWalker, Visitor as SqlVisitor, NodePath as SqlNodePath } from '@pgsql/traverse'; import type { - PLpgSQLParseResult, - PLpgSQLFunctionNode, - PLpgSQL_function, - PLpgSQLDatum, - PLpgSQLStmtNode, - PLpgSQLExprNode, - PLpgSQLTypeNode, - PLpgSQL_stmt_block, - PLpgSQL_stmt_if, - PLpgSQL_stmt_case, - PLpgSQL_stmt_loop, - PLpgSQL_stmt_while, - PLpgSQL_stmt_fori, - PLpgSQL_stmt_fors, - PLpgSQL_stmt_forc, - PLpgSQL_stmt_foreach_a, - PLpgSQL_stmt_return_query, - PLpgSQL_stmt_raise, - PLpgSQL_stmt_dynexecute, - PLpgSQL_stmt_dynfors, - PLpgSQL_stmt_open, - PLpgSQLElsifNode, - PLpgSQLCaseWhenNode, - PLpgSQLException, - PLpgSQLRaiseOption, -} from 'plpgsql-deparser'; -import type { - HydratedExprQuery, - HydratedExprSqlStmt, - HydratedExprSqlExpr, - HydratedExprAssign, -} from 'plpgsql-deparser'; - -// PL/pgSQL node tag types -export type PLpgSQLNodeTag = - | 'PLpgSQL_function' - | 'PLpgSQL_var' - | 'PLpgSQL_rec' - | 'PLpgSQL_row' - | 'PLpgSQL_recfield' - | 'PLpgSQL_type' - | 'PLpgSQL_expr' - | 'PLpgSQL_stmt_block' - | 'PLpgSQL_stmt_assign' - | 'PLpgSQL_stmt_if' - | 'PLpgSQL_stmt_case' - | 'PLpgSQL_stmt_loop' - | 'PLpgSQL_stmt_while' - | 'PLpgSQL_stmt_fori' - | 'PLpgSQL_stmt_fors' - | 'PLpgSQL_stmt_forc' - | 'PLpgSQL_stmt_foreach_a' - | 'PLpgSQL_stmt_exit' - | 'PLpgSQL_stmt_return' - | 'PLpgSQL_stmt_return_next' - | 'PLpgSQL_stmt_return_query' - | 'PLpgSQL_stmt_raise' - | 'PLpgSQL_stmt_assert' - | 'PLpgSQL_stmt_execsql' - | 'PLpgSQL_stmt_dynexecute' - | 'PLpgSQL_stmt_dynfors' - | 'PLpgSQL_stmt_getdiag' - | 'PLpgSQL_stmt_open' - | 'PLpgSQL_stmt_fetch' - | 'PLpgSQL_stmt_close' - | 'PLpgSQL_stmt_perform' - | 'PLpgSQL_stmt_call' - | 'PLpgSQL_stmt_commit' - | 'PLpgSQL_stmt_rollback' - | 'PLpgSQL_stmt_set' - | 'PLpgSQL_if_elsif' - | 'PLpgSQL_case_when' - | 'PLpgSQL_exception' - | 'PLpgSQL_condition' - | 'PLpgSQL_raise_option' - | 'PLpgSQL_diag_item'; - -export class PLpgSQLNodePath { - constructor( - public tag: TTag, - public node: any, - public parent: PLpgSQLNodePath | null = null, - public keyPath: readonly (string | number)[] = [] - ) {} - - get path(): (string | number)[] { - return [...this.keyPath]; - } - - get key(): string | number { - return this.keyPath[this.keyPath.length - 1] ?? ''; - } -} - -export type PLpgSQLWalker = ( - path: TNodePath, -) => boolean | void; + UnifiedVisitor, + UnifiedWalker, + WalkOptions, + WalkResult, +} from '@pgsql/traverse'; +import { walk } from '@pgsql/traverse'; -export type PLpgSQLVisitor = { - [key: string]: PLpgSQLWalker; -}; +import { parseSync } from './parse'; -export interface WalkOptions { +export interface WalkSqlOptions extends WalkOptions { /** - * Whether to recurse into hydrated SQL expressions using @pgsql/traverse. - * Default: true + * Hydrate and walk PL/pgSQL function bodies. Turning this off skips the + * PL/pgSQL parse entirely, not just the traversal. Default: true */ - walkSqlExpressions?: boolean; - - /** - * SQL visitor to use when walking hydrated SQL expressions. - * Only used if walkSqlExpressions is true. - */ - sqlVisitor?: SqlVisitor | SqlWalker; + walkFunctionBodies?: boolean; } /** - * Walks the tree of PL/pgSQL AST nodes using a visitor pattern. - * - * If a callback returns `false`, the walk will continue to the next sibling - * node, rather than recurse into the children of the current node. - * - * @param root - The PL/pgSQL AST node to traverse - * @param callback - A walker function or visitor object - * @param options - Walk options - * @param parent - Parent NodePath (for internal use) - * @param keyPath - Current key path (for internal use) + * Parse a SQL string, hydrate its PL/pgSQL function bodies, and walk the whole + * thing with the given visitors. + * + * ```ts + * const result = walkSql(sql, { + * RangeVar: (path, ctx) => { + * if (ctx.isWrite && path.node.schemaname === 'audit') { + * ctx.abort('audit schema is read-only'); + * } + * }, + * PLpgSQL_stmt_dynexecute: (_path, ctx) => ctx.abort('dynamic EXECUTE is not allowed'), + * }); + * ``` + * + * Unparseable input is reported as an abort rather than a thrown error, so a + * validator can treat "rejected" and "could not be understood" uniformly. */ -export function walk( - root: any, - callback: PLpgSQLWalker | PLpgSQLVisitor, - options: WalkOptions = {}, - parent: PLpgSQLNodePath | null = null, - keyPath: readonly (string | number)[] = [], -): void { - const { walkSqlExpressions = true, sqlVisitor } = options; - - const actualCallback: PLpgSQLWalker = typeof callback === 'function' - ? callback - : (path: PLpgSQLNodePath) => { - const visitor = callback as PLpgSQLVisitor; - const visitFn = visitor[path.tag]; - return visitFn ? visitFn(path) : undefined; - }; +export function walkSql( + sql: string, + visitors: UnifiedVisitor | UnifiedWalker | Array, + options: WalkSqlOptions = {}, +): WalkResult { + const walkFunctionBodies = options.walkFunctionBodies ?? true; - if (Array.isArray(root)) { - root.forEach((node, index) => { - walk(node, actualCallback, options, parent, [...keyPath, index]); - }); - } else if (typeof root === 'object' && root !== null) { - const keys = Object.keys(root); - - // Check if this is a PL/pgSQL node (single key starting with PLpgSQL_) - if (keys.length === 1 && keys[0].startsWith('PLpgSQL_')) { - const tag = keys[0]; - const nodeData = root[tag]; - const path = new PLpgSQLNodePath(tag, nodeData, parent, keyPath); - - if (actualCallback(path) === false) { - return; - } - - // Recurse into child nodes based on node type - walkNodeChildren(tag, nodeData, actualCallback, options, path); - } else { - // Not a PL/pgSQL node wrapper, check for nested structures - for (const key of keys) { - const value = root[key]; - if (typeof value === 'object' && value !== null) { - walk(value, actualCallback, options, parent, [...keyPath, key]); - } - } - } - } - - // Helper function to walk into hydrated SQL expressions - function walkHydratedExpr(expr: any, exprPath: PLpgSQLNodePath) { - if (!walkSqlExpressions || !expr) return; - - // Check if this is a hydrated expression - if (expr.query && typeof expr.query === 'object' && 'kind' in expr.query) { - const hydratedQuery = expr.query as HydratedExprQuery; - - if (hydratedQuery.kind === 'sql-stmt') { - const sqlStmt = hydratedQuery as HydratedExprSqlStmt; - if (sqlStmt.parseResult && sqlVisitor) { - walkSql(sqlStmt.parseResult, sqlVisitor); - } - } else if (hydratedQuery.kind === 'sql-expr') { - const sqlExpr = hydratedQuery as HydratedExprSqlExpr; - if (sqlExpr.expr && sqlVisitor) { - walkSql(sqlExpr.expr, sqlVisitor); - } - } else if (hydratedQuery.kind === 'assign') { - const assignExpr = hydratedQuery as HydratedExprAssign; - if (assignExpr.targetExpr && sqlVisitor) { - walkSql(assignExpr.targetExpr, sqlVisitor); - } - if (assignExpr.valueExpr && sqlVisitor) { - walkSql(assignExpr.valueExpr, sqlVisitor); - } - } - } + if (!sql || sql.trim().length === 0) { + return { aborted: false, reason: undefined, reasons: [] }; } - - // Helper function to walk children based on node type - function walkNodeChildren( - tag: string, - nodeData: any, - cb: PLpgSQLWalker, - opts: WalkOptions, - parentPath: PLpgSQLNodePath - ) { - switch (tag) { - case 'PLpgSQL_function': { - const fn = nodeData as PLpgSQL_function; - if (fn.datums) { - fn.datums.forEach((datum, i) => { - walk(datum, cb, opts, parentPath, [...parentPath.keyPath, 'datums', i]); - }); - } - if (fn.action) { - walk(fn.action, cb, opts, parentPath, [...parentPath.keyPath, 'action']); - } - break; - } - - case 'PLpgSQL_var': { - if (nodeData.datatype) { - walk(nodeData.datatype, cb, opts, parentPath, [...parentPath.keyPath, 'datatype']); - } - if (nodeData.default_val) { - walk(nodeData.default_val, cb, opts, parentPath, [...parentPath.keyPath, 'default_val']); - } - if (nodeData.cursor_explicit_expr) { - walk(nodeData.cursor_explicit_expr, cb, opts, parentPath, [...parentPath.keyPath, 'cursor_explicit_expr']); - } - break; - } - - case 'PLpgSQL_expr': { - // This is where we recurse into SQL expressions - walkHydratedExpr(nodeData, parentPath); - break; - } - - case 'PLpgSQL_stmt_block': { - const block = nodeData as PLpgSQL_stmt_block; - if (block.body) { - block.body.forEach((stmt, i) => { - walk(stmt, cb, opts, parentPath, [...parentPath.keyPath, 'body', i]); - }); - } - if (block.exceptions?.exc_list) { - block.exceptions.exc_list.forEach((exc, i) => { - walk(exc, cb, opts, parentPath, [...parentPath.keyPath, 'exceptions', 'exc_list', i]); - }); - } - break; - } - - case 'PLpgSQL_stmt_assign': { - if (nodeData.expr) { - walk(nodeData.expr, cb, opts, parentPath, [...parentPath.keyPath, 'expr']); - } - break; - } - - case 'PLpgSQL_stmt_if': { - const ifStmt = nodeData as PLpgSQL_stmt_if; - if (ifStmt.cond) { - walk(ifStmt.cond, cb, opts, parentPath, [...parentPath.keyPath, 'cond']); - } - if (ifStmt.then_body) { - ifStmt.then_body.forEach((stmt, i) => { - walk(stmt, cb, opts, parentPath, [...parentPath.keyPath, 'then_body', i]); - }); - } - if (ifStmt.elsif_list) { - ifStmt.elsif_list.forEach((elsif, i) => { - walk(elsif, cb, opts, parentPath, [...parentPath.keyPath, 'elsif_list', i]); - }); - } - if (ifStmt.else_body) { - ifStmt.else_body.forEach((stmt, i) => { - walk(stmt, cb, opts, parentPath, [...parentPath.keyPath, 'else_body', i]); - }); - } - break; - } - - case 'PLpgSQL_if_elsif': { - if (nodeData.cond) { - walk(nodeData.cond, cb, opts, parentPath, [...parentPath.keyPath, 'cond']); - } - if (nodeData.stmts) { - nodeData.stmts.forEach((stmt: any, i: number) => { - walk(stmt, cb, opts, parentPath, [...parentPath.keyPath, 'stmts', i]); - }); - } - break; - } - - case 'PLpgSQL_stmt_case': { - const caseStmt = nodeData as PLpgSQL_stmt_case; - if (caseStmt.t_expr) { - walk(caseStmt.t_expr, cb, opts, parentPath, [...parentPath.keyPath, 't_expr']); - } - if (caseStmt.case_when_list) { - caseStmt.case_when_list.forEach((when, i) => { - walk(when, cb, opts, parentPath, [...parentPath.keyPath, 'case_when_list', i]); - }); - } - if (caseStmt.else_stmts) { - caseStmt.else_stmts.forEach((stmt, i) => { - walk(stmt, cb, opts, parentPath, [...parentPath.keyPath, 'else_stmts', i]); - }); - } - break; - } - - case 'PLpgSQL_case_when': { - if (nodeData.expr) { - walk(nodeData.expr, cb, opts, parentPath, [...parentPath.keyPath, 'expr']); - } - if (nodeData.stmts) { - nodeData.stmts.forEach((stmt: any, i: number) => { - walk(stmt, cb, opts, parentPath, [...parentPath.keyPath, 'stmts', i]); - }); - } - break; - } - - case 'PLpgSQL_stmt_loop': { - const loop = nodeData as PLpgSQL_stmt_loop; - if (loop.body) { - loop.body.forEach((stmt, i) => { - walk(stmt, cb, opts, parentPath, [...parentPath.keyPath, 'body', i]); - }); - } - break; - } - - case 'PLpgSQL_stmt_while': { - const whileStmt = nodeData as PLpgSQL_stmt_while; - if (whileStmt.cond) { - walk(whileStmt.cond, cb, opts, parentPath, [...parentPath.keyPath, 'cond']); - } - if (whileStmt.body) { - whileStmt.body.forEach((stmt, i) => { - walk(stmt, cb, opts, parentPath, [...parentPath.keyPath, 'body', i]); - }); - } - break; - } - - case 'PLpgSQL_stmt_fori': { - const fori = nodeData as PLpgSQL_stmt_fori; - if (fori.var) { - walk(fori.var, cb, opts, parentPath, [...parentPath.keyPath, 'var']); - } - if (fori.lower) { - walk(fori.lower, cb, opts, parentPath, [...parentPath.keyPath, 'lower']); - } - if (fori.upper) { - walk(fori.upper, cb, opts, parentPath, [...parentPath.keyPath, 'upper']); - } - if (fori.step) { - walk(fori.step, cb, opts, parentPath, [...parentPath.keyPath, 'step']); - } - if (fori.body) { - fori.body.forEach((stmt, i) => { - walk(stmt, cb, opts, parentPath, [...parentPath.keyPath, 'body', i]); - }); - } - break; - } - - case 'PLpgSQL_stmt_fors': { - const fors = nodeData as PLpgSQL_stmt_fors; - if (fors.var) { - walk(fors.var, cb, opts, parentPath, [...parentPath.keyPath, 'var']); - } - if (fors.query) { - walk(fors.query, cb, opts, parentPath, [...parentPath.keyPath, 'query']); - } - if (fors.body) { - fors.body.forEach((stmt, i) => { - walk(stmt, cb, opts, parentPath, [...parentPath.keyPath, 'body', i]); - }); - } - break; - } - - case 'PLpgSQL_stmt_forc': { - const forc = nodeData as PLpgSQL_stmt_forc; - if (forc.var) { - walk(forc.var, cb, opts, parentPath, [...parentPath.keyPath, 'var']); - } - if (forc.argquery) { - walk(forc.argquery, cb, opts, parentPath, [...parentPath.keyPath, 'argquery']); - } - if (forc.body) { - forc.body.forEach((stmt, i) => { - walk(stmt, cb, opts, parentPath, [...parentPath.keyPath, 'body', i]); - }); - } - break; - } - - case 'PLpgSQL_stmt_foreach_a': { - const foreach = nodeData as PLpgSQL_stmt_foreach_a; - if (foreach.expr) { - walk(foreach.expr, cb, opts, parentPath, [...parentPath.keyPath, 'expr']); - } - if (foreach.body) { - foreach.body.forEach((stmt, i) => { - walk(stmt, cb, opts, parentPath, [...parentPath.keyPath, 'body', i]); - }); - } - break; - } - - case 'PLpgSQL_stmt_exit': { - if (nodeData.cond) { - walk(nodeData.cond, cb, opts, parentPath, [...parentPath.keyPath, 'cond']); - } - break; - } - - case 'PLpgSQL_stmt_return': { - if (nodeData.expr) { - walk(nodeData.expr, cb, opts, parentPath, [...parentPath.keyPath, 'expr']); - } - break; - } - - case 'PLpgSQL_stmt_return_next': { - if (nodeData.expr) { - walk(nodeData.expr, cb, opts, parentPath, [...parentPath.keyPath, 'expr']); - } - break; - } - - case 'PLpgSQL_stmt_return_query': { - const retQuery = nodeData as PLpgSQL_stmt_return_query; - if (retQuery.query) { - walk(retQuery.query, cb, opts, parentPath, [...parentPath.keyPath, 'query']); - } - if (retQuery.dynquery) { - walk(retQuery.dynquery, cb, opts, parentPath, [...parentPath.keyPath, 'dynquery']); - } - if (retQuery.params) { - retQuery.params.forEach((param, i) => { - walk(param, cb, opts, parentPath, [...parentPath.keyPath, 'params', i]); - }); - } - break; - } - - case 'PLpgSQL_stmt_raise': { - const raise = nodeData as PLpgSQL_stmt_raise; - if (raise.params) { - raise.params.forEach((param, i) => { - walk(param, cb, opts, parentPath, [...parentPath.keyPath, 'params', i]); - }); - } - if (raise.options) { - raise.options.forEach((opt, i) => { - walk(opt, cb, opts, parentPath, [...parentPath.keyPath, 'options', i]); - }); - } - break; - } - - case 'PLpgSQL_raise_option': { - if (nodeData.expr) { - walk(nodeData.expr, cb, opts, parentPath, [...parentPath.keyPath, 'expr']); - } - break; - } - - case 'PLpgSQL_stmt_assert': { - if (nodeData.cond) { - walk(nodeData.cond, cb, opts, parentPath, [...parentPath.keyPath, 'cond']); - } - if (nodeData.message) { - walk(nodeData.message, cb, opts, parentPath, [...parentPath.keyPath, 'message']); - } - break; - } - - case 'PLpgSQL_stmt_execsql': { - if (nodeData.sqlstmt) { - walk(nodeData.sqlstmt, cb, opts, parentPath, [...parentPath.keyPath, 'sqlstmt']); - } - if (nodeData.target) { - walk(nodeData.target, cb, opts, parentPath, [...parentPath.keyPath, 'target']); - } - break; - } - - case 'PLpgSQL_stmt_dynexecute': { - const dynexec = nodeData as PLpgSQL_stmt_dynexecute; - if (dynexec.query) { - walk(dynexec.query, cb, opts, parentPath, [...parentPath.keyPath, 'query']); - } - if (dynexec.target) { - walk(dynexec.target, cb, opts, parentPath, [...parentPath.keyPath, 'target']); - } - if (dynexec.params) { - dynexec.params.forEach((param, i) => { - walk(param, cb, opts, parentPath, [...parentPath.keyPath, 'params', i]); - }); - } - break; - } - - case 'PLpgSQL_stmt_dynfors': { - const dynfors = nodeData as PLpgSQL_stmt_dynfors; - if (dynfors.var) { - walk(dynfors.var, cb, opts, parentPath, [...parentPath.keyPath, 'var']); - } - if (dynfors.query) { - walk(dynfors.query, cb, opts, parentPath, [...parentPath.keyPath, 'query']); - } - if (dynfors.params) { - dynfors.params.forEach((param, i) => { - walk(param, cb, opts, parentPath, [...parentPath.keyPath, 'params', i]); - }); - } - if (dynfors.body) { - dynfors.body.forEach((stmt, i) => { - walk(stmt, cb, opts, parentPath, [...parentPath.keyPath, 'body', i]); - }); - } - break; - } - - case 'PLpgSQL_stmt_open': { - const open = nodeData as PLpgSQL_stmt_open; - if (open.argquery) { - walk(open.argquery, cb, opts, parentPath, [...parentPath.keyPath, 'argquery']); - } - if (open.query) { - walk(open.query, cb, opts, parentPath, [...parentPath.keyPath, 'query']); - } - if (open.dynquery) { - walk(open.dynquery, cb, opts, parentPath, [...parentPath.keyPath, 'dynquery']); - } - if (open.params) { - open.params.forEach((param, i) => { - walk(param, cb, opts, parentPath, [...parentPath.keyPath, 'params', i]); - }); - } - break; - } - - case 'PLpgSQL_stmt_fetch': { - if (nodeData.target) { - walk(nodeData.target, cb, opts, parentPath, [...parentPath.keyPath, 'target']); - } - if (nodeData.expr) { - walk(nodeData.expr, cb, opts, parentPath, [...parentPath.keyPath, 'expr']); - } - break; - } - - case 'PLpgSQL_stmt_perform': { - if (nodeData.expr) { - walk(nodeData.expr, cb, opts, parentPath, [...parentPath.keyPath, 'expr']); - } - break; - } - - case 'PLpgSQL_stmt_call': { - if (nodeData.expr) { - walk(nodeData.expr, cb, opts, parentPath, [...parentPath.keyPath, 'expr']); - } - if (nodeData.target) { - walk(nodeData.target, cb, opts, parentPath, [...parentPath.keyPath, 'target']); - } - break; - } - - case 'PLpgSQL_stmt_set': { - if (nodeData.expr) { - walk(nodeData.expr, cb, opts, parentPath, [...parentPath.keyPath, 'expr']); - } - break; - } - - case 'PLpgSQL_exception': { - if (nodeData.conditions) { - nodeData.conditions.forEach((cond: any, i: number) => { - walk(cond, cb, opts, parentPath, [...parentPath.keyPath, 'conditions', i]); - }); - } - if (nodeData.action) { - nodeData.action.forEach((stmt: any, i: number) => { - walk(stmt, cb, opts, parentPath, [...parentPath.keyPath, 'action', i]); - }); - } - break; - } - - // Nodes with no children to traverse - case 'PLpgSQL_rec': - case 'PLpgSQL_row': - case 'PLpgSQL_recfield': - case 'PLpgSQL_type': - case 'PLpgSQL_stmt_getdiag': - case 'PLpgSQL_stmt_close': - case 'PLpgSQL_stmt_commit': - case 'PLpgSQL_stmt_rollback': - case 'PLpgSQL_condition': - case 'PLpgSQL_diag_item': - // No children to traverse - break; - - default: - // Unknown node type - try to traverse any object/array children - for (const key in nodeData) { - const value = nodeData[key]; - if (Array.isArray(value)) { - value.forEach((item, index) => { - if (typeof item === 'object' && item !== null) { - walk(item, cb, opts, parentPath, [...parentPath.keyPath, key, index]); - } - }); - } else if (typeof value === 'object' && value !== null) { - walk(value, cb, opts, parentPath, [...parentPath.keyPath, key]); - } - } - } - } -} -/** - * Convenience function to walk a parsed script from plpgsql-parser. - * Walks both the SQL statements and PL/pgSQL function bodies. - */ -export function walkParsedScript( - parsed: { sql: any; functions: Array<{ plpgsql: { hydrated: any } }> }, - plpgsqlVisitor: PLpgSQLVisitor | PLpgSQLWalker, - sqlVisitor?: SqlVisitor | SqlWalker, -): void { - // Walk SQL statements - if (sqlVisitor && parsed.sql) { - walkSql(parsed.sql, sqlVisitor); - } - - // Walk PL/pgSQL function bodies - for (const fn of parsed.functions) { - if (fn.plpgsql?.hydrated) { - walk(fn.plpgsql.hydrated, plpgsqlVisitor, { - walkSqlExpressions: true, - sqlVisitor - }); - } + let parsed; + try { + parsed = parseSync(sql, { hydrate: walkFunctionBodies }); + } catch (err) { + const reason = err instanceof Error ? err.message : 'Unparseable SQL'; + return { aborted: true, reason, reasons: [reason] }; } + + return walk(parsed, visitors, { ...options, walkFunctionBodies }); } diff --git a/packages/transform/README.md b/packages/transform/README.md index 50deb332..efef5b55 100644 --- a/packages/transform/README.md +++ b/packages/transform/README.md @@ -58,6 +58,29 @@ Qualify unqualified object references against an inventory of known objects, wit `normalizeTree` / `cleanTree` / `validateRoundTrip` — dependency-free AST normalization and mutation-aware parse→deparse→re-parse validation. +## Relationship to `@pgsql/traverse` + +This package owns **policy**, not traversal. Every entry point above parses once, +then drives `walk` from `@pgsql/traverse` over the parsed script — statements and +hydrated PL/pgSQL bodies alike — with a visitor built here: + +| Concern | Lives here | +|---|---| +| Which schema/role/extension a name maps to | `SchemaRouter`, `RoleRouter`, extension routes | +| What counts as a reference, a creation, a security-relevant statement | `classifyStatements` | +| When an unqualified name should be qualified | `qualifyUnqualified` | +| Whether the rewritten SQL still parses to the same tree | `validateRoundTrip` | +| How to reach every node of a SQL or PL/pgSQL AST | **`@pgsql/traverse`** | + +So mapping logic never moves into the walker, and traversal logic never moves in +here. Some passes keep per-statement state (CTE names in `qualifyUnqualified`, +per-statement facts in `classifyStatements`) and drive the `walkSqlAst` / +`walkPlpgsqlAst` primitives directly with a visitor per statement. + +Consequence for contributors: `__fixtures__/output/` is the regression gate for +traversal changes made *anywhere* in the ecosystem. A walker change that alters +these golden files is a behavior change, not a refactor. + ## Scripts - `npm run fixtures` — regenerate `__fixtures__/output/` golden files from `__fixtures__/input/` diff --git a/packages/transform/package.json b/packages/transform/package.json index 5027ac69..c04127cc 100644 --- a/packages/transform/package.json +++ b/packages/transform/package.json @@ -1,6 +1,6 @@ { "name": "@pgsql/transform", - "version": "18.11.1", + "version": "18.12.0", "author": "Constructive ", "description": "AST-based SQL transformation, qualification, classification and closure analysis for PostgreSQL", "main": "index.js", diff --git a/packages/transform/src/extension-transform.ts b/packages/transform/src/extension-transform.ts index a5b7c4db..3cb7493e 100644 --- a/packages/transform/src/extension-transform.ts +++ b/packages/transform/src/extension-transform.ts @@ -15,8 +15,8 @@ * (and symbols that have graduated into core) are left untouched. */ -import { walk as walkSql } from '@pgsql/traverse'; -import { Deparser, parseSql, transformSync, walk as walkPlpgsql } from 'plpgsql-parser'; +import { walk, walkSqlAst } from '@pgsql/traverse'; +import { Deparser, parseSql, transformSync } from 'plpgsql-parser'; import type { ExtensionRouteSpec, ExtensionRouterOptions, ExtensionSymbolNamespace } from './extension-router'; import { ExtensionRouter } from './extension-router'; @@ -191,7 +191,7 @@ function rewriteSqlBodyString( const pieces: string[] = []; for (const stmt of stmts) { if (!stmt?.stmt) continue; - walkSql(stmt.stmt, visitor); + walkSqlAst(stmt.stmt, visitor); pieces.push(Deparser.deparse(stmt.stmt)); } return pieces.join(';\n'); @@ -211,21 +211,11 @@ export function transformExtensions( const result = createExtensionResult(); const out = transformSync(sql, (ctx: any) => { - const stmts: any[] = ctx.sql?.stmts ?? []; - const visitor = createExtensionVisitor(resolved, result); - for (const stmt of stmts) { - if (stmt?.stmt) walkSql(stmt.stmt, visitor); - } - if (resolved.hasSymbolRoutes()) { - for (const fn of ctx.functions ?? []) { - if (fn.plpgsql?.hydrated) { - walkPlpgsql(fn.plpgsql.hydrated, {}, { - walkSqlExpressions: true, - sqlVisitor: visitor - }); - } - } - } + walk(ctx, createExtensionVisitor(resolved, result), { + // Bodies only matter when a symbol (function/operator/type) is routed; + // a bare extension rename never appears inside one. + walkFunctionBodies: resolved.hasSymbolRoutes() + }); }, { hydrate: true, pretty: true }); return { sql: out, result }; diff --git a/packages/transform/src/facts.ts b/packages/transform/src/facts.ts index 50ee1adf..e3242160 100644 --- a/packages/transform/src/facts.ts +++ b/packages/transform/src/facts.ts @@ -1,5 +1,5 @@ -import { walk as walkSql } from '@pgsql/traverse'; -import { parseSql, transformSync, walk as walkPlpgsql } from 'plpgsql-parser'; +import { walkSqlAst } from '@pgsql/traverse'; +import { parseSql, transformSync, walkPlpgsqlAst } from 'plpgsql-parser'; /** * A (possibly schema-qualified) object name extracted from a statement. @@ -473,7 +473,7 @@ function collectSqlBodyReferences(node: any, facts: StatementFacts): void { const stmts: any[] = parseSql(body)?.stmts ?? []; const visitor = createFactsVisitor(facts, facts.bodyReferences); for (const stmt of stmts) { - if (stmt?.stmt) walkSql(stmt.stmt, visitor); + if (stmt?.stmt) walkSqlAst(stmt.stmt, visitor); } } catch { // A non-parseable body (C symbol name, etc.) contributes no references. @@ -502,7 +502,7 @@ export function classifyStatements(sql: string): StatementFacts[] { if (stmtNode) { facts.stmt = stmtNode; - walkSql(stmtNode, createFactsVisitor(facts)); + walkSqlAst(stmtNode, createFactsVisitor(facts)); } if (nodeTag === 'CreateFunctionStmt') { collectSqlBodyReferences(node, facts); @@ -518,7 +518,7 @@ export function classifyStatements(sql: string): StatementFacts[] { const outer = new Set( facts.references.map(r => `${r.schema ?? '?'}.${r.name}`) ); - walkPlpgsql(fn.plpgsql.hydrated, { + walkPlpgsqlAst(fn.plpgsql.hydrated, { PLpgSQL_stmt_dynexecute: () => { facts.dynamicSql = true; }, PLpgSQL_stmt_dynfors: () => { facts.dynamicSql = true; } }, { diff --git a/packages/transform/src/qualify.ts b/packages/transform/src/qualify.ts index 1f7c5e50..eb27916e 100644 --- a/packages/transform/src/qualify.ts +++ b/packages/transform/src/qualify.ts @@ -15,8 +15,8 @@ * with a routed relation. */ -import { walk as walkSql } from '@pgsql/traverse'; -import { Deparser, parseSql, transformSync, walk as walkPlpgsql } from 'plpgsql-parser'; +import { walkSqlAst } from '@pgsql/traverse'; +import { Deparser, parseSql, transformSync, walkPlpgsqlAst } from 'plpgsql-parser'; import { classifyStatements } from './facts'; @@ -189,7 +189,7 @@ function resolveRoutes(sql: string, options: QualifyUnqualifiedOptions): Qualify function collectCteNames(stmt: any): Set { const names = new Set(); - walkSql(stmt, { + walkSqlAst(stmt, { CommonTableExpr: (path: any) => { if (path.node?.ctename) names.add(path.node.ctename); } @@ -217,7 +217,7 @@ function qualifySqlBodyString( { routes, cteNames: collectCteNames(stmt.stmt) }, result ); - walkSql(stmt.stmt, visitor); + walkSqlAst(stmt.stmt, visitor); pieces.push(Deparser.deparse(stmt.stmt)); } return pieces.join(';\n'); @@ -365,14 +365,14 @@ export function qualifyUnqualified( { routes, cteNames: collectCteNames(stmt.stmt) }, result ); - walkSql(stmt.stmt, visitor); + walkSqlAst(stmt.stmt, visitor); } } const bodyVisitor = createQualifyVisitor({ routes }, result); for (const fn of ctx.functions) { if (fn.plpgsql?.hydrated) { - walkPlpgsql(fn.plpgsql.hydrated, {}, { + walkPlpgsqlAst(fn.plpgsql.hydrated, {}, { walkSqlExpressions: true, sqlVisitor: bodyVisitor }); diff --git a/packages/transform/src/role-transform.ts b/packages/transform/src/role-transform.ts index ae67bf3f..2e888372 100644 --- a/packages/transform/src/role-transform.ts +++ b/packages/transform/src/role-transform.ts @@ -21,7 +21,7 @@ * rewritten. */ -import { walk as walkSql } from '@pgsql/traverse'; +import { walkSqlAst } from '@pgsql/traverse'; import { transformSync } from 'plpgsql-parser'; import type { RoleRouteSpec } from './role-router'; @@ -130,7 +130,7 @@ export function transformRoles( const stmts: any[] = ctx.sql?.stmts ?? []; const visitor = createRoleVisitor(resolved, result); for (const stmt of stmts) { - if (stmt?.stmt) walkSql(stmt.stmt, visitor); + if (stmt?.stmt) walkSqlAst(stmt.stmt, visitor); } }, { hydrate: true, pretty: true }); diff --git a/packages/transform/src/transform.ts b/packages/transform/src/transform.ts index b13519c0..937e5c27 100644 --- a/packages/transform/src/transform.ts +++ b/packages/transform/src/transform.ts @@ -25,8 +25,8 @@ */ import { QuoteUtils } from '@pgsql/quotes'; -import { walk as walkSql } from '@pgsql/traverse'; -import { Deparser,parseSql, transformSync, walk as walkPlpgsql } from 'plpgsql-parser'; +import { walk, walkSqlAst } from '@pgsql/traverse'; +import { Deparser,parseSql, transformSync, walkPlpgsqlAst } from 'plpgsql-parser'; import type { QualifyUnqualifiedOptions } from './qualify'; import { qualifyUnqualified } from './qualify'; @@ -880,7 +880,7 @@ export function transformPlpgsqlTypeAst( } }; - walkSql(parseResult.stmts[0].stmt, sqlVisitor); + walkSqlAst(parseResult.stmts[0].stmt, sqlVisitor); const deparsed = Deparser.deparse(parseResult.stmts[0].stmt); @@ -958,7 +958,7 @@ export function walkPlpgsqlForSchemas( if (Array.isArray(expr.query.parseResult?.stmts)) { for (const stmt of expr.query.parseResult.stmts) { if (stmt?.stmt) { - walkSql(stmt.stmt, sqlVisitor); + walkSqlAst(stmt.stmt, sqlVisitor); } } } @@ -981,7 +981,7 @@ export function walkPlpgsqlForSchemas( if (plType.typname) { if (typeof plType.typname === 'object' && plType.typname.kind === 'type-name') { // Transform schema names directly in the typeNameNode.names array. - // We cannot use walkSql here because the raw typeNameNode is not + // We cannot use walkSqlAst here because the raw typeNameNode is not // wrapped in the expected AST envelope that the traverse walker needs. if (plType.typname.typeNameNode?.names) { transformNameList(plType.typname.typeNameNode.names, schemaMapping, result, 'type'); @@ -1035,7 +1035,7 @@ function transformSqlBodyString( const pieces: string[] = []; for (const stmt of stmts) { if (!stmt?.stmt) continue; - walkSql(stmt.stmt, visitor); + walkSqlAst(stmt.stmt, visitor); pieces.push(Deparser.deparse(stmt.stmt)); } return pieces.join(';\n'); @@ -1044,6 +1044,27 @@ function transformSqlBodyString( } } +/** + * Apply schema renaming to a hydrated parse context: every top-level statement + * and every hydrated PL/pgSQL body in one walk. `walkPlpgsqlForSchemas` handles + * what the SQL visitor cannot see — PL/pgSQL-only nodes such as declared + * variable types, which carry schema-qualified names as plain strings. + */ +function transformSchemasInContext( + ctx: any, + schemaMapping: SchemaMappingInput, + result: SchemaTransformResult, + visitorOptions?: { assumeSchemasExist?: Set } +): void { + walk(ctx, createSqlVisitor(schemaMapping, result, visitorOptions)); + + for (const fn of ctx.functions ?? []) { + if (fn.plpgsql?.hydrated) { + walkPlpgsqlForSchemas(fn.plpgsql.hydrated, schemaMapping, result); + } + } +} + /** * Escape a string for use in a regular expression */ @@ -1290,26 +1311,7 @@ export function transformSqlStatement( try { const transformed = transformSync(sql, (ctx) => { - const sqlVisitor = createSqlVisitor(schemaMapping, r); - - if (ctx.sql?.stmts) { - for (const stmt of ctx.sql.stmts) { - if (stmt?.stmt) { - walkSql(stmt.stmt, sqlVisitor); - } - } - } - - for (const fn of ctx.functions) { - if (fn.plpgsql?.hydrated) { - walkPlpgsql(fn.plpgsql.hydrated, {}, { - walkSqlExpressions: true, - sqlVisitor: sqlVisitor - }); - - walkPlpgsqlForSchemas(fn.plpgsql.hydrated, schemaMapping, r); - } - } + transformSchemasInContext(ctx, schemaMapping, r); }, { hydrate: true, pretty: true }); return { sql: transformed, result: r }; @@ -1362,26 +1364,7 @@ function transformSqlContentAst( let before: CapturedAsts | undefined; try { transformedBody = transformSync(transformedBody, (ctx) => { - const sqlVisitor = createSqlVisitor(schemaMapping, result, { assumeSchemasExist }); - - if (ctx.sql?.stmts) { - for (const stmt of ctx.sql.stmts) { - if (stmt?.stmt) { - walkSql(stmt.stmt, sqlVisitor); - } - } - } - - for (const fn of ctx.functions) { - if (fn.plpgsql?.hydrated) { - walkPlpgsql(fn.plpgsql.hydrated, {}, { - walkSqlExpressions: true, - sqlVisitor: sqlVisitor - }); - - walkPlpgsqlForSchemas(fn.plpgsql.hydrated, schemaMapping, result); - } - } + transformSchemasInContext(ctx, schemaMapping, result, { assumeSchemasExist }); if (roundTrip) { before = captureTransformAsts(ctx); diff --git a/packages/traverse/README.md b/packages/traverse/README.md index da189225..62dc47f6 100644 --- a/packages/traverse/README.md +++ b/packages/traverse/README.md @@ -13,7 +13,7 @@

-PostgreSQL AST traversal utilities for pgsql-parser. This package provides a visitor pattern for traversing PostgreSQL Abstract Syntax Tree (AST) nodes, similar to Babel's traverse functionality but specifically designed for PostgreSQL AST structures. +AST traversal for the pgsql-parser ecosystem: a Babel-style visitor pattern for PostgreSQL SQL ASTs **and** PL/pgSQL ASTs. Traversal only — nothing here parses SQL, so the package stays free of WASM. ## Installation @@ -21,43 +21,144 @@ PostgreSQL AST traversal utilities for pgsql-parser. This package provides a vis npm install @pgsql/traverse ``` +## Which function do I want? + +| Function | Walks | Use when | +| --- | --- | --- | +| `walk(ast, visitors, opts?)` | **any** AST — parsed script, `ParseResult`, SQL node, PL/pgSQL node | default choice | +| `walkSqlAst(ast, visitor)` | SQL AST only | you want the raw primitive, no statement context | +| `walkPlpgsqlAst(ast, visitor, opts?)` | PL/pgSQL AST only | you already have a hydrated function body | +| `traverse(ast, mutableVisitor)` | SQL AST, mutation-safe | you need to replace/insert/remove nodes | +| [`walkSql(text, ...)`](../plpgsql-parser) in `plpgsql-parser` | SQL **text** | you start from a string and need parse + hydrate | + ## Usage -### Walk API (Recommended) +### `walk` — one entry point for any AST -The `walk` function provides improved traversal with NodePath context and early return support: +`walk` dispatches on the shape of what you hand it: + +| Input | Walked | +| --- | --- | +| `{ sql, functions }` (a `ParsedScript` from `plpgsql-parser`) | every statement, then every hydrated PL/pgSQL body | +| `{ version, stmts }` (a `ParseResult`) | every statement, with statement context | +| `{ PLpgSQL_*: ... }` / `{ plpgsql_funcs: [...] }` | the PL/pgSQL body, descending into its hydrated SQL expressions | +| any other SQL node | that node | +| an array | each element | ```typescript -import { walk, NodePath } from '@pgsql/traverse'; -import type { Walker, Visitor } from '@pgsql/traverse'; +import { walk } from '@pgsql/traverse'; +import type { NodePath, Visitor, Walker } from '@pgsql/traverse'; -// Using a simple walker function -const walker: Walker = (path: NodePath) => { +// A walker function fires on every node. +walk(ast, (path) => { console.log(`Visiting ${path.tag} at path:`, path.path); - - // Return false to skip traversing children if (path.tag === 'SelectStmt') { - return false; // Skip SELECT statement children + return false; // skip this node's children } -}; +}); -walk(ast, walker); +// A visitor object fires per node tag. SQL and PL/pgSQL tags may be mixed. +walk(ast, { + SelectStmt: (path) => console.log('SELECT statement:', path.node), + RangeVar: (path) => console.log('Table:', path.node.relname), + PLpgSQL_stmt_dynexecute: (path) => console.log('dynamic EXECUTE:', path.node) +}); +``` -// Using a visitor object (recommended for multiple node types) -const visitor: Visitor = { - SelectStmt: (path) => { - console.log('SELECT statement:', path.node); - }, - RangeVar: (path) => { - console.log('Table:', path.node.relname); - console.log('Path to table:', path.path); - console.log('Parent node:', path.parent?.tag); +#### Statement context + +Every callback gets a second argument describing the statement the node belongs +to. `NodePath` tells you *where* a node sits; `WalkContext` tells you *what it is +part of* — which is what a `RangeVar` handler needs to know it is the write +target of an `UPDATE`, or that it came from inside `auth.login`'s body. + +```typescript +interface WalkContext { + stmtTag: string | null; // enclosing top-level statement, e.g. 'UpdateStmt' + stmtIndex: number; // its index in the script, -1 if unknown + isWrite: boolean; // INSERT/UPDATE/DELETE/MERGE/TRUNCATE/COPY + isRead: boolean; // SELECT/EXPLAIN/DECLARE/FETCH + insideFunction: boolean; // node came from a PL/pgSQL body + functionName: string | null; // dotted name of that function + abort(reason?: string): void; +} + +walk(parsed, { + RangeVar: (path, ctx) => { + if (ctx.isWrite) console.log('writes to', path.node.relname); + if (ctx.insideFunction) console.log('...from inside', ctx.functionName); } +}); +``` + +The reserved `statement` key fires once per top-level statement, before its +children: + +```typescript +walk(parseResult, { + statement: (path, ctx) => console.log(ctx.stmtIndex, path.tag) +}); +``` + +#### Composing visitors + +Pass an array to run several independent policies in a single pass: + +```typescript +walk(parsed, [blockedSchemas, readOnlySchemas, blockedFunctions]); +``` + +#### Skip vs. abort + +`return false` skips that node's children. `ctx.abort(reason?)` ends the entire +walk — what a validator wants when it has already decided to reject: + +```typescript +const result = walk(parsed, { + PLpgSQL_stmt_dynexecute: (_path, ctx) => ctx.abort('dynamic EXECUTE is not allowed') +}); + +result.aborted; // true +result.reason; // 'dynamic EXECUTE is not allowed' +result.reasons; // every reason, in call order +``` + +### `walkSqlAst` — the SQL-only primitive + +Recursion is derived from PostgreSQL's runtime schema, so it knows exactly which +fields hold nodes, including untagged typed fields such as +`CreatePolicyStmt.table`. `walk` is built on it. + +```typescript +import { walkSqlAst } from '@pgsql/traverse'; +import type { Visitor, Walker } from '@pgsql/traverse'; + +const visitor: Visitor = { + RangeVar: (path) => console.log('Table:', path.node.relname) }; -walk(ast, visitor); +walkSqlAst(ast, visitor); +``` + +### `walkPlpgsqlAst` — the PL/pgSQL-only primitive + +Walks the PL/pgSQL node universe (`PLpgSQL_stmt_block`, `PLpgSQL_stmt_if`, +`PLpgSQL_var`, ...), which the SQL parser never produces. A PL/pgSQL body is a +control-flow skeleton whose leaves are SQL expressions, so it bridges into +`walkSqlAst` for every hydrated expression: + +```typescript +import { walkPlpgsqlAst } from '@pgsql/traverse'; + +walkPlpgsqlAst(hydratedBody, { PLpgSQL_var: (path) => console.log(path.node.refname) }, { + walkSqlExpressions: true, + sqlVisitor: { RangeVar: (path) => console.log('table in body:', path.node.relname) } +}); ``` +Hydration itself lives in `plpgsql-parser`. + + ### NodePath Class The `NodePath` class provides rich context information: @@ -74,32 +175,6 @@ class NodePath { } ``` -### Runtime Schema Integration - -The new implementation uses PostgreSQL's runtime schema to precisely determine which fields contain Node types that need traversal, eliminating guesswork and improving accuracy. - -### Visit API - -The original `visit` function is still available for backward compatibility: - -```typescript -import { visit } from '@pgsql/traverse'; - -const visitor = { - SelectStmt: (node, ctx) => { - console.log('Found SELECT statement:', node); - console.log('Path:', ctx.path); - }, - RangeVar: (node, ctx) => { - console.log('Found table reference:', node.relname); - } -}; - -// Parse some SQL and traverse the AST -const ast = /* your parsed AST */; -visit(ast, visitor); -``` - ### Working with ParseResult ```typescript @@ -118,9 +193,9 @@ const visitor = { walk(parseResult, visitor); ``` -### Using Visitor Context +### Reading a node's position -The visitor context provides information about the current traversal state: +`NodePath` describes where a node sits in the tree: ```typescript const visitor = { @@ -136,30 +211,28 @@ const visitor = { ### Collecting Information During Traversal ```typescript -import { visit } from '@pgsql/traverse'; +import { walk } from '@pgsql/traverse'; import type { Visitor } from '@pgsql/traverse'; const tableNames: string[] = []; const columnRefs: string[] = []; const visitor: Visitor = { - RangeVar: (node) => { - if (node.relname) { - tableNames.push(node.relname); + RangeVar: (path) => { + if (path.node.relname) { + tableNames.push(path.node.relname); } }, - ColumnRef: (node) => { - if (node.fields) { - node.fields.forEach(field => { - if (field.String?.sval) { - columnRefs.push(field.String.sval); - } - }); + ColumnRef: (path) => { + for (const field of path.node.fields ?? []) { + if (field.String?.sval) { + columnRefs.push(field.String.sval); + } } } }; -visit(ast, visitor); +walk(ast, visitor); console.log('Tables referenced:', tableNames); console.log('Columns referenced:', columnRefs); @@ -167,9 +240,24 @@ console.log('Columns referenced:', columnRefs); ## API -### `walk(root, callback, parent?, keyPath?)` +### `walk(root, visitors, options?): WalkResult` -Walks the tree of PostgreSQL AST nodes using runtime schema for precise traversal. +Walks any AST — SQL, PL/pgSQL, or a parsed script — with one or more visitors. + +**Parameters:** +- `root`: the AST, parse result, or parsed script to traverse +- `visitors`: a walker function, a visitor object, or an array of either +- `options?`: + - `walkFunctionBodies` (default `true`) — walk hydrated PL/pgSQL bodies of a parsed script + - `walkSqlExpressions` (default `true`) — recurse into hydrated SQL expressions inside bodies + - `sqlVisitor` — override the visitor used for those SQL expressions + +**Returns** `{ aborted, reason?, reasons }`. + +### `walkSqlAst(root, callback, parent?, keyPath?)` + +The SQL-only primitive. Walks PostgreSQL AST nodes using the runtime schema for +precise traversal. **Parameters:** - `root`: The AST node to traverse @@ -177,29 +265,51 @@ Walks the tree of PostgreSQL AST nodes using runtime schema for precise traversa - `parent?`: Optional parent NodePath (for internal use) - `keyPath?`: Optional key path array (for internal use) -**Example:** +### `walkPlpgsqlAst(root, callback, options?, parent?, keyPath?)` + +The PL/pgSQL-only primitive. + +**Parameters:** +- `root`: The PL/pgSQL AST node to traverse +- `callback`: A walker function or visitor object keyed by `PLpgSQL_*` tags +- `options?`: `{ walkSqlExpressions?, sqlVisitor? }` + +### `traverse(root, mutableVisitor)` + +The mutation-capable walker: `enter`/`exit` hooks and sibling insert/remove/replace +through `MutablePath`. Use it when the walk needs to change the tree; `walk` is +read-only. + +### Types + +#### `WalkContext` + +Statement context threaded into every `walk` callback: + ```typescript -walk(ast, { - SelectStmt: (path) => { - // Handle SELECT statements - // Return false to skip children - }, - RangeVar: (path) => { - // Handle table references - } -}); +interface WalkContext { + readonly stmtTag: string | null; + readonly stmtIndex: number; + readonly isWrite: boolean; + readonly isRead: boolean; + readonly insideFunction: boolean; + readonly functionName: string | null; + abort(reason?: string): void; +} ``` -### `visit(node, visitor, ctx?)` (Legacy) +#### `WalkResult` -Recursively visits a PostgreSQL AST node, calling any matching visitor functions. Maintained for backward compatibility. +What `walk` returns: -**Parameters:** -- `node`: The AST node to traverse -- `visitor`: An object with visitor functions for different node types -- `ctx?`: Optional initial visitor context +```typescript +interface WalkResult { + aborted: boolean; // a visitor called ctx.abort() + reason?: string; // the first reason given + reasons: string[]; // every reason, in call order +} +``` -### Types #### `Visitor` @@ -237,18 +347,6 @@ class NodePath { } ``` -#### `VisitorContext` (Legacy) - -Context information provided to legacy visitor functions: - -```typescript -type VisitorContext = { - path: (string | number)[]; // Path to current node - parent: any; // Parent node - key: string | number; // Key in parent node -}; -``` - #### `NodeTag` Union type of all PostgreSQL AST node type names: @@ -278,19 +376,20 @@ This package is designed to work seamlessly with the pgsql-parser ecosystem: ```typescript import { parse } from 'pgsql-parser'; -import { visit } from '@pgsql/traverse'; +import { walk } from '@pgsql/traverse'; const sql = 'SELECT name, email FROM users WHERE age > 18'; const ast = await parse(sql); -const visitor = { - RangeVar: (node) => { - console.log('Table:', node.relname); +walk(ast, { + RangeVar: (path) => { + console.log('Table:', path.node.relname); }, - ColumnRef: (node) => { - console.log('Column:', node.fields?.[0]?.String?.sval); + ColumnRef: (path) => { + console.log('Column:', path.node.fields?.[0]?.String?.sval); } -}; +}); +``` -visit(ast, visitor); -``` \ No newline at end of file +Starting from SQL text? `plpgsql-parser`'s `walkSql` does the parse and the +PL/pgSQL hydration for you, then hands off to `walk`. \ No newline at end of file diff --git a/packages/traverse/__test__/traverse.test.ts b/packages/traverse/__test__/traverse.test.ts index 1768c61e..41ee2769 100644 --- a/packages/traverse/__test__/traverse.test.ts +++ b/packages/traverse/__test__/traverse.test.ts @@ -1,5 +1,5 @@ import type { Visitor, Walker } from '../src'; -import { NodePath,visit, walk } from '../src'; +import { NodePath, walk } from '../src'; describe('traverse', () => { it('should visit SelectStmt nodes with new walk API', () => { @@ -157,16 +157,16 @@ describe('traverse', () => { expect(visitedNodes).toContain('A_Const'); }); - it('should maintain backward compatibility with visit function', () => { + it('should dispatch a visitor object per node tag', () => { const selectStmtVisited: any[] = []; const rangeVarVisited: any[] = []; - - const visitor = { - SelectStmt: (node: any, ctx: any) => { - selectStmtVisited.push({ node, ctx }); + + const visitor: Visitor = { + SelectStmt: (path: NodePath) => { + selectStmtVisited.push(path); }, - RangeVar: (node: any, ctx: any) => { - rangeVarVisited.push({ node, ctx }); + RangeVar: (path: NodePath) => { + rangeVarVisited.push(path); } }; @@ -184,7 +184,7 @@ describe('traverse', () => { } }; - visit(ast, visitor); + walk(ast, visitor); expect(selectStmtVisited).toHaveLength(1); expect(rangeVarVisited).toHaveLength(1); @@ -196,15 +196,15 @@ describe('traverse', () => { const rawStmtVisited: any[] = []; const selectStmtVisited: any[] = []; - const visitor = { - ParseResult: (node: any, ctx: any) => { - parseResultVisited.push({ node, ctx }); + const visitor: Visitor = { + ParseResult: (path: NodePath) => { + parseResultVisited.push(path); }, - RawStmt: (node: any, ctx: any) => { - rawStmtVisited.push({ node, ctx }); + RawStmt: (path: NodePath) => { + rawStmtVisited.push(path); }, - SelectStmt: (node: any, ctx: any) => { - selectStmtVisited.push({ node, ctx }); + SelectStmt: (path: NodePath) => { + selectStmtVisited.push(path); } }; @@ -227,7 +227,7 @@ describe('traverse', () => { } }; - visit(parseResult, visitor); + walk(parseResult, visitor); expect(parseResultVisited).toHaveLength(1); expect(rawStmtVisited).toHaveLength(1); @@ -237,12 +237,12 @@ describe('traverse', () => { expect(selectStmtVisited[0].node.limitOption).toBe('LIMIT_OPTION_DEFAULT'); }); - it('should provide correct visitor context', () => { - const contexts: any[] = []; - - const visitor = { - RangeVar: (node: any, ctx: any) => { - contexts.push(ctx); + it('should provide the key path of each node', () => { + const paths: NodePath[] = []; + + const visitor: Visitor = { + RangeVar: (path: NodePath) => { + paths.push(path); } }; @@ -260,11 +260,11 @@ describe('traverse', () => { } }; - visit(ast, visitor); + walk(ast, visitor); - expect(contexts).toHaveLength(1); - expect(contexts[0].path).toEqual(['fromClause', 0]); - expect(contexts[0].key).toBe(0); + expect(paths).toHaveLength(1); + expect(paths[0].path).toEqual(['fromClause', 0]); + expect(paths[0].key).toBe(0); }); it('should handle null and undefined nodes gracefully', () => { @@ -274,22 +274,22 @@ describe('traverse', () => { } }; - expect(() => visit(null as any, visitor)).not.toThrow(); - expect(() => visit(undefined as any, visitor)).not.toThrow(); - expect(() => visit('string' as any, visitor)).not.toThrow(); + expect(() => walk(null as any, visitor)).not.toThrow(); + expect(() => walk(undefined as any, visitor)).not.toThrow(); + expect(() => walk('string' as any, visitor)).not.toThrow(); }); it('should handle nested complex AST structures', () => { const visitedNodes: string[] = []; - const visitor = { - SelectStmt: () => visitedNodes.push('SelectStmt'), - ResTarget: () => visitedNodes.push('ResTarget'), - ColumnRef: () => visitedNodes.push('ColumnRef'), - A_Star: () => visitedNodes.push('A_Star'), - RangeVar: () => visitedNodes.push('RangeVar'), - A_Expr: () => visitedNodes.push('A_Expr'), - A_Const: () => visitedNodes.push('A_Const') + const visitor: Visitor = { + SelectStmt: () => { visitedNodes.push('SelectStmt'); }, + ResTarget: () => { visitedNodes.push('ResTarget'); }, + ColumnRef: () => { visitedNodes.push('ColumnRef'); }, + A_Star: () => { visitedNodes.push('A_Star'); }, + RangeVar: () => { visitedNodes.push('RangeVar'); }, + A_Expr: () => { visitedNodes.push('A_Expr'); }, + A_Const: () => { visitedNodes.push('A_Const'); } }; const complexAst = { @@ -335,7 +335,7 @@ describe('traverse', () => { } }; - visit(complexAst, visitor); + walk(complexAst, visitor); expect(visitedNodes).toContain('SelectStmt'); expect(visitedNodes).toContain('ResTarget'); @@ -349,9 +349,9 @@ describe('traverse', () => { it('should handle arrays of nodes correctly', () => { const targetListVisited: any[] = []; - const visitor = { - ResTarget: (node: any, ctx: any) => { - targetListVisited.push({ node, ctx }); + const visitor: Visitor = { + ResTarget: (path: NodePath) => { + targetListVisited.push(path); } }; @@ -382,13 +382,13 @@ describe('traverse', () => { } }; - visit(ast, visitor); + walk(ast, visitor); expect(targetListVisited).toHaveLength(2); - expect(targetListVisited[0].ctx.key).toBe(0); - expect(targetListVisited[1].ctx.key).toBe(1); - expect(targetListVisited[0].ctx.path).toEqual(['targetList', 0]); - expect(targetListVisited[1].ctx.path).toEqual(['targetList', 1]); + expect(targetListVisited[0].key).toBe(0); + expect(targetListVisited[1].key).toBe(1); + expect(targetListVisited[0].path).toEqual(['targetList', 0]); + expect(targetListVisited[1].path).toEqual(['targetList', 1]); }); it('should traverse WithClause nodes', () => { diff --git a/packages/traverse/__tests__/typed-fields.test.ts b/packages/traverse/__tests__/typed-fields.test.ts index 314dc356..d6795392 100644 --- a/packages/traverse/__tests__/typed-fields.test.ts +++ b/packages/traverse/__tests__/typed-fields.test.ts @@ -1,5 +1,5 @@ import type { Visitor } from '../src'; -import { NodePath, visit, walk } from '../src'; +import { NodePath, walk } from '../src'; // AST literals below are real libpg_query (PG18) parse output. Concrete // typed embedded fields (e.g. CreatePolicyStmt.table: RangeVar) are bare @@ -271,58 +271,3 @@ describe('walk — typed embedded fields (untagged)', () => { ]); }); }); - -describe('visit — typed embedded fields (untagged)', () => { - it('visits RangeVar for CreatePolicyStmt.table', () => { - const visited: any[] = []; - visit(createPolicyAst, { - RangeVar: (node, ctx) => { visited.push({ node, ctx }); } - }); - - expect(visited).toHaveLength(1); - expect(visited[0].node.relname).toBe('posts'); - expect(visited[0].ctx.path).toEqual(['table']); - }); - - it('visits RawStmt for the array-typed ParseResult.stmts field', () => { - const parseResult = { - ParseResult: { - version: 180004, - stmts: [{ stmt: indexStmtAst, stmt_len: 34 }] - } - }; - const rawStmts: any[] = []; - const rangeVars: any[] = []; - visit(parseResult, { - RawStmt: (node) => { rawStmts.push(node); }, - RangeVar: (node) => { rangeVars.push(node); } - }); - - expect(rawStmts).toHaveLength(1); - expect(rangeVars).toHaveLength(1); - expect(rangeVars[0].relname).toBe('posts'); - }); - - it('detects a bare libpg-query parse result at the root', () => { - const parseResult = { - version: 180004, - stmts: [{ stmt: createPolicyAst, stmt_len: 52 }] - }; - const visited: string[] = []; - visit(parseResult, { - ParseResult: () => { visited.push('ParseResult'); }, - RawStmt: () => { visited.push('RawStmt'); }, - RangeVar: (node) => { visited.push(`RangeVar:${node.relname}`); } - }); - expect(visited).toEqual(['ParseResult', 'RawStmt', 'RangeVar:posts']); - }); - - it('keeps existing behavior for tagged nodes', () => { - const visited: string[] = []; - visit(createPolicyAst, { - CreatePolicyStmt: (node) => { visited.push(`CreatePolicyStmt:${node.policy_name}`); }, - RoleSpec: (node) => { visited.push(`RoleSpec:${node.roletype}`); } - }); - expect(visited).toEqual(['CreatePolicyStmt:p', 'RoleSpec:ROLESPEC_PUBLIC']); - }); -}); diff --git a/packages/traverse/package.json b/packages/traverse/package.json index 2b275a47..34638872 100644 --- a/packages/traverse/package.json +++ b/packages/traverse/package.json @@ -1,6 +1,6 @@ { "name": "@pgsql/traverse", - "version": "18.4.0", + "version": "18.5.0", "author": "Constructive ", "description": "PostgreSQL AST traversal utilities for pgsql-parser", "main": "index.js", @@ -32,7 +32,8 @@ }, "dependencies": { "@pgsql/types": "^18.0.0", - "pg-proto-parser": "workspace:*" + "pg-proto-parser": "workspace:*", + "plpgsql-deparser": "workspace:*" }, "devDependencies": { "makage": "^0.1.8" diff --git a/packages/traverse/src/index.ts b/packages/traverse/src/index.ts index f9c4d4fc..c01f3854 100644 --- a/packages/traverse/src/index.ts +++ b/packages/traverse/src/index.ts @@ -1,4 +1,19 @@ export type { EnterExit, MutableVisitor, MutableWalker } from './mutate'; export { MutablePath, traverse } from './mutate'; -export type { NodeTag,Visitor, VisitorContext, Walker } from './traverse'; -export { NodePath,visit, walk } from './traverse'; +export type { + PlpgsqlNodeTag, + PlpgsqlVisitor, + PlpgsqlWalker, + PlpgsqlWalkOptions, +} from './plpgsql'; +export { PlpgsqlNodePath, walkPlpgsqlAst } from './plpgsql'; +export type { NodeTag, Visitor, Walker } from './traverse'; +export { NodePath, walkSqlAst } from './traverse'; +export type { + UnifiedVisitor, + UnifiedWalker, + WalkContext, + WalkOptions, + WalkResult, +} from './walk'; +export { READ_STATEMENTS, walk, WRITE_STATEMENTS } from './walk'; diff --git a/packages/traverse/src/plpgsql.ts b/packages/traverse/src/plpgsql.ts new file mode 100644 index 00000000..9f0a92a9 --- /dev/null +++ b/packages/traverse/src/plpgsql.ts @@ -0,0 +1,652 @@ +/** + * PL/pgSQL AST Traversal + * + * Provides a visitor pattern for traversing PL/pgSQL ASTs, sharing the node-path + * shape of the SQL walker in this package. A PL/pgSQL body is a skeleton of + * control flow whose leaves are SQL expressions, so hydrated expressions are + * handed to `walkSqlAst` — this is the bridge between the two node universes. + * + * Only PL/pgSQL *types* are needed here (from `plpgsql-deparser`), never the + * parser, which keeps this package free of the WASM parser dependency. + */ + +import type { NodePath as SqlNodePath,Visitor as SqlVisitor, Walker as SqlWalker } from './traverse'; +import { walkSqlAst } from './traverse'; +import type { + PLpgSQLParseResult, + PLpgSQLFunctionNode, + PLpgSQL_function, + PLpgSQLDatum, + PLpgSQLStmtNode, + PLpgSQLExprNode, + PLpgSQLTypeNode, + PLpgSQL_stmt_block, + PLpgSQL_stmt_if, + PLpgSQL_stmt_case, + PLpgSQL_stmt_loop, + PLpgSQL_stmt_while, + PLpgSQL_stmt_fori, + PLpgSQL_stmt_fors, + PLpgSQL_stmt_forc, + PLpgSQL_stmt_foreach_a, + PLpgSQL_stmt_return_query, + PLpgSQL_stmt_raise, + PLpgSQL_stmt_dynexecute, + PLpgSQL_stmt_dynfors, + PLpgSQL_stmt_open, + PLpgSQLElsifNode, + PLpgSQLCaseWhenNode, + PLpgSQLException, + PLpgSQLRaiseOption, +} from 'plpgsql-deparser'; +import type { + HydratedExprQuery, + HydratedExprSqlStmt, + HydratedExprSqlExpr, + HydratedExprAssign, +} from 'plpgsql-deparser'; + +// PL/pgSQL node tag types +export type PlpgsqlNodeTag = + | 'PLpgSQL_function' + | 'PLpgSQL_var' + | 'PLpgSQL_rec' + | 'PLpgSQL_row' + | 'PLpgSQL_recfield' + | 'PLpgSQL_type' + | 'PLpgSQL_expr' + | 'PLpgSQL_stmt_block' + | 'PLpgSQL_stmt_assign' + | 'PLpgSQL_stmt_if' + | 'PLpgSQL_stmt_case' + | 'PLpgSQL_stmt_loop' + | 'PLpgSQL_stmt_while' + | 'PLpgSQL_stmt_fori' + | 'PLpgSQL_stmt_fors' + | 'PLpgSQL_stmt_forc' + | 'PLpgSQL_stmt_foreach_a' + | 'PLpgSQL_stmt_exit' + | 'PLpgSQL_stmt_return' + | 'PLpgSQL_stmt_return_next' + | 'PLpgSQL_stmt_return_query' + | 'PLpgSQL_stmt_raise' + | 'PLpgSQL_stmt_assert' + | 'PLpgSQL_stmt_execsql' + | 'PLpgSQL_stmt_dynexecute' + | 'PLpgSQL_stmt_dynfors' + | 'PLpgSQL_stmt_getdiag' + | 'PLpgSQL_stmt_open' + | 'PLpgSQL_stmt_fetch' + | 'PLpgSQL_stmt_close' + | 'PLpgSQL_stmt_perform' + | 'PLpgSQL_stmt_call' + | 'PLpgSQL_stmt_commit' + | 'PLpgSQL_stmt_rollback' + | 'PLpgSQL_stmt_set' + | 'PLpgSQL_if_elsif' + | 'PLpgSQL_case_when' + | 'PLpgSQL_exception' + | 'PLpgSQL_condition' + | 'PLpgSQL_raise_option' + | 'PLpgSQL_diag_item'; + +export class PlpgsqlNodePath { + constructor( + public tag: TTag, + public node: any, + public parent: PlpgsqlNodePath | null = null, + public keyPath: readonly (string | number)[] = [] + ) {} + + get path(): (string | number)[] { + return [...this.keyPath]; + } + + get key(): string | number { + return this.keyPath[this.keyPath.length - 1] ?? ''; + } +} + +export type PlpgsqlWalker = ( + path: TNodePath, +) => boolean | void; + +export type PlpgsqlVisitor = { + [key: string]: PlpgsqlWalker; +}; + +export interface PlpgsqlWalkOptions { + /** + * Whether to recurse into hydrated SQL expressions using @pgsql/traverse. + * Default: true + */ + walkSqlExpressions?: boolean; + + /** + * SQL visitor to use when walking hydrated SQL expressions. + * Only used if walkSqlExpressions is true. + */ + sqlVisitor?: SqlVisitor | SqlWalker; +} + +/** + * Walks the tree of PL/pgSQL AST nodes using a visitor pattern. + * + * If a callback returns `false`, the walk will continue to the next sibling + * node, rather than recurse into the children of the current node. + * + * @param root - The PL/pgSQL AST node to traverse + * @param callback - A walker function or visitor object + * @param options - Walk options + * @param parent - Parent NodePath (for internal use) + * @param keyPath - Current key path (for internal use) + */ +export function walkPlpgsqlAst( + root: any, + callback: PlpgsqlWalker | PlpgsqlVisitor, + options: PlpgsqlWalkOptions = {}, + parent: PlpgsqlNodePath | null = null, + keyPath: readonly (string | number)[] = [], +): void { + const { walkSqlExpressions = true, sqlVisitor } = options; + + const actualCallback: PlpgsqlWalker = typeof callback === 'function' + ? callback + : (path: PlpgsqlNodePath) => { + const visitor = callback as PlpgsqlVisitor; + const visitFn = visitor[path.tag]; + return visitFn ? visitFn(path) : undefined; + }; + + if (Array.isArray(root)) { + root.forEach((node, index) => { + walkPlpgsqlAst(node, actualCallback, options, parent, [...keyPath, index]); + }); + } else if (typeof root === 'object' && root !== null) { + const keys = Object.keys(root); + + // Check if this is a PL/pgSQL node (single key starting with PLpgSQL_) + if (keys.length === 1 && keys[0].startsWith('PLpgSQL_')) { + const tag = keys[0]; + const nodeData = root[tag]; + const path = new PlpgsqlNodePath(tag, nodeData, parent, keyPath); + + if (actualCallback(path) === false) { + return; + } + + // Recurse into child nodes based on node type + walkNodeChildren(tag, nodeData, actualCallback, options, path); + } else { + // Not a PL/pgSQL node wrapper, check for nested structures + for (const key of keys) { + const value = root[key]; + if (typeof value === 'object' && value !== null) { + walkPlpgsqlAst(value, actualCallback, options, parent, [...keyPath, key]); + } + } + } + } + + // Helper function to walk into hydrated SQL expressions + function walkHydratedExpr(expr: any, exprPath: PlpgsqlNodePath) { + if (!walkSqlExpressions || !expr) return; + + // Check if this is a hydrated expression + if (expr.query && typeof expr.query === 'object' && 'kind' in expr.query) { + const hydratedQuery = expr.query as HydratedExprQuery; + + if (hydratedQuery.kind === 'sql-stmt') { + const sqlStmt = hydratedQuery as HydratedExprSqlStmt; + if (sqlStmt.parseResult && sqlVisitor) { + walkSqlAst(sqlStmt.parseResult, sqlVisitor); + } + } else if (hydratedQuery.kind === 'sql-expr') { + const sqlExpr = hydratedQuery as HydratedExprSqlExpr; + if (sqlExpr.expr && sqlVisitor) { + walkSqlAst(sqlExpr.expr, sqlVisitor); + } + } else if (hydratedQuery.kind === 'assign') { + const assignExpr = hydratedQuery as HydratedExprAssign; + if (assignExpr.targetExpr && sqlVisitor) { + walkSqlAst(assignExpr.targetExpr, sqlVisitor); + } + if (assignExpr.valueExpr && sqlVisitor) { + walkSqlAst(assignExpr.valueExpr, sqlVisitor); + } + } + } + } + + // Helper function to walk children based on node type + function walkNodeChildren( + tag: string, + nodeData: any, + cb: PlpgsqlWalker, + opts: PlpgsqlWalkOptions, + parentPath: PlpgsqlNodePath + ) { + switch (tag) { + case 'PLpgSQL_function': { + const fn = nodeData as PLpgSQL_function; + if (fn.datums) { + fn.datums.forEach((datum, i) => { + walkPlpgsqlAst(datum, cb, opts, parentPath, [...parentPath.keyPath, 'datums', i]); + }); + } + if (fn.action) { + walkPlpgsqlAst(fn.action, cb, opts, parentPath, [...parentPath.keyPath, 'action']); + } + break; + } + + case 'PLpgSQL_var': { + if (nodeData.datatype) { + walkPlpgsqlAst(nodeData.datatype, cb, opts, parentPath, [...parentPath.keyPath, 'datatype']); + } + if (nodeData.default_val) { + walkPlpgsqlAst(nodeData.default_val, cb, opts, parentPath, [...parentPath.keyPath, 'default_val']); + } + if (nodeData.cursor_explicit_expr) { + walkPlpgsqlAst(nodeData.cursor_explicit_expr, cb, opts, parentPath, [...parentPath.keyPath, 'cursor_explicit_expr']); + } + break; + } + + case 'PLpgSQL_expr': { + // This is where we recurse into SQL expressions + walkHydratedExpr(nodeData, parentPath); + break; + } + + case 'PLpgSQL_stmt_block': { + const block = nodeData as PLpgSQL_stmt_block; + if (block.body) { + block.body.forEach((stmt, i) => { + walkPlpgsqlAst(stmt, cb, opts, parentPath, [...parentPath.keyPath, 'body', i]); + }); + } + if (block.exceptions?.exc_list) { + block.exceptions.exc_list.forEach((exc, i) => { + walkPlpgsqlAst(exc, cb, opts, parentPath, [...parentPath.keyPath, 'exceptions', 'exc_list', i]); + }); + } + break; + } + + case 'PLpgSQL_stmt_assign': { + if (nodeData.expr) { + walkPlpgsqlAst(nodeData.expr, cb, opts, parentPath, [...parentPath.keyPath, 'expr']); + } + break; + } + + case 'PLpgSQL_stmt_if': { + const ifStmt = nodeData as PLpgSQL_stmt_if; + if (ifStmt.cond) { + walkPlpgsqlAst(ifStmt.cond, cb, opts, parentPath, [...parentPath.keyPath, 'cond']); + } + if (ifStmt.then_body) { + ifStmt.then_body.forEach((stmt, i) => { + walkPlpgsqlAst(stmt, cb, opts, parentPath, [...parentPath.keyPath, 'then_body', i]); + }); + } + if (ifStmt.elsif_list) { + ifStmt.elsif_list.forEach((elsif, i) => { + walkPlpgsqlAst(elsif, cb, opts, parentPath, [...parentPath.keyPath, 'elsif_list', i]); + }); + } + if (ifStmt.else_body) { + ifStmt.else_body.forEach((stmt, i) => { + walkPlpgsqlAst(stmt, cb, opts, parentPath, [...parentPath.keyPath, 'else_body', i]); + }); + } + break; + } + + case 'PLpgSQL_if_elsif': { + if (nodeData.cond) { + walkPlpgsqlAst(nodeData.cond, cb, opts, parentPath, [...parentPath.keyPath, 'cond']); + } + if (nodeData.stmts) { + nodeData.stmts.forEach((stmt: any, i: number) => { + walkPlpgsqlAst(stmt, cb, opts, parentPath, [...parentPath.keyPath, 'stmts', i]); + }); + } + break; + } + + case 'PLpgSQL_stmt_case': { + const caseStmt = nodeData as PLpgSQL_stmt_case; + if (caseStmt.t_expr) { + walkPlpgsqlAst(caseStmt.t_expr, cb, opts, parentPath, [...parentPath.keyPath, 't_expr']); + } + if (caseStmt.case_when_list) { + caseStmt.case_when_list.forEach((when, i) => { + walkPlpgsqlAst(when, cb, opts, parentPath, [...parentPath.keyPath, 'case_when_list', i]); + }); + } + if (caseStmt.else_stmts) { + caseStmt.else_stmts.forEach((stmt, i) => { + walkPlpgsqlAst(stmt, cb, opts, parentPath, [...parentPath.keyPath, 'else_stmts', i]); + }); + } + break; + } + + case 'PLpgSQL_case_when': { + if (nodeData.expr) { + walkPlpgsqlAst(nodeData.expr, cb, opts, parentPath, [...parentPath.keyPath, 'expr']); + } + if (nodeData.stmts) { + nodeData.stmts.forEach((stmt: any, i: number) => { + walkPlpgsqlAst(stmt, cb, opts, parentPath, [...parentPath.keyPath, 'stmts', i]); + }); + } + break; + } + + case 'PLpgSQL_stmt_loop': { + const loop = nodeData as PLpgSQL_stmt_loop; + if (loop.body) { + loop.body.forEach((stmt, i) => { + walkPlpgsqlAst(stmt, cb, opts, parentPath, [...parentPath.keyPath, 'body', i]); + }); + } + break; + } + + case 'PLpgSQL_stmt_while': { + const whileStmt = nodeData as PLpgSQL_stmt_while; + if (whileStmt.cond) { + walkPlpgsqlAst(whileStmt.cond, cb, opts, parentPath, [...parentPath.keyPath, 'cond']); + } + if (whileStmt.body) { + whileStmt.body.forEach((stmt, i) => { + walkPlpgsqlAst(stmt, cb, opts, parentPath, [...parentPath.keyPath, 'body', i]); + }); + } + break; + } + + case 'PLpgSQL_stmt_fori': { + const fori = nodeData as PLpgSQL_stmt_fori; + if (fori.var) { + walkPlpgsqlAst(fori.var, cb, opts, parentPath, [...parentPath.keyPath, 'var']); + } + if (fori.lower) { + walkPlpgsqlAst(fori.lower, cb, opts, parentPath, [...parentPath.keyPath, 'lower']); + } + if (fori.upper) { + walkPlpgsqlAst(fori.upper, cb, opts, parentPath, [...parentPath.keyPath, 'upper']); + } + if (fori.step) { + walkPlpgsqlAst(fori.step, cb, opts, parentPath, [...parentPath.keyPath, 'step']); + } + if (fori.body) { + fori.body.forEach((stmt, i) => { + walkPlpgsqlAst(stmt, cb, opts, parentPath, [...parentPath.keyPath, 'body', i]); + }); + } + break; + } + + case 'PLpgSQL_stmt_fors': { + const fors = nodeData as PLpgSQL_stmt_fors; + if (fors.var) { + walkPlpgsqlAst(fors.var, cb, opts, parentPath, [...parentPath.keyPath, 'var']); + } + if (fors.query) { + walkPlpgsqlAst(fors.query, cb, opts, parentPath, [...parentPath.keyPath, 'query']); + } + if (fors.body) { + fors.body.forEach((stmt, i) => { + walkPlpgsqlAst(stmt, cb, opts, parentPath, [...parentPath.keyPath, 'body', i]); + }); + } + break; + } + + case 'PLpgSQL_stmt_forc': { + const forc = nodeData as PLpgSQL_stmt_forc; + if (forc.var) { + walkPlpgsqlAst(forc.var, cb, opts, parentPath, [...parentPath.keyPath, 'var']); + } + if (forc.argquery) { + walkPlpgsqlAst(forc.argquery, cb, opts, parentPath, [...parentPath.keyPath, 'argquery']); + } + if (forc.body) { + forc.body.forEach((stmt, i) => { + walkPlpgsqlAst(stmt, cb, opts, parentPath, [...parentPath.keyPath, 'body', i]); + }); + } + break; + } + + case 'PLpgSQL_stmt_foreach_a': { + const foreach = nodeData as PLpgSQL_stmt_foreach_a; + if (foreach.expr) { + walkPlpgsqlAst(foreach.expr, cb, opts, parentPath, [...parentPath.keyPath, 'expr']); + } + if (foreach.body) { + foreach.body.forEach((stmt, i) => { + walkPlpgsqlAst(stmt, cb, opts, parentPath, [...parentPath.keyPath, 'body', i]); + }); + } + break; + } + + case 'PLpgSQL_stmt_exit': { + if (nodeData.cond) { + walkPlpgsqlAst(nodeData.cond, cb, opts, parentPath, [...parentPath.keyPath, 'cond']); + } + break; + } + + case 'PLpgSQL_stmt_return': { + if (nodeData.expr) { + walkPlpgsqlAst(nodeData.expr, cb, opts, parentPath, [...parentPath.keyPath, 'expr']); + } + break; + } + + case 'PLpgSQL_stmt_return_next': { + if (nodeData.expr) { + walkPlpgsqlAst(nodeData.expr, cb, opts, parentPath, [...parentPath.keyPath, 'expr']); + } + break; + } + + case 'PLpgSQL_stmt_return_query': { + const retQuery = nodeData as PLpgSQL_stmt_return_query; + if (retQuery.query) { + walkPlpgsqlAst(retQuery.query, cb, opts, parentPath, [...parentPath.keyPath, 'query']); + } + if (retQuery.dynquery) { + walkPlpgsqlAst(retQuery.dynquery, cb, opts, parentPath, [...parentPath.keyPath, 'dynquery']); + } + if (retQuery.params) { + retQuery.params.forEach((param, i) => { + walkPlpgsqlAst(param, cb, opts, parentPath, [...parentPath.keyPath, 'params', i]); + }); + } + break; + } + + case 'PLpgSQL_stmt_raise': { + const raise = nodeData as PLpgSQL_stmt_raise; + if (raise.params) { + raise.params.forEach((param, i) => { + walkPlpgsqlAst(param, cb, opts, parentPath, [...parentPath.keyPath, 'params', i]); + }); + } + if (raise.options) { + raise.options.forEach((opt, i) => { + walkPlpgsqlAst(opt, cb, opts, parentPath, [...parentPath.keyPath, 'options', i]); + }); + } + break; + } + + case 'PLpgSQL_raise_option': { + if (nodeData.expr) { + walkPlpgsqlAst(nodeData.expr, cb, opts, parentPath, [...parentPath.keyPath, 'expr']); + } + break; + } + + case 'PLpgSQL_stmt_assert': { + if (nodeData.cond) { + walkPlpgsqlAst(nodeData.cond, cb, opts, parentPath, [...parentPath.keyPath, 'cond']); + } + if (nodeData.message) { + walkPlpgsqlAst(nodeData.message, cb, opts, parentPath, [...parentPath.keyPath, 'message']); + } + break; + } + + case 'PLpgSQL_stmt_execsql': { + if (nodeData.sqlstmt) { + walkPlpgsqlAst(nodeData.sqlstmt, cb, opts, parentPath, [...parentPath.keyPath, 'sqlstmt']); + } + if (nodeData.target) { + walkPlpgsqlAst(nodeData.target, cb, opts, parentPath, [...parentPath.keyPath, 'target']); + } + break; + } + + case 'PLpgSQL_stmt_dynexecute': { + const dynexec = nodeData as PLpgSQL_stmt_dynexecute; + if (dynexec.query) { + walkPlpgsqlAst(dynexec.query, cb, opts, parentPath, [...parentPath.keyPath, 'query']); + } + if (dynexec.target) { + walkPlpgsqlAst(dynexec.target, cb, opts, parentPath, [...parentPath.keyPath, 'target']); + } + if (dynexec.params) { + dynexec.params.forEach((param, i) => { + walkPlpgsqlAst(param, cb, opts, parentPath, [...parentPath.keyPath, 'params', i]); + }); + } + break; + } + + case 'PLpgSQL_stmt_dynfors': { + const dynfors = nodeData as PLpgSQL_stmt_dynfors; + if (dynfors.var) { + walkPlpgsqlAst(dynfors.var, cb, opts, parentPath, [...parentPath.keyPath, 'var']); + } + if (dynfors.query) { + walkPlpgsqlAst(dynfors.query, cb, opts, parentPath, [...parentPath.keyPath, 'query']); + } + if (dynfors.params) { + dynfors.params.forEach((param, i) => { + walkPlpgsqlAst(param, cb, opts, parentPath, [...parentPath.keyPath, 'params', i]); + }); + } + if (dynfors.body) { + dynfors.body.forEach((stmt, i) => { + walkPlpgsqlAst(stmt, cb, opts, parentPath, [...parentPath.keyPath, 'body', i]); + }); + } + break; + } + + case 'PLpgSQL_stmt_open': { + const open = nodeData as PLpgSQL_stmt_open; + if (open.argquery) { + walkPlpgsqlAst(open.argquery, cb, opts, parentPath, [...parentPath.keyPath, 'argquery']); + } + if (open.query) { + walkPlpgsqlAst(open.query, cb, opts, parentPath, [...parentPath.keyPath, 'query']); + } + if (open.dynquery) { + walkPlpgsqlAst(open.dynquery, cb, opts, parentPath, [...parentPath.keyPath, 'dynquery']); + } + if (open.params) { + open.params.forEach((param, i) => { + walkPlpgsqlAst(param, cb, opts, parentPath, [...parentPath.keyPath, 'params', i]); + }); + } + break; + } + + case 'PLpgSQL_stmt_fetch': { + if (nodeData.target) { + walkPlpgsqlAst(nodeData.target, cb, opts, parentPath, [...parentPath.keyPath, 'target']); + } + if (nodeData.expr) { + walkPlpgsqlAst(nodeData.expr, cb, opts, parentPath, [...parentPath.keyPath, 'expr']); + } + break; + } + + case 'PLpgSQL_stmt_perform': { + if (nodeData.expr) { + walkPlpgsqlAst(nodeData.expr, cb, opts, parentPath, [...parentPath.keyPath, 'expr']); + } + break; + } + + case 'PLpgSQL_stmt_call': { + if (nodeData.expr) { + walkPlpgsqlAst(nodeData.expr, cb, opts, parentPath, [...parentPath.keyPath, 'expr']); + } + if (nodeData.target) { + walkPlpgsqlAst(nodeData.target, cb, opts, parentPath, [...parentPath.keyPath, 'target']); + } + break; + } + + case 'PLpgSQL_stmt_set': { + if (nodeData.expr) { + walkPlpgsqlAst(nodeData.expr, cb, opts, parentPath, [...parentPath.keyPath, 'expr']); + } + break; + } + + case 'PLpgSQL_exception': { + if (nodeData.conditions) { + nodeData.conditions.forEach((cond: any, i: number) => { + walkPlpgsqlAst(cond, cb, opts, parentPath, [...parentPath.keyPath, 'conditions', i]); + }); + } + if (nodeData.action) { + nodeData.action.forEach((stmt: any, i: number) => { + walkPlpgsqlAst(stmt, cb, opts, parentPath, [...parentPath.keyPath, 'action', i]); + }); + } + break; + } + + // Nodes with no children to traverse + case 'PLpgSQL_rec': + case 'PLpgSQL_row': + case 'PLpgSQL_recfield': + case 'PLpgSQL_type': + case 'PLpgSQL_stmt_getdiag': + case 'PLpgSQL_stmt_close': + case 'PLpgSQL_stmt_commit': + case 'PLpgSQL_stmt_rollback': + case 'PLpgSQL_condition': + case 'PLpgSQL_diag_item': + // No children to traverse + break; + + default: + // Unknown node type - try to traverse any object/array children + for (const key in nodeData) { + const value = nodeData[key]; + if (Array.isArray(value)) { + value.forEach((item, index) => { + if (typeof item === 'object' && item !== null) { + walkPlpgsqlAst(item, cb, opts, parentPath, [...parentPath.keyPath, key, index]); + } + }); + } else if (typeof value === 'object' && value !== null) { + walkPlpgsqlAst(value, cb, opts, parentPath, [...parentPath.keyPath, key]); + } + } + } + } +} diff --git a/packages/traverse/src/traverse.ts b/packages/traverse/src/traverse.ts index aa2a27ec..d86b9758 100644 --- a/packages/traverse/src/traverse.ts +++ b/packages/traverse/src/traverse.ts @@ -35,11 +35,15 @@ export type Visitor = { /** * Walks the tree of PostgreSQL AST nodes using runtime schema for precise traversal. - * + * + * This is the SQL-only primitive: it knows nothing about PL/pgSQL bodies or + * statement context. Prefer `walk`, which dispatches over any AST shape and + * threads a {@link WalkContext} through the visitors. + * * If a callback returns `false`, the walk will continue to the next sibling * node, rather than recurse into the children of the current node. */ -export function walk( +export function walkSqlAst( root: any, callback: Walker | Visitor, parent: NodePath | null = null, @@ -55,7 +59,7 @@ export function walk( if (Array.isArray(root)) { root.forEach((node, index) => { - walk(node, actualCallback, parent, [...keyPath, index]); + walkSqlAst(node, actualCallback, parent, [...keyPath, index]); }); } else if (typeof root === 'object' && root !== null) { const keys = Object.keys(root); @@ -71,7 +75,7 @@ export function walk( } } for (const key of keys) { - walk(root[key], actualCallback, parent, [...keyPath, key]); + walkSqlAst(root[key], actualCallback, parent, [...keyPath, key]); } } } @@ -138,11 +142,11 @@ function walkNode( if (Array.isArray(value)) { value.forEach((item, index) => { if (typeof item === 'object' && item !== null) { - walk(item, actualCallback, path, [...path.keyPath, key, index]); + walkSqlAst(item, actualCallback, path, [...path.keyPath, key, index]); } }); } else if (typeof value === 'object' && value !== null) { - walk(value, actualCallback, path, [...path.keyPath, key]); + walkSqlAst(value, actualCallback, path, [...path.keyPath, key]); } } } @@ -168,81 +172,6 @@ function walkFieldValue( ) { walkNode(declaredType, value, actualCallback, parent, keyPath); } else { - walk(value, actualCallback, parent, keyPath); - } -} - -export type VisitorContext = { - path: (string | number)[]; - parent: any; - key: string | number; -}; - -export function visit( - node: any, - visitor: { [key: string]: (node: any, ctx: VisitorContext) => void }, - ctx: VisitorContext = { path: [], parent: null, key: '' } -): void { - if (node == null || typeof node !== 'object') return; - - const rootTag = detectUntaggedRootTag(node); - if (rootTag) { - visitNode(rootTag, node, visitor, ctx); - return; - } - - const nodeType = Object.keys(node)[0] as string; - const nodeData = node[nodeType]; - - visitNode(nodeType, nodeData, visitor, ctx); -} - -function visitNode( - nodeType: string, - nodeData: any, - visitor: { [key: string]: (node: any, ctx: VisitorContext) => void }, - ctx: VisitorContext -): void { - const visitFn = visitor[nodeType]; - if (visitFn) { - visitFn(nodeData, ctx); - } - - const nodeSpec = schemaMap.get(nodeType); - - for (const key in nodeData) { - const value = (nodeData as any)[key]; - const field = nodeSpec?.fields.find((f) => f.name === key); - const concreteType = - field && field.type !== 'Node' && schemaMap.has(field.type) - ? field.type - : null; - if (Array.isArray(value)) { - value.forEach((item, index) => { - if (typeof item === 'object' && item !== null) { - const itemCtx: VisitorContext = { - parent: value, - key: index, - path: [...ctx.path, key, index], - }; - if (concreteType && !isTaggedNode(item)) { - visitNode(concreteType, item, visitor, itemCtx); - } else if (Object.keys(item).length === 1) { - visit(item, visitor, itemCtx); - } - } - }); - } else if (typeof value === 'object' && value !== null) { - const valueCtx: VisitorContext = { - parent: nodeData, - key, - path: [...ctx.path, key], - }; - if (concreteType && !isTaggedNode(value)) { - visitNode(concreteType, value, visitor, valueCtx); - } else if (Object.keys(value).length === 1) { - visit(value, visitor, valueCtx); - } - } + walkSqlAst(value, actualCallback, parent, keyPath); } } diff --git a/packages/traverse/src/walk.ts b/packages/traverse/src/walk.ts new file mode 100644 index 00000000..ad474722 --- /dev/null +++ b/packages/traverse/src/walk.ts @@ -0,0 +1,369 @@ +/** + * Unified AST traversal. + * + * `walk` is the single entry point for every AST shape in the ecosystem — a + * parsed script, a `ParseResult`, a lone SQL node, or a PL/pgSQL function body — + * layered over the two primitives in this package: + * + * - `walkSqlAst` SQL nodes, driven by the runtime schema + * - `walkPlpgsqlAst` PL/pgSQL nodes, bridging into hydrated SQL expressions + * + * On top of dispatch it adds the three things the primitives cannot provide, + * because they see one node at a time and never the statement it belongs to: + * + * 1. a {@link WalkContext} on every callback (`stmtTag`, `isWrite`, `isRead`, + * `insideFunction`, `functionName`), refined by the nearest enclosing + * statement — so an `INSERT` nested in a CTE or in a PL/pgSQL body reports + * `isWrite`, not the context of whatever statement contains it + * 2. visitor composition — N visitors in a single pass, each free to mix SQL + * and `PLpgSQL_*` tags in the same object + * 3. `ctx.abort()`, which ends the whole walk (returning `false` from a + * handler still means "skip this node's children") + * + * Parsing lives one layer up: `plpgsql-parser`'s `walkSql(sqlText, ...)` parses + * and hydrates, then delegates here. + */ + +import type { PlpgsqlNodePath, PlpgsqlVisitor,PlpgsqlWalker } from './plpgsql'; +import { walkPlpgsqlAst } from './plpgsql'; +import type { Visitor as SqlVisitor, Walker as SqlWalker } from './traverse'; +import { NodePath, walkSqlAst } from './traverse'; + +/** Statement tags that write data. */ +export const WRITE_STATEMENTS: ReadonlySet = new Set([ + 'InsertStmt', + 'UpdateStmt', + 'DeleteStmt', + 'MergeStmt', + 'TruncateStmt', + 'CopyStmt', +]); + +/** Statement tags that only read data. */ +export const READ_STATEMENTS: ReadonlySet = new Set([ + 'SelectStmt', + 'ExplainStmt', + 'DeclareCursorStmt', + 'FetchStmt', +]); + +/** + * Statement-level context threaded through every visitor callback. The node + * path tells you *where* a node sits; the context tells you *what it is part + * of* — which statement, whether that statement writes, and whether you are + * inside a function body. + */ +export interface WalkContext { + /** Tag of the nearest enclosing statement, or `null` when walking a bare node. */ + readonly stmtTag: string | null; + /** Index of the enclosing statement within the script, or `-1` if unknown. */ + readonly stmtIndex: number; + /** Whether the enclosing statement writes data. */ + readonly isWrite: boolean; + /** Whether the enclosing statement only reads data. */ + readonly isRead: boolean; + /** Whether this node came from inside a PL/pgSQL function body. */ + readonly insideFunction: boolean; + /** Dotted name of the enclosing function, when known. */ + readonly functionName: string | null; + /** End the entire walk. Unlike returning `false`, no further nodes are visited. */ + abort(reason?: string): void; +} + +export type UnifiedWalker = ( + path: NodePath | PlpgsqlNodePath, + ctx: WalkContext, +) => boolean | void; + +/** + * Visitor keyed by node tag. SQL tags (`SelectStmt`, `RangeVar`, ...) and + * PL/pgSQL tags (`PLpgSQL_stmt_dynexecute`, ...) may be mixed freely. The + * reserved `statement` key fires once per top-level statement, before its + * children are visited. + */ +export type UnifiedVisitor = { + [tag: string]: UnifiedWalker; +}; + +export interface WalkOptions { + /** Walk hydrated PL/pgSQL function bodies of a parsed script. Default: true */ + walkFunctionBodies?: boolean; + /** Recurse into hydrated SQL expressions inside PL/pgSQL bodies. Default: true */ + walkSqlExpressions?: boolean; + /** + * Visitor for hydrated SQL expressions inside PL/pgSQL bodies. Defaults to + * the visitors passed to `walk`, so one visitor object covers both universes. + */ + sqlVisitor?: SqlVisitor | SqlWalker; +} + +export interface WalkResult { + /** Whether a visitor called `ctx.abort()`. */ + aborted: boolean; + /** The first abort reason, if one was given. */ + reason?: string; + /** Every abort reason, in call order. */ + reasons: string[]; +} + +/** Thrown internally by `ctx.abort()` to unwind the primitives, which have no abort of their own. */ +class AbortWalk extends Error { + constructor() { + super('walk aborted'); + this.name = 'AbortWalk'; + } +} + +const RESERVED_STATEMENT_KEY = 'statement'; + +/** + * Walk any AST with one or more visitors. + * + * Dispatch is by shape: + * + * | Input | Walked | + * | ---------------------------------------- | --------------------------------------------- | + * | `{ sql, functions }` (parsed script) | every statement, then every hydrated body | + * | `{ stmts: [...] }` (parse result) | every statement, with statement context | + * | `{ PLpgSQL_*: ... }` | PL/pgSQL body, descending into hydrated SQL | + * | any other SQL node | the SQL node | + * | an array | each element | + */ +export function walk( + root: any, + visitors: UnifiedVisitor | UnifiedWalker | Array, + options: WalkOptions = {}, +): WalkResult { + const list = Array.isArray(visitors) ? visitors : [visitors]; + const result: WalkResult = { aborted: false, reason: undefined, reasons: [] }; + + const ctxBase = { + abort(reason?: string) { + if (reason !== undefined) { + result.reasons.push(reason); + } + if (!result.aborted) { + result.aborted = true; + result.reason = reason; + } + throw new AbortWalk(); + }, + }; + + const makeContext = (fields: Omit): WalkContext => + Object.assign({}, fields, ctxBase); + + /** + * Statement context per node, derived from the path chain and memoized. + * + * A node's context is its parent's, except when the node is itself a + * statement: then it becomes the enclosing statement for its whole subtree. + * Nesting is what makes this necessary — the write target of an `INSERT` + * inside a CTE, a rule action, or a PL/pgSQL body is a write, even though + * the statement the script starts with may only read. + */ + const contexts = new WeakMap(); + + const contextFor = ( + path: NodePath | PlpgsqlNodePath, + base: WalkContext, + ): WalkContext => { + const cached = contexts.get(path); + if (cached) return cached; + + const parent = (path as { parent?: NodePath | PlpgsqlNodePath | null }).parent; + const inherited = parent ? contextFor(parent, base) : base; + const tag = path.tag; + const isWrite = WRITE_STATEMENTS.has(tag); + const isRead = READ_STATEMENTS.has(tag); + + const ctx = + isWrite || isRead + ? makeContext({ + stmtTag: tag, + stmtIndex: inherited.stmtIndex, + isWrite, + isRead, + insideFunction: inherited.insideFunction, + functionName: inherited.functionName, + }) + : inherited; + + contexts.set(path, ctx); + return ctx; + }; + + /** Fire every visitor for one node. Returns false if any asks to skip children. */ + const fire = (path: NodePath | PlpgsqlNodePath, ctx: WalkContext): boolean => { + let descend = true; + for (const visitor of list) { + const handler: UnifiedWalker | undefined = + typeof visitor === 'function' ? visitor : visitor[path.tag]; + if (!handler) continue; + if (handler(path, ctx) === false) descend = false; + } + return descend; + }; + + /** + * Fire the reserved `statement` hook. Bare walker functions are called for + * every node already, so they are not called again here. + */ + const fireStatement = (path: NodePath, ctx: WalkContext): boolean => { + let descend = true; + for (const visitor of list) { + if (typeof visitor === 'function') continue; + const handler = visitor[RESERVED_STATEMENT_KEY]; + if (handler && handler(path, ctx) === false) descend = false; + } + return descend; + }; + + const sqlCallback = (base: WalkContext): SqlWalker => (path: NodePath) => + fire(path, contextFor(path, base)) ? undefined : false; + + const plpgsqlCallback = (base: WalkContext): PlpgsqlWalker => (path: PlpgsqlNodePath) => + fire(path, contextFor(path, base)) ? undefined : false; + + const bareContext = (): WalkContext => + makeContext({ + stmtTag: null, + stmtIndex: -1, + isWrite: false, + isRead: false, + insideFunction: false, + functionName: null, + }); + + /** + * Walk a whole `ParseResult` in one `walkSqlAst` pass — the SQL walker knows + * how to reach untagged `RawStmt` entries and every other typed field, which + * hand-rolled per-statement iteration does not. Statement context is derived + * on the fly: traversal is depth-first and statements are sequential, so + * every node after a `RawStmt` belongs to that statement. + */ + const walkParseResult = (parseResult: any): void => { + const base = bareContext(); + + const callback: SqlWalker = (path: NodePath) => { + if (path.tag !== 'RawStmt') { + return fire(path, contextFor(path, base)) ? undefined : false; + } + + const stmt = path.node?.stmt; + const stmtTag = stmt ? Object.keys(stmt)[0] : null; + const last = path.keyPath[path.keyPath.length - 1]; + const ctx = makeContext({ + stmtTag, + stmtIndex: typeof last === 'number' ? last : -1, + isWrite: stmtTag != null && WRITE_STATEMENTS.has(stmtTag), + isRead: stmtTag != null && READ_STATEMENTS.has(stmtTag), + insideFunction: false, + functionName: null, + }); + contexts.set(path, ctx); + + let descend = fire(path, ctx); + if (stmt && stmtTag) { + const stmtPath = new NodePath(stmtTag, stmt[stmtTag], path, [...path.keyPath, 'stmt']); + if (!fireStatement(stmtPath, ctx)) descend = false; + } + return descend ? undefined : false; + }; + + walkSqlAst(parseResult, callback); + }; + + const walkBody = (hydrated: any, ctx: WalkContext): void => { + walkPlpgsqlAst(hydrated, plpgsqlCallback(ctx) as PlpgsqlWalker | PlpgsqlVisitor, { + walkSqlExpressions: options.walkSqlExpressions ?? true, + sqlVisitor: options.sqlVisitor ?? sqlCallback(ctx), + }); + }; + + const walkAny = (node: any): void => { + if (node == null || typeof node !== 'object') return; + + if (Array.isArray(node)) { + node.forEach(walkAny); + return; + } + + if (isParsedScript(node)) { + if (node.sql) walkParseResult(node.sql); + if (options.walkFunctionBodies ?? true) { + for (const fn of node.functions ?? []) { + const hydrated = fn?.plpgsql?.hydrated; + if (!hydrated) continue; + walkBody( + hydrated, + makeContext({ + stmtTag: 'CreateFunctionStmt', + stmtIndex: -1, + isWrite: false, + isRead: false, + insideFunction: true, + functionName: resolveFunctionName(fn.stmt), + }), + ); + } + } + return; + } + + // libpg-query returns `{ version, stmts }`; the tagged form appears in + // fixtures and in ASTs round-tripped through the protobuf types. + if (Array.isArray(node.stmts)) { + walkParseResult(node); + return; + } + if (Array.isArray(node.ParseResult?.stmts)) { + walkParseResult(node); + return; + } + + if (isPlpgsqlNode(node)) { + walkBody(node, bareContext()); + return; + } + + walkSqlAst(node, sqlCallback(bareContext())); + }; + + try { + walkAny(root); + } catch (err) { + if (!(err instanceof AbortWalk)) throw err; + } + + return result; +} + +function isParsedScript(node: any): boolean { + return Array.isArray(node.functions) && node.sql != null; +} + +/** + * A hydrated body arrives either as a `{ plpgsql_funcs: [...] }` parse result or + * as a lone tagged node. Both belong to the PL/pgSQL walker: the SQL walker + * would descend them generically and never bridge their SQL expressions. + */ +function isPlpgsqlNode(node: any): boolean { + if (Array.isArray(node.plpgsql_funcs)) return true; + const keys = Object.keys(node); + return keys.length === 1 && keys[0].startsWith('PLpgSQL_'); +} + +/** + * `ParsedFunction.stmt` is the inner `CreateFunctionStmt` node, but accept the + * tagged form too so callers can hand us either. + */ +function resolveFunctionName(stmt: any): string | null { + const funcname = stmt?.funcname ?? stmt?.CreateFunctionStmt?.funcname; + if (!funcname) return null; + const name = funcname + .map((part: any) => part?.String?.sval ?? '') + .filter(Boolean) + .join('.'); + return name || null; +} diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index b871e3f4..c38c1746 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -453,6 +453,9 @@ importers: pg-proto-parser: specifier: workspace:* version: link:../proto-parser/dist + plpgsql-deparser: + specifier: workspace:* + version: link:../plpgsql-deparser/dist devDependencies: makage: specifier: ^0.1.8