Port StatementRewrite::sharding_key_update to the new parser

pgdog/Cargo.toml

@@ -80,7 +80,7 @@ smallvec = "1"
 reqwest = { version = "0.12", default-features = false, features = ["rustls-tls-webpki-roots-no-provider", "json"] }
 hex = "0.4"
 x509-parser = "0.18"
-pg_raw_parse = { git = "https://github.com/pgdogdev/pg_raw_parse.git", rev = "6f5d35c", optional = true }
+pg_raw_parse = { git = "https://github.com/pgdogdev/pg_raw_parse.git", rev = "34e5987", optional = true }
 itertools = "0.15.0"
 
 [target.'cfg(unix)'.dependencies]

pgdog/src/frontend/client/query_engine/multi_step/update.rs

@@ -99,7 +99,7 @@ impl<'a> UpdateMulti<'a> {
         context: &mut QueryEngineContext<'_>,
         row: Row,
     ) -> Result<(), Error> {
-        let mut request = self.rewrite.insert.build_request(
+        let mut request = self.rewrite.build_insert_request(
             context.client_request,
             &row.row_description,
             &row.data_row,
@@ -140,12 +140,8 @@ impl<'a> UpdateMulti<'a> {
             }
 
             self.delete_row(context).await?;
-            self.execute_request_internal(
-                context,
-                &mut request,
-                self.rewrite.insert.is_returning(),
-            )
-            .await?;
+            self.execute_request_internal(context, &mut request, self.rewrite.is_returning())
+                .await?;
 
             self.engine
                 .process_server_message(context, CommandComplete::new("UPDATE 1").message()?) // We only allow to update one row at a time.
@@ -167,9 +163,7 @@ impl<'a> UpdateMulti<'a> {
     ) -> Result<bool, Error> {
         let cluster = self.engine.backend.cluster()?;
         let schema = cluster.schema();
-        let Some(table) = self.rewrite.target_table() else {
-            return Ok(false);
-        };
+        let table = self.rewrite.target_table();
 
         let Some(relation) = schema.table(table, cluster.user(), context.params.search_path())
         else {

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

@@ -1,5 +1,9 @@
+#[cfg(feature = "new_parser")]
+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 pgdog_config::QueryParserEngine;
 use std::fmt::Debug;
 use std::ops::Deref;
@@ -40,7 +44,7 @@ pub struct AstInner {
     pub ast: ParseResult,
     // FIXME(sage): Rename to ast when parser port is done
     #[cfg(feature = "new_parser")]
-    pub new_ast: pg_raw_parse::ParseResult,
+    pub new_ast: Owned<StmtList>,
     /// AST stats.
     pub stats: Mutex<Stats>,
     /// Rewrite plan.
@@ -53,10 +57,29 @@ pub struct AstInner {
 
 impl AstInner {
     /// Create new AST record, with no rewrite or comment routing.
-    pub fn new(ast: ParseResult) -> Self {
+    #[cfg(feature = "new_parser")]
+    pub(crate) fn new(ast: Owned<StmtList>) -> Self {
+        let deparse_results = ast
+            .iter()
+            .map(|s| pg_raw_parse::deparse(s))
+            .collect::<Result<Vec<_>, _>>()
+            .unwrap();
+        Self {
+            new_ast: ast,
+            ast: pg_query::parse(&deparse_results.iter().map(|r| r.as_str()).join("; ")).unwrap(),
+            stats: Mutex::new(Stats::new()),
+            rewrite_plan: RewritePlan::default(),
+            fingerprint: OnceCell::new(),
+            query_without_comment: "".into(),
+        }
+    }
+
+    pub(crate) fn old(ast: ParseResult) -> Self {
         Self {
             #[cfg(feature = "new_parser")]
-            new_ast: pg_raw_parse::parse(&ast.deparse().unwrap()).unwrap(),
+            new_ast: pg_raw_parse::parse(&ast.deparse().unwrap())
+                .unwrap()
+                .into_inner(),
             ast,
             stats: Mutex::new(Stats::new()),
             rewrite_plan: RewritePlan::default(),
@@ -133,7 +156,9 @@ impl Ast {
             inner: Arc::new(AstInner {
                 stats: Mutex::new(stats),
                 #[cfg(feature = "new_parser")]
-                new_ast: pg_raw_parse::parse(&ast.deparse().unwrap()).unwrap(),
+                new_ast: pg_raw_parse::parse(&ast.deparse().unwrap())
+                    .unwrap()
+                    .into_inner(),
                 ast,
                 rewrite_plan,
                 fingerprint: OnceCell::new(),
@@ -171,18 +196,30 @@ impl Ast {
             comment_role: None,
             comment_shard: None,
             query_parser_engine,
-            inner: Arc::new(AstInner::new(ast)),
+            inner: Arc::new(AstInner::old(ast)),
         })
     }
 
+    /// Create new AST from a parse result.
+    #[cfg(feature = "new_parser")]
+    pub fn from_raw_stmts(stmts: Owned<StmtList>) -> Self {
+        Self {
+            cached: true,
+            comment_role: None,
+            comment_shard: None,
+            query_parser_engine: QueryParserEngine::default(),
+            inner: Arc::new(AstInner::new(stmts)),
+        }
+    }
+
     /// Create new AST from a parse result.
     pub fn from_parse_result(parse_result: ParseResult) -> Self {
         Self {
             cached: true,
             comment_role: None,
             comment_shard: None,
             query_parser_engine: QueryParserEngine::default(),
-            inner: Arc::new(AstInner::new(parse_result)),
+            inner: Arc::new(AstInner::old(parse_result)),
         }
     }