From dce8bba2d8a8401943d85044d2385e8e70be336c Mon Sep 17 00:00:00 2001
From: Sage Griffin <sage@sagetheprogrammer.com>
Date: Thu, 25 Jun 2026 16:13:34 -0600
Subject: [PATCH] Update `StatementParser::run_walk` to not recurse twice

I changed the behavior of the parser to not pass the root node to the
callback. I found recursing based on pointer equality to be too painful
as I work on porting the routing functions.

As a side effect, this should be more efficient. We did a full AST walk
of each subselect looking for advisory locks, and then walked it again
looking for tables. With this change they're now a single walk, and the
code should be easier to follow
---
 Cargo.lock                                    |  2 +-
 pgdog/Cargo.toml                              |  2 +-
 pgdog/src/frontend/router/parser/statement.rs | 65 ++++++++++---------
 3 files changed, 37 insertions(+), 32 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 4bc6db03e..a5d28f239 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -3224,7 +3224,7 @@ dependencies = [
 [[package]]
 name = "pg_raw_parse"
 version = "0.1.0"
-source = "git+https://github.com/pgdogdev/pg_raw_parse.git?rev=a5b12f3#a5b12f330f10ecf309c5621624fc8ee2063e63aa"
+source = "git+https://github.com/pgdogdev/pg_raw_parse.git?rev=adbfcae#adbfcae6fd22441b9e32e97be088e2a2c893cf72"
 dependencies = [
  "bindgen 0.72.1",
  "cc",
diff --git a/pgdog/Cargo.toml b/pgdog/Cargo.toml
index 3916e2993..c4a4319cd 100644
--- a/pgdog/Cargo.toml
+++ b/pgdog/Cargo.toml
@@ -79,7 +79,7 @@ smallvec = "1"
 reqwest = { version = "0.12", default-features = false, features = ["rustls-tls-webpki-roots-no-provider"] }
 hex = "0.4"
 x509-parser = "0.18"
-pg_raw_parse = { git = "https://github.com/pgdogdev/pg_raw_parse.git", rev = "a5b12f3", optional = true }
+pg_raw_parse = { git = "https://github.com/pgdogdev/pg_raw_parse.git", rev = "adbfcae", optional = true }
 itertools = "0.15.0"
 
 [target.'cfg(unix)'.dependencies]
diff --git a/pgdog/src/frontend/router/parser/statement.rs b/pgdog/src/frontend/router/parser/statement.rs
index 3447da973..c364849eb 100644
--- a/pgdog/src/frontend/router/parser/statement.rs
+++ b/pgdog/src/frontend/router/parser/statement.rs
@@ -18,6 +18,8 @@ use pg_query::{
     },
 };
 #[cfg(feature = "new_parser")]
+use pg_raw_parse::walk::Recurse;
+#[cfg(feature = "new_parser")]
 use pg_raw_parse::{Node, nodes, walk};
 
 #[cfg(feature = "new_parser")]
@@ -915,41 +917,44 @@ impl<'a, 'b, 'c> StatementParser<'a, 'b, 'c> {
     // Are we running? Or walking? MAKE UP YOUR MIND DAMMIT
     #[cfg(feature = "new_parser")]
     fn run_walk(&self) -> Walk<'a> {
-        use pg_raw_parse::walk::Recurse;
-        use std::ops::ControlFlow;
-        use std::ptr;
-
         let mut walk = Walk::default();
-        walk::walk(self.new_stmt, |node| match node {
-            Node::SelectStmt(s) => {
-                let values_columns = collect_values_columns(s);
-
-                // Extract any advisory locks that exist in this, and only
-                // this select statement, using the values columns that appear
-                // in its FROM clause if any
-                walk::walk_manual::<()>(node, |node| match node {
-                    Node::FuncCall(func) => {
-                        walk.advisory_locks.extend(advisory_locks_from_func_call(
-                            func,
-                            self.bind,
-                            values_columns.as_ref(),
-                        ));
-                        ControlFlow::Continue(Recurse::No)
-                    }
-                    // Don't recurse into subselects, otherwise we'll resolve
-                    // advisory locks with the wrong set of values columns
-                    Node::SelectStmt(s2) => Recurse::recurse_if(ptr::eq(s, s2)),
-                    _ => ControlFlow::Continue(Recurse::Yes),
-                })
-                .expect("PG received an unrecognized node");
+        self.walk_stmt(self.new_stmt, &mut walk);
+        walk
+    }
+
+    #[cfg(feature = "new_parser")]
+    fn walk_stmt(&self, stmt: Node<'a>, walk: &mut Walk<'a>) {
+        let values_columns = match stmt {
+            Node::SelectStmt(s) => collect_values_columns(s),
+            _ => None,
+        };
+        walk::walk_manual::<()>(stmt, |node| match node {
+            // Walk the select statement separately, as it may have its own
+            // set of VALUES columns.
+            Node::SelectStmt(_) => {
+                self.walk_stmt(node, walk);
+                Recurse::no()
+            }
+
+            // Extract any advisory locks that this function call may
+            // represent, using the values clause of the current statement
+            Node::FuncCall(func) => {
+                walk.advisory_locks.extend(advisory_locks_from_func_call(
+                    func,
+                    self.bind,
+                    values_columns.as_ref(),
+                ));
+                Recurse::no()
             }
 
-            Node::RangeVar(r) => walk.tables.push(Table::from(r)),
-            _ => (),
+            Node::RangeVar(r) => {
+                walk.tables.push(Table::from(r));
+                Recurse::yes()
+            }
+
+            _ => Recurse::yes(),
         })
         .expect("PG received an unrecognized node");
-
-        walk
     }
 
     #[cfg(not(feature = "new_parser"))]