Port StatementRewrite::rewrite_unique_id

pgdog/Cargo.toml

@@ -80,7 +80,7 @@ smallvec = "1"
 reqwest.workspace = true
 hex = "0.4"
 x509-parser = "0.18"
-pg_raw_parse = { git = "https://github.com/pgdogdev/pg_raw_parse.git", rev = "a4a2a2c", optional = true }
+pg_raw_parse = { git = "https://github.com/pgdogdev/pg_raw_parse.git", rev = "b2a15ec", optional = true }
 itertools = "0.15.0"
 
 [target.'cfg(unix)'.dependencies]

pgdog/src/frontend/router/parser/rewrite/statement/aggregate/engine.rs

@@ -26,7 +26,7 @@ impl AggregatesRewrite {
     /// Rewrite a SELECT query in-place, adding helper aggregates when necessary.
     #[cfg(feature = "new_parser")]
     pub(crate) fn rewrite_select<'a>(
-        mut select: nodes::SelectStmtMut<'a, '_>,
+        select: &mut nodes::SelectStmtMut<'a, '_>,
         mem: make::MemoryToken<'a>,
         aggregate: &Aggregate,
     ) -> RewriteOutput {
@@ -455,7 +455,7 @@ mod tests {
 
             let aggregate = Aggregate::parse(&*stmt, &Default::default());
             output = Some(AggregatesRewrite::rewrite_select(
-                stmt.as_mut(),
+                &mut stmt.as_mut(),
                 mem,
                 &aggregate,
             ));

pgdog/src/frontend/router/parser/rewrite/statement/aggregate/mod.rs

@@ -17,7 +17,7 @@ impl StatementRewrite<'_> {
     #[cfg(feature = "new_parser")]
     pub(super) fn rewrite_aggregates<'a>(
         &mut self,
-        select: SelectStmtMut<'a, '_>,
+        select: &mut SelectStmtMut<'a, '_>,
         mem: MemoryToken<'a>,
         plan: &mut RewritePlan,
         schema: &Schema,

pgdog/src/frontend/router/parser/rewrite/statement/mod.rs

@@ -1,9 +1,10 @@
 //! Statement rewriter.
 
+#[cfg(not(feature = "new_parser"))]
 use pg_query::Node as PgNode;
 use pg_query::protobuf::ParseResult;
 #[cfg(feature = "new_parser")]
-use pg_raw_parse::{Node, NodeMut, make, walk};
+use pg_raw_parse::{Node, NodeMut, make, transform, walk};
 #[cfg(not(feature = "new_parser"))]
 use pgdog_config::QueryParserEngine;
 
@@ -142,47 +143,45 @@ impl<'a> StatementRewrite<'a> {
 
             // Track the next parameter number to use
             let mut next_param = plan.params as i32 + 1;
-
-            self.stmt.stmts[0] = pg_query::parse(pg_raw_parse::deparse(node.as_ref())?.as_str())?
-                .protobuf
-                .stmts
-                .pop()
-                .unwrap();
-            let extended = self.extended;
-            visitor::visit_and_mutate_nodes(self.stmt, |node| -> Result<Option<PgNode>, Error> {
-                match Self::rewrite_unique_id(node, extended, &mut next_param)? {
-                    Some(replacement) => {
-                        plan.unique_ids += 1;
-                        self.rewritten = true;
-                        Ok(Some(replacement))
-                    }
-                    None => Ok(None),
+            let generator = crate::unique_id::UniqueId::generator()?;
+            transform::transform(&mut node, |node| {
+                if let Some(replacement) = Self::rewrite_unique_id(
+                    node.as_ref(),
+                    mem,
+                    generator,
+                    self.extended,
+                    &mut next_param,
+                ) {
+                    plan.unique_ids += 1;
+                    self.rewritten = true;
+                    node.replace(replacement);
+                    None
+                } else {
+                    Some(node)
                 }
-            })?;
-
-            let reparsed = pg_raw_parse::parse(&pg_query::deparse(self.stmt)?)?;
-            let mut node = reparsed.stmts().next().unwrap();
-            let mut select_stmt;
-            if let Node::SelectStmt(select) = node {
-                select_stmt = mem.make_unique(select);
-                self.rewrite_aggregates(select_stmt.as_mut(), mem, &mut plan, self.db_schema)?;
-                self.limit_offset(&select_stmt, &mut plan);
-                node = (&*select_stmt).into();
-                self.stmt.stmts = pg_query::parse(pg_raw_parse::deparse(node)?.as_str())?
-                    .protobuf
-                    .stmts;
+            });
+
+            if let NodeMut::SelectStmt(mut select) = node.as_mut() {
+                self.rewrite_aggregates(&mut select, mem, &mut plan, self.db_schema)?;
+                self.limit_offset(&select, &mut plan);
             }
 
             if self.rewritten {
-                plan.stmt = Some(pg_raw_parse::deparse(node)?.as_str().to_owned());
+                plan.stmt = Some(pg_raw_parse::deparse(node.as_ref())?.as_str().to_owned());
             }
 
             self.split_insert(&mut plan)?;
-            if let Node::UpdateStmt(stmt) = node {
+            if let Node::UpdateStmt(stmt) = node.as_ref() {
                 self.sharding_key_update(stmt, &mut plan)?;
             }
 
-            Ok::<_, Error>(mem.make_raw_stmt(mem.make_unique(node)))
+            self.stmt.stmts[0] = pg_query::parse(pg_raw_parse::deparse(node.as_ref())?.as_str())?
+                .protobuf
+                .stmts
+                .pop()
+                .unwrap();
+
+            Ok::<_, Error>(mem.make_raw_stmt(node))
         })?;
 
         Ok(plan)