Merge pull request #9 from supabase-community/feat/node-utils

gregnr · web-flow · commit 165ba3bf24fc · 2025-06-09T11:51:59.000-06:00
feat: node utils
diff --git a/README.md b/README.md
@@ -231,6 +231,92 @@ If the parse fails, `PgParser` will return an `error` of type `ParseError` with
 
   **Note:** This is relative to the entire SQL string, not just the statement being parsed or line numbers within a statement. If you are parsing a multi-statement query, the position will be relative to the entire query string, where newlines are counted as single characters.
 
+### Utility functions
+
+The following utility functions are available to help with parsing:
+
+#### `unwrapNode()`
+
+Extracts the node type and nested value while preserving type information.
+
+```typescript
+import { unwrapNode } from '@supabase/pg-parser';
+
+const wrappedStatement = result.tree.stmts[0].stmt;
+// { SelectStmt: { ... } }
+
+const { type, node } = unwrapNode(wrappedStatement);
+// { type: 'SelectStmt', node: { ... } }
+```
+
+**Background:** The AST structure produced by Postgres ([libpg_query](https://github.com/pganalyze/libpg_query)) can be complex due to nesting. For example, a `SELECT` statement is represented as:
+
+```typescript
+{
+  version: 170004,
+  stmts: [
+    {
+      stmt: {
+        SelectStmt: {
+          targetList: [ ... ],
+          fromClause: [ ... ],
+          whereClause: { ... },
+          ...
+        }
+      }
+    }
+  ]
+}
+```
+
+In order to determine which statement type is being parsed, you'd have to use the `in` operator to check for the presence of a specific key:
+
+```typescript
+const wrappedStatement = result.tree.stmts[0].stmt;
+// e.g. { SelectStmt: { ... } }
+
+if ('SelectStmt' in wrappedStatement) {
+  // It's a SELECT statement
+  const selectStmt = wrappedStatement.SelectStmt;
+} else if ('InsertStmt' in wrappedStatement) {
+  // It's an INSERT statement
+  const insertStmt = wrappedStatement.InsertStmt;
+} else if ('UpdateStmt' in wrappedStatement) {
+  // It's an UPDATE statement
+  const updateStmt = wrappedStatement.UpdateStmt;
+} else if ('DeleteStmt' in wrappedStatement) {
+  // It's a DELETE statement
+  const deleteStmt = wrappedStatement.DeleteStmt;
+}
+```
+
+`unwrapNode()` simplifies this by extracting the node type and nested value in a single step:
+
+```typescript
+const { type, node } = unwrapNode(wrappedStatement);
+```
+
+You can then use `type` to determine which statement it is and narrow the type of `node` accordingly:
+
+```typescript
+const { type, node } = unwrapNode(wrappedStatement);
+
+switch (type) {
+  case 'SelectStmt':
+    // Now `node` is narrowed to `SelectStmt`
+    break;
+  case 'InsertStmt':
+    // Now `node` is narrowed to `InsertStmt`
+    break;
+  case 'UpdateStmt':
+    // Now `node` is narrowed to `UpdateStmt`
+    break;
+  case 'DeleteStmt':
+    // Now `node` is narrowed to `DeleteStmt`
+    break;
+}
+```
+
 ## Roadmap
 
 - [ ] Deparse SQL queries (AST -> SQL)
diff --git a/packages/pg-parser/src/index.ts b/packages/pg-parser/src/index.ts
@@ -15,5 +15,6 @@ export {
   getSupportedVersions,
   isParseResultVersion,
   isSupportedVersion,
+  unwrapNode,
   unwrapParseResult,
 } from './util.js';
diff --git a/packages/pg-parser/src/pg-parser.test.ts b/packages/pg-parser/src/pg-parser.test.ts
@@ -4,7 +4,12 @@ import { stripIndent } from 'common-tags';
 import { describe, expect, it } from 'vitest';
 import { PgParser } from './pg-parser.js';
 import type { ParseResult } from './types/index.js';
-import { isParseResultVersion, unwrapParseResult } from './util.js';
+import {
+  assertAndUnwrapNode,
+  assertDefined,
+  isParseResultVersion,
+  unwrapParseResult,
+} from './util.js';
 
 import sqlDump from '../test/fixtures/dump.sql';
 
@@ -36,10 +41,7 @@ describe('pg-parser', () => {
     const pgParser = new PgParser();
     const result = await unwrapParseResult(pgParser.parse(sqlDump));
 
-    if (!result.stmts) {
-      throw new Error('stmts not found');
-    }
-
+    assertDefined(result.stmts, 'stmts not found');
     expect(result.stmts.length).toBeGreaterThan(0);
   });
 
@@ -79,106 +81,55 @@ describe('pg-parser', () => {
 
     expect(result.version).toBe(170004);
 
-    if (!result.stmts) {
-      throw new Error('stmts not found');
-    }
-
-    const [firstStmt] = result.stmts;
-    if (!firstStmt) {
-      throw new Error('stmts are empty');
-    }
-
     // Use type narrowing to ensure the result is of the expected type
     // These should produce compile-time errors if the types are incorrect
-    if (!firstStmt.stmt) {
-      throw new Error('stmt not found');
-    }
-
-    if (!('SelectStmt' in firstStmt.stmt) || !firstStmt.stmt.SelectStmt) {
-      throw new Error('SelectStmt not found');
-    }
-
-    if (!('targetList' in firstStmt.stmt.SelectStmt)) {
-      throw new Error('targetList not found');
-    }
-
-    if (!Array.isArray(firstStmt.stmt.SelectStmt.targetList)) {
-      throw new Error('targetList is not an array');
-    }
-
-    const [firstTarget] = firstStmt.stmt.SelectStmt.targetList;
-    if (!firstTarget) {
-      throw new Error('targetList is empty');
-    }
-
-    if (!('ResTarget' in firstTarget) || !firstTarget.ResTarget) {
-      throw new Error('ResTarget not found');
-    }
+    assertDefined(result.stmts, 'stmts not found');
 
-    expect(firstTarget.ResTarget.name).toBe('sum');
+    const [firstStmt] = result.stmts;
+    assertDefined(firstStmt, 'stmts are empty');
 
-    if (!firstTarget.ResTarget.val) {
-      throw new Error('val not found');
-    }
+    assertDefined(firstStmt.stmt, 'stmt not found');
 
-    if (
-      !('A_Expr' in firstTarget.ResTarget.val) ||
-      !firstTarget.ResTarget.val.A_Expr
-    ) {
-      throw new Error('A_Expr not found');
-    }
+    const selectStmt = assertAndUnwrapNode(firstStmt.stmt, 'SelectStmt');
+    assertDefined(selectStmt.targetList, 'targetList not found');
 
-    expect(firstTarget.ResTarget.val.A_Expr.kind).toBe('AEXPR_OP');
+    const [firstTarget] = selectStmt.targetList;
+    assertDefined(firstTarget, 'targetList is empty');
 
-    if (!firstTarget.ResTarget.val.A_Expr.name) {
-      throw new Error('name not found');
-    }
+    const resTarget = assertAndUnwrapNode(firstTarget, 'ResTarget');
 
-    const [firstName] = firstTarget.ResTarget.val.A_Expr.name;
+    expect(resTarget.name).toBe('sum');
+    assertDefined(resTarget.val, 'val not found');
 
-    if (!firstName) {
-      throw new Error('expression name is empty');
-    }
+    const aExpr = assertAndUnwrapNode(resTarget.val, 'A_Expr');
 
-    if (!('String' in firstName) || !firstName.String) {
-      throw new Error('expression name should be String');
-    }
+    expect(aExpr.kind).toBe('AEXPR_OP');
+    assertDefined(aExpr.name, 'name not found');
 
-    expect(firstName.String.sval).toBe('+');
+    const [firstName] = aExpr.name;
+    assertDefined(firstName, 'expression name is empty');
 
-    if (!firstTarget.ResTarget.val.A_Expr.lexpr) {
-      throw new Error('lexpr not found');
-    }
+    const name = assertAndUnwrapNode(firstName, 'String');
 
-    if (
-      !('A_Const' in firstTarget.ResTarget.val.A_Expr.lexpr) ||
-      !firstTarget.ResTarget.val.A_Expr.lexpr.A_Const
-    ) {
-      throw new Error('left side of expression should be A_Const');
-    }
+    expect(name.sval).toBe('+');
 
-    if (!firstTarget.ResTarget.val.A_Expr.lexpr.A_Const.ival) {
-      throw new Error('expected left side constant to be an integer');
-    }
-
-    expect(firstTarget.ResTarget.val.A_Expr.lexpr.A_Const.ival.ival).toBe(1);
-
-    if (!firstTarget.ResTarget.val.A_Expr.rexpr) {
-      throw new Error('rexpr not found');
-    }
+    assertDefined(aExpr.lexpr, 'left side of expression not found');
 
-    if (
-      !('A_Const' in firstTarget.ResTarget.val.A_Expr.rexpr) ||
-      !firstTarget.ResTarget.val.A_Expr.rexpr.A_Const
-    ) {
-      throw new Error('right side of expression should be A_Const');
-    }
+    const leftConst = assertAndUnwrapNode(aExpr.lexpr, 'A_Const');
+    assertDefined(
+      leftConst.ival,
+      'expected left side constant to be an integer'
+    );
+    expect(leftConst.ival.ival).toBe(1);
 
-    if (!firstTarget.ResTarget.val.A_Expr.rexpr.A_Const.ival) {
-      throw new Error('expected right side constant to be an integer');
-    }
+    assertDefined(aExpr.rexpr, 'right side of expression not found');
 
-    expect(firstTarget.ResTarget.val.A_Expr.rexpr.A_Const.ival.ival).toBe(1);
+    const rightConst = assertAndUnwrapNode(aExpr.rexpr, 'A_Const');
+    assertDefined(
+      rightConst.ival,
+      'expected right side constant to be an integer'
+    );
+    expect(rightConst.ival.ival).toBe(1);
   });
 
   it('narrows type using isParseResultVersion', async () => {
diff --git a/packages/pg-parser/src/types/15.ts b/packages/pg-parser/src/types/15.ts
@@ -1,5 +1,7 @@
-import type { ParseResult } from '../../wasm/15/pg-parser-types.js';
+import type { Node, ParseResult } from '../../wasm/15/pg-parser-types.js';
+
 export type ParseResult15 = ParseResult;
+export type Node15 = Node;
 
-export * from '../../wasm/15/pg-parser-types.js';
 export * from '../../wasm/15/pg-parser-enums.js';
+export * from '../../wasm/15/pg-parser-types.js';
diff --git a/packages/pg-parser/src/types/16.ts b/packages/pg-parser/src/types/16.ts
@@ -1,5 +1,7 @@
-import type { ParseResult } from '../../wasm/16/pg-parser-types.js';
+import type { Node, ParseResult } from '../../wasm/16/pg-parser-types.js';
+
 export type ParseResult16 = ParseResult;
+export type Node16 = Node;
 
-export * from '../../wasm/16/pg-parser-types.js';
 export * from '../../wasm/16/pg-parser-enums.js';
+export * from '../../wasm/16/pg-parser-types.js';
diff --git a/packages/pg-parser/src/types/17.ts b/packages/pg-parser/src/types/17.ts
@@ -1,5 +1,7 @@
-import type { ParseResult } from '../../wasm/17/pg-parser-types.js';
+import type { Node, ParseResult } from '../../wasm/17/pg-parser-types.js';
+
 export type ParseResult17 = ParseResult;
+export type Node17 = Node;
 
-export * from '../../wasm/17/pg-parser-types.js';
 export * from '../../wasm/17/pg-parser-enums.js';
+export * from '../../wasm/17/pg-parser-types.js';
diff --git a/packages/pg-parser/src/types/index.ts b/packages/pg-parser/src/types/index.ts
@@ -2,9 +2,9 @@ import type { MainModule as MainModule15 } from '../../wasm/15/pg-parser.js';
 import type { MainModule as MainModule16 } from '../../wasm/16/pg-parser.js';
 import type { MainModule as MainModule17 } from '../../wasm/17/pg-parser.js';
 
-import type { ParseResult15 } from './15.js';
-import type { ParseResult16 } from './16.js';
-import type { ParseResult17 } from './17.js';
+import type { Node15, ParseResult15 } from './15.js';
+import type { Node16, ParseResult16 } from './16.js';
+import type { Node17, ParseResult17 } from './17.js';
 
 import type { SUPPORTED_VERSIONS } from '../constants.js';
 import type { ParseError } from '../errors.js';
@@ -23,13 +23,23 @@ type ParseResultVersionMap = {
   17: ParseResult17;
 };
 
+type NodeVersionMap = {
+  15: Node15;
+  16: Node16;
+  17: Node17;
+};
+
 export type MainModule<Version extends SupportedVersion> =
   ModuleVersionMap[Version];
 export type PgParserModule<T extends SupportedVersion> = (
   options?: unknown
 ) => Promise<MainModule<T>>;
 
-export type ParseResult<T extends SupportedVersion> = ParseResultVersionMap[T];
+export type ParseResult<T extends SupportedVersion = SupportedVersion> =
+  ParseResultVersionMap[T];
+
+export type Node<Version extends SupportedVersion = SupportedVersion> =
+  NodeVersionMap[Version];
 
 export type WrappedParseSuccess<Version extends SupportedVersion> = {
   tree: ParseResult<Version>;
diff --git a/packages/pg-parser/src/util.test.ts b/packages/pg-parser/src/util.test.ts
diff --git a/packages/pg-parser/src/util.ts b/packages/pg-parser/src/util.ts
