Remove the old parser's AST from StatementRewrite

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 = "6421337", optional = true }
+pg_raw_parse = { git = "https://github.com/pgdogdev/pg_raw_parse.git", rev = "6b424c4", optional = true }
 itertools = "0.15.0"
 
 [target.'cfg(unix)'.dependencies]

pgdog/src/frontend/router/parser/cache/ast.rs

@@ -3,7 +3,7 @@ use itertools::*;
 use once_cell::sync::OnceCell;
 use pg_query::{NodeEnum, ParseResult, parse, parse_raw};
 #[cfg(feature = "new_parser")]
-use pg_raw_parse::{Owned, StmtList};
+use pg_raw_parse::{Owned, StmtList, make};
 use pgdog_config::QueryParserEngine;
 use std::fmt::Debug;
 use std::ops::Deref;
@@ -108,19 +108,21 @@ impl Ast {
         search_path: Option<&ParameterValue>,
     ) -> Result<Self, Error> {
         let now = Instant::now();
+        #[cfg(not(feature = "new_parser"))]
         let mut ast = match schema.query_parser_engine {
             QueryParserEngine::PgQueryProtobuf => parse(query.query_without_comment),
             QueryParserEngine::PgQueryRaw => parse_raw(query.query_without_comment),
         }
         .map_err(Error::PgQuery)?;
         #[cfg(feature = "new_parser")]
-        let new_ast = pg_raw_parse::parse(query.query_without_comment).map_err(Error::Parse)?;
+        let ast = pg_raw_parse::parse(query.query_without_comment).map_err(Error::Parse)?;
 
         // Run the rewrite unconditionally. Even when a shard comment will
         // route the query to a specific shard, we need to know whether the
         // same query body (without the comment) would require a rewrite, so
         // `Cache::query` can decide whether this entry is safe to cache.
         let mut rewriter = StatementRewrite::new(StatementRewriteContext {
+            #[cfg(not(feature = "new_parser"))]
             stmt: &mut ast.protobuf,
             extended: query.original_query.extended(),
             prepared: query.original_query.prepared(),
@@ -131,11 +133,18 @@ impl Ast {
             search_path,
         });
         #[cfg(feature = "new_parser")]
-        let rewrite_plan = if let Some(node) = new_ast.stmts().next() {
-            rewriter.maybe_rewrite(node)?
-        } else {
-            Default::default()
-        };
+        let mut rewrite_plan = Default::default();
+        #[cfg(feature = "new_parser")]
+        let ast = make::try_owned(|mem| {
+            // FIXME(sage): We should have a parse function on mem so we don't
+            // need to parse and then copy the parsed tree just to throw the
+            // original away
+            let mut copy = mem.make_unique(&*ast.into_inner());
+            if let Some(stmt) = copy.as_mut().into_iter().next() {
+                rewrite_plan = rewriter.maybe_rewrite(stmt, mem)?;
+            }
+            Ok::<_, Error>(copy)
+        })?;
 
         #[cfg(not(feature = "new_parser"))]
         let rewrite_plan = rewriter.maybe_rewrite()?;
@@ -154,6 +163,10 @@ impl Ast {
             );
         }
 
+        #[cfg(feature = "new_parser")]
+        let old_ast =
+            pg_query::parse(&pg_raw_parse::deparse_stmts(&*ast)?).map_err(Error::PgQuery)?;
+
         Ok(Self {
             cached: true,
             comment_shard: None,
@@ -162,9 +175,10 @@ impl Ast {
             inner: Arc::new(AstInner {
                 stats: Mutex::new(stats),
                 #[cfg(feature = "new_parser")]
-                new_ast: pg_raw_parse::parse(&ast.deparse().unwrap())
-                    .unwrap()
-                    .into_inner(),
+                new_ast: ast,
+                #[cfg(feature = "new_parser")]
+                ast: old_ast,
+                #[cfg(not(feature = "new_parser"))]
                 ast,
                 rewrite_plan,
                 fingerprint: OnceCell::new(),

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

@@ -557,32 +557,8 @@ mod tests {
         db_schema: &Schema,
         mode: RewriteMode,
     ) -> Result<(String, RewritePlan), Error> {
-        let _guard = set_env_var("NODE_ID", "pgdog-1");
-        let mut ast = pg_query::parse(sql).unwrap().protobuf;
-        #[cfg(feature = "new_parser")]
-        let stmt = pg_raw_parse::parse(sql).unwrap();
-        let mut prepared = PreparedStatements::default();
         let schema = sharding_schema_with_mode(mode);
-        let mut rewriter = StatementRewrite::new(StatementRewriteContext {
-            stmt: &mut ast,
-            extended: false,
-            prepared: false,
-            prepared_statements: &mut prepared,
-            schema: &schema,
-            db_schema,
-            user: "",
-            search_path: None,
-        });
-        let plan = rewriter.maybe_rewrite(
-            #[cfg(feature = "new_parser")]
-            stmt.stmts().next().unwrap(),
-        )?;
-        let result = if plan.stmt.is_some() {
-            plan.stmt.clone().unwrap()
-        } else {
-            ast.deparse().unwrap()
-        };
-        Ok((result, plan))
+        rewrite_sql_with_sharding_schema(sql, db_schema, &schema)
     }
 
     #[test]
@@ -768,18 +744,16 @@ mod tests {
         }
     }
 
+    #[cfg(feature = "new_parser")]
     fn rewrite_sql_with_sharding_schema(
         sql: &str,
         db_schema: &Schema,
         schema: &ShardingSchema,
     ) -> Result<(String, RewritePlan), Error> {
         let _guard = set_env_var("NODE_ID", "pgdog-1");
-        let mut ast = pg_query::parse(sql).unwrap().protobuf;
-        #[cfg(feature = "new_parser")]
-        let new_ast = pg_raw_parse::parse(sql).unwrap();
+        let ast = pg_raw_parse::parse(sql).unwrap();
         let mut prepared = PreparedStatements::default();
         let mut rewriter = StatementRewrite::new(StatementRewriteContext {
-            stmt: &mut ast,
             extended: false,
             prepared: false,
             prepared_statements: &mut prepared,
@@ -788,16 +762,46 @@ mod tests {
             user: "",
             search_path: None,
         });
-        let plan = rewriter.maybe_rewrite(
-            #[cfg(feature = "new_parser")]
-            new_ast.stmts().next().unwrap(),
-        )?;
-        let result = if plan.stmt.is_some() {
-            plan.stmt.clone().unwrap()
-        } else {
-            ast.deparse().unwrap()
-        };
-        Ok((result, plan))
+        let mut plan = Default::default();
+        let ast = make::try_owned(|mem| {
+            let mut copy = mem.make_unique(&*ast.into_inner());
+            plan = rewriter.maybe_rewrite(copy.as_mut().into_iter().next().unwrap(), mem)?;
+            Ok::<_, Error>(copy)
+        })?;
+        let sql = pg_raw_parse::deparse_stmts(&*ast)?;
+        Ok((sql, plan))
+    }
+
+    cfg_select! {
+        not(feature = "new_parser") => {
+            fn rewrite_sql_with_sharding_schema(
+                sql: &str,
+                db_schema: &Schema,
+                schema: &ShardingSchema,
+            ) -> Result<(String, RewritePlan), Error> {
+                let _guard = set_env_var("NODE_ID", "pgdog-1");
+                let mut ast = pg_query::parse(sql).unwrap().protobuf;
+                let mut prepared = PreparedStatements::default();
+                let mut rewriter = StatementRewrite::new(StatementRewriteContext {
+                    stmt: &mut ast,
+                    extended: false,
+                    prepared: false,
+                    prepared_statements: &mut prepared,
+                    schema,
+                    db_schema,
+                    user: "",
+                    search_path: None,
+                });
+                let plan = rewriter.maybe_rewrite()?;
+                let result = if plan.stmt.is_some() {
+                    plan.stmt.clone().unwrap()
+                } else {
+                    ast.deparse().unwrap()
+                };
+                Ok((result, plan))
+            }
+        }
+        _ => {}
     }
 
     #[test]

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

@@ -464,6 +464,7 @@ mod tests {
     }
 
     fn parse_and_split(sql: &str) -> Vec<InsertSplit> {
+        #[cfg(not(feature = "new_parser"))]
         let mut ast = pg_query::parse(sql).unwrap().protobuf;
         #[cfg(feature = "new_parser")]
         let root = pg_raw_parse::parse(sql).unwrap();
@@ -476,6 +477,7 @@ mod tests {
         let schema = default_schema();
         let db_schema = default_db_schema();
         let mut rewriter = StatementRewrite::new(StatementRewriteContext {
+            #[cfg(not(feature = "new_parser"))]
             stmt: &mut ast,
             extended: false,
             prepared: false,

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

@@ -2,9 +2,10 @@
 
 #[cfg(not(feature = "new_parser"))]
 use pg_query::Node as PgNode;
+#[cfg(not(feature = "new_parser"))]
 use pg_query::protobuf::ParseResult;
 #[cfg(feature = "new_parser")]
-use pg_raw_parse::{Node, NodeMut, make, transform, walk};
+use pg_raw_parse::{Node, NodeMut, make, nodes, transform, walk};
 #[cfg(not(feature = "new_parser"))]
 use pgdog_config::QueryParserEngine;
 
@@ -35,6 +36,7 @@ pub(crate) use update::*;
 #[derive(Debug)]
 pub struct StatementRewriteContext<'a> {
     /// The AST of the statement we are rewriting.
+    #[cfg(not(feature = "new_parser"))]
     pub stmt: &'a mut ParseResult,
     /// The statement is using the extended protocol with placeholders.
     pub extended: bool,
@@ -56,6 +58,7 @@ pub struct StatementRewriteContext<'a> {
 #[derive(Debug)]
 pub struct StatementRewrite<'a> {
     /// SQL statement.
+    #[cfg(not(feature = "new_parser"))]
     stmt: &'a mut ParseResult,
     /// The statement was rewritten.
     rewritten: bool,
@@ -85,6 +88,7 @@ impl<'a> StatementRewrite<'a> {
     ///
     pub fn new(ctx: StatementRewriteContext<'a>) -> Self {
         Self {
+            #[cfg(not(feature = "new_parser"))]
             stmt: ctx.stmt,
             rewritten: false,
             extended: ctx.extended,
@@ -110,14 +114,18 @@ impl<'a> StatementRewrite<'a> {
     /// Maybe rewrite the statement and produce a rewrite plan
     /// we can apply to Bind messages.
     #[cfg(feature = "new_parser")]
-    pub fn maybe_rewrite(&mut self, node: Node<'a>) -> Result<RewritePlan, Error> {
+    pub fn maybe_rewrite<'mem>(
+        &mut self,
+        mut stmt: nodes::RawStmtMut<'mem, '_>,
+        mem: make::MemoryToken<'mem>,
+    ) -> Result<RewritePlan, Error> {
         let mut plan = RewritePlan::default();
 
-        match node {
+        match stmt.stmt() {
             Node::InsertStmt(_)
             | Node::SelectStmt(_)
             | Node::UpdateStmt(_)
-            | Node::DeleteStmt(_) => walk::walk(node, |node| match node {
+            | Node::DeleteStmt(_) => walk::walk(stmt.stmt(), |node| match node {
                 Node::ParamRef(param) => plan.params = plan.params.max(param.number as u16),
                 _ => {}
             }),
@@ -126,27 +134,26 @@ impl<'a> StatementRewrite<'a> {
             _ => return Ok(plan),
         }
 
-        make::try_owned(|mem| {
-            let mut node = mem.make_unique(node);
-
-            // Handle top-level PREPARE/EXECUTE statements.
-            let prepared_result = self.rewrite_simple_prepared(node.as_mut(), mem)?;
-            if prepared_result.rewritten {
-                self.rewritten = true;
-                plan.prepares = prepared_result.prepares;
-            }
+        // Handle top-level PREPARE/EXECUTE statements.
+        let prepared_result = self.rewrite_simple_prepared(stmt.stmt_mut(), mem)?;
+        if prepared_result.rewritten {
+            self.rewritten = true;
+            plan.prepares = prepared_result.prepares;
+        }
 
-            // Inject pgdog.unique_id() for missing BIGINT primary keys.
-            // This must run BEFORE the unique_id rewriter so the injected
-            // function calls get processed.
-            if let NodeMut::InsertStmt(insert) = node.as_mut() {
-                self.inject_auto_id(insert, mem, &mut plan)?;
-            }
+        // Inject pgdog.unique_id() for missing BIGINT primary keys.
+        // This must run BEFORE the unique_id rewriter so the injected
+        // function calls get processed.
+        if let NodeMut::InsertStmt(insert) = stmt.stmt_mut() {
+            self.inject_auto_id(insert, mem, &mut plan)?;
+        }
 
-            // Track the next parameter number to use
-            let mut next_param = plan.params as i32 + 1;
-            let generator = crate::unique_id::UniqueId::generator()?;
-            transform::transform(&mut node, |node| {
+        // Track the next parameter number to use
+        let mut next_param = plan.params as i32 + 1;
+        let generator = crate::unique_id::UniqueId::generator()?;
+        transform::transform_node(
+            stmt.stmt_mut(),
+            &mut transform::TransformClosure::new(|node| {
                 if let Some(replacement) = Self::rewrite_unique_id(
                     node.as_ref(),
                     mem,
@@ -161,33 +168,25 @@ impl<'a> StatementRewrite<'a> {
                 } else {
                     Some(node)
                 }
-            });
-
-            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_ref())?.as_str().to_owned());
-            }
+            }),
+        );
 
-            if let Node::InsertStmt(insert) = node.as_ref() {
-                self.split_insert(insert, &mut plan)?;
-            }
+        if let NodeMut::SelectStmt(mut select) = stmt.stmt_mut() {
+            self.rewrite_aggregates(&mut select, mem, &mut plan, self.db_schema)?;
+            self.limit_offset(&select, &mut plan);
+        }
 
-            if let Node::UpdateStmt(stmt) = node.as_ref() {
-                self.sharding_key_update(stmt, &mut plan)?;
-            }
+        if self.rewritten {
+            plan.stmt = Some(pg_raw_parse::deparse(&*stmt)?.as_str().to_owned());
+        }
 
-            self.stmt.stmts[0] = pg_query::parse(pg_raw_parse::deparse(node.as_ref())?.as_str())?
-                .protobuf
-                .stmts
-                .pop()
-                .unwrap();
+        if let Node::InsertStmt(insert) = stmt.stmt() {
+            self.split_insert(insert, &mut plan)?;
+        }
 
-            Ok::<_, Error>(mem.make_raw_stmt(node))
-        })?;
+        if let Node::UpdateStmt(stmt) = stmt.stmt() {
+            self.sharding_key_update(stmt, &mut plan)?;
+        }
 
         Ok(plan)
     }

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

@@ -322,13 +322,15 @@ mod tests {
     }
 
     fn run_limit_offset(sql: &str, schema: &ShardingSchema) -> RewritePlan {
+        #[cfg(not(feature = "new_parser"))]
         let mut ast = pg_query::parse(sql).unwrap();
         #[cfg(feature = "new_parser")]
         let stmt = pg_raw_parse::parse(sql).unwrap();
         let db_schema = Schema::default();
         let mut ps = PreparedStatements::default();
         #[cfg_attr(feature = "new_parser", allow(unused_mut))]
         let mut rewrite = StatementRewrite::new(StatementRewriteContext {
+            #[cfg(not(feature = "new_parser"))]
             stmt: &mut ast.protobuf,
             extended: false,
             prepared: false,