Merge pull request #314 from constructive-io/feat/traverse-bare-parse-result-roots

feat(traverse): dispatch on bare ParseResult/ScanResult roots + payload guards

Author: pyramation

packages/traverse/__tests__/typed-fields.test.ts

@@ -229,6 +229,33 @@ describe('walk — typed embedded fields (untagged)', () => {
     expect(rawStmts).toHaveLength(2);
   });
 
+  it('detects a bare libpg-query parse result at the root', () => {
+    // libpg-query parseSync() returns {version, stmts} with no ParseResult wrapper
+    const parseResult = {
+      version: 180004,
+      stmts: [{ stmt: createPolicyAst, stmt_len: 52 }]
+    };
+    const tags: string[] = [];
+    walk(parseResult, (path) => { tags.push(path.tag); });
+
+    expect(tags[0]).toBe('ParseResult');
+    expect(tags[1]).toBe('RawStmt');
+    expect(tags).toContain('CreatePolicyStmt');
+    expect(tags).toContain('RangeVar');
+  });
+
+  it('does not treat non-root {version, stmts} shapes as ParseResult', () => {
+    const tags: string[] = [];
+    walk({ wrapper: { version: 180004, stmts: [] as any[] } }, (path) => { tags.push(path.tag); });
+    expect(tags).toEqual([]);
+  });
+
+  it('ignores tagged wrappers whose payload is not an object', () => {
+    const visited: NodePath[] = [];
+    walk({ Weird: 'not-an-object' }, (path) => { visited.push(path); });
+    expect(visited).toEqual([]);
+  });
+
   it('keeps existing behavior for tagged nodes', () => {
     const visited: string[] = [];
     const visitor: Visitor = {
@@ -276,6 +303,20 @@ describe('visit — typed embedded fields (untagged)', () => {
     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, {

packages/traverse/src/traverse.ts

@@ -61,12 +61,34 @@ export function walk(
     const keys = Object.keys(root);
     if (keys.length === 1 && /^[A-Z]/.test(keys[0])) {
       walkNode(keys[0], root[keys[0]], actualCallback, parent, keyPath);
-    } else {
-      for (const key of keys) {
-        walk(root[key], actualCallback, parent, [...keyPath, key]);
+      return;
+    }
+    if (parent === null && keyPath.length === 0) {
+      const rootTag = detectUntaggedRootTag(root);
+      if (rootTag) {
+        walkNode(rootTag, root, actualCallback, parent, keyPath);
+        return;
       }
     }
+    for (const key of keys) {
+      walk(root[key], actualCallback, parent, [...keyPath, key]);
+    }
+  }
+}
+
+/**
+ * libpg-query returns the top-level ParseResult/ScanResult as a bare object
+ * (no `{ParseResult: {...}}` wrapper). Detect those shapes at the root so
+ * their visitors — and typed descendants like RawStmt — are dispatched.
+ */
+function detectUntaggedRootTag(root: any): string | null {
+  if (Array.isArray(root.stmts) && typeof root.version === 'number') {
+    return 'ParseResult';
   }
+  if (Array.isArray(root.tokens) && typeof root.version === 'number') {
+    return 'ScanResult';
+  }
+  return null;
 }
 
 function isTaggedNode(value: any): boolean {
@@ -84,6 +106,10 @@ function walkNode(
   parent: NodePath | null,
   keyPath: readonly (string | number)[],
 ): void {
+  if (typeof nodeData !== 'object' || nodeData === null) {
+    return;
+  }
+
   const path = new NodePath(tag, nodeData, parent, keyPath);
 
   if (actualCallback(path) === false) {
@@ -159,6 +185,12 @@ export function visit(
 ): 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];