refactor(query_engine_context): remove rw result from context field

pgdog/src/frontend/client/query_engine/context.rs

@@ -3,7 +3,6 @@ use crate::{
     frontend::{
         Client, ClientRequest, PreparedStatements,
         client::{Sticky, TransactionType, timeouts::Timeouts},
-        router::parser::rewrite::statement::plan::RewriteResult,
     },
     net::{FrontendPid, Parameters, Stream},
 };
@@ -36,8 +35,6 @@ pub struct QueryEngineContext<'a> {
     pub(super) rollback: bool,
     /// Sticky config:
     pub(super) sticky: Sticky,
-    /// Rewrite result.
-    pub(super) rewrite_result: Option<RewriteResult>,
     /// Log queries to stdout.
     pub(super) query_log_stdout: bool,
     /// Maximum query message size before a warning is logged.
@@ -62,7 +59,6 @@ impl<'a> QueryEngineContext<'a> {
             requests_left: 0,
             rollback: false,
             sticky: client.sticky,
-            rewrite_result: None,
             query_log_stdout: client.query_log_stdout,
             query_size_limit: client.query_size_limit,
         }
@@ -90,7 +86,6 @@ impl<'a> QueryEngineContext<'a> {
             requests_left: 0,
             rollback: false,
             sticky: Sticky::new(),
-            rewrite_result: None,
             query_log_stdout: false,
             query_size_limit: None,
         }

pgdog/src/frontend/client/query_engine/end_transaction.rs

@@ -91,7 +91,7 @@ impl QueryEngine {
                 self.notify_buffer.clear();
             }
             context.rollback = rollback;
-            self.execute(context).await?;
+            self.execute(context, None).await?;
         }
 
         Ok(())

pgdog/src/frontend/client/query_engine/mod.rs

@@ -40,6 +40,7 @@ pub mod two_pc;
 
 use self::query::ExplainResponseState;
 use self::query_log_stdout::log_query_stdout;
+use crate::frontend::router::parser::rewrite::statement::plan::RewriteResult;
 pub(crate) use advisory_lock::AdvisoryLocks;
 pub use context::QueryEngineContext;
 use notify_buffer::NotifyBuffer;
@@ -126,15 +127,14 @@ impl QueryEngine {
         }
 
         // Rewrite statement if necessary.
-        match self.parse_and_rewrite(context) {
-            Ok(true) => {}
-            Ok(false) => return Ok(()),
+        let rewrite_result = match self.parse_and_rewrite(context) {
+            Ok(rewrite_result) => rewrite_result,
             Err(e) => {
                 self.error_response(context, ErrorResponse::syntax(e.to_string()))
                     .await?;
                 return Ok(());
             }
-        }
+        };
 
         // Intercept commands we don't have to forward to a server.
         if self.intercept_incomplete(context).await? {
@@ -143,7 +143,7 @@ impl QueryEngine {
         }
 
         // Route transaction to the right servers.
-        if !self.route_query(context).await? {
+        if !self.route_query(context, rewrite_result.as_ref()).await? {
             self.update_stats(context);
             debug!("query has nowhere to go");
             return Ok(());
@@ -221,11 +221,11 @@ impl QueryEngine {
 
                 context.params.rollback();
             }
-            Command::Query(_) => self.execute(context).await?,
+            Command::Query(_) => self.execute(context, rewrite_result).await?,
             Command::Listen { .. } | Command::Notify { .. } | Command::Unlisten(_)
                 if self.backend.session_mode() =>
             {
-                self.execute(context).await?
+                self.execute(context, rewrite_result).await?
             }
             Command::Listen { channel, shard } => {
                 self.listen(context, &channel.clone(), shard.clone())
@@ -251,7 +251,7 @@ impl QueryEngine {
             Command::ResetAll => {
                 self.reset_all(context).await?;
             }
-            Command::Copy(_) => self.execute(context).await?,
+            Command::Copy(_) => self.execute(context, rewrite_result).await?,
             Command::Deallocate => self.deallocate(context).await?,
             Command::Discard { extended } => self.discard(context, *extended).await?,
         }

pgdog/src/frontend/client/query_engine/query.rs

@@ -23,6 +23,7 @@ impl QueryEngine {
     pub(super) async fn execute(
         &mut self,
         context: &mut QueryEngineContext<'_>,
+        query_planner: Option<RewriteResult>,
     ) -> Result<(), Error> {
         // Check that we're not in a transaction error state.
         if !self.transaction_error_check(context).await? {
@@ -59,7 +60,11 @@ impl QueryEngine {
         }
 
         let query_timeout = context.timeouts.query_timeout(&State::Active);
-        let result = safe_timeout(query_timeout, self.client_server_exchange(context)).await;
+        let result = safe_timeout(
+            query_timeout,
+            self.client_server_exchange(context, query_planner),
+        )
+        .await;
 
         match result {
             Ok(response) => response?,
@@ -77,8 +82,9 @@ impl QueryEngine {
     async fn client_server_exchange(
         &mut self,
         context: &mut QueryEngineContext<'_>,
+        rewrite_result: Option<RewriteResult>,
     ) -> Result<(), Error> {
-        match context.rewrite_result.take() {
+        match rewrite_result {
             Some(RewriteResult::InsertSplit(requests)) => {
                 Box::pin(multi_step::InsertMulti::from_engine(self, requests).execute(context))
                     .await?;

pgdog/src/frontend/client/query_engine/rewrite.rs

@@ -1,6 +1,6 @@
-use crate::frontend::router::parser::{AstContext, Cache};
-
 use super::*;
+use crate::frontend::router::parser::rewrite::statement::plan::RewriteResult;
+use crate::frontend::router::parser::{AstContext, Cache};
 
 impl QueryEngine {
     /// Rewrite extended protocol messages.
@@ -23,15 +23,15 @@ impl QueryEngine {
     pub(super) fn parse_and_rewrite(
         &mut self,
         context: &mut QueryEngineContext<'_>,
-    ) -> Result<bool, Error> {
+    ) -> Result<Option<RewriteResult>, Error> {
         let use_parser = self
             .backend
             .cluster()
             .map(|cluster| cluster.use_query_parser(context.client_request))
             .unwrap_or(false);
 
         if !use_parser {
-            return Ok(true);
+            return Ok(None);
         }
 
         let query = context.client_request.query()?;
@@ -40,10 +40,12 @@ impl QueryEngine {
             let ast_ctx = AstContext::from_cluster(cluster, context.params);
             let ast = Cache::get().query(&query, &ast_ctx, context.prepared_statements)?;
 
-            context.rewrite_result = Some(ast.rewrite_plan.apply(context.client_request)?);
+            let rewrite_result = ast.rewrite_plan.apply(context.client_request)?;
+            //context.rewrite_result = Some(ast.rewrite_plan.apply(context.client_request)?);
             context.client_request.ast = Some(ast);
+            Ok(Some(rewrite_result))
+        } else {
+            Ok(None)
         }
-
-        Ok(true)
     }
 }

pgdog/src/frontend/client/query_engine/route_query.rs

@@ -70,6 +70,7 @@ impl QueryEngine {
     pub(super) async fn route_query(
         &mut self,
         context: &mut QueryEngineContext<'_>,
+        rewrite_result: Option<&RewriteResult>,
     ) -> Result<bool, Error> {
         // Check that we can route this transaction at all.
         if self.backend.pooler_mode() == PoolerMode::Statement && context.client_request.is_begin()
@@ -151,7 +152,7 @@ impl QueryEngine {
                 );
 
                 // Apply post-parser rewrites, e.g. offset/limit.
-                if let Some(rewrite_result) = &context.rewrite_result {
+                if let Some(rewrite_result) = rewrite_result {
                     rewrite_result.apply_after_parser(context.client_request)?;
                 }
 

pgdog/src/frontend/client/query_engine/set.rs

@@ -56,7 +56,7 @@ impl QueryEngine {
         }
 
         if self.backend.connected() {
-            self.execute(context).await?;
+            self.execute(context, None).await?;
         } else {
             let values_to_return =
                 behave_like_select.then(|| params.iter().map(|p| p.value.as_ref()));
@@ -99,7 +99,7 @@ impl QueryEngine {
         context.params.reset_all();
 
         if self.backend.connected() {
-            self.execute(context).await?;
+            self.execute(context, None).await?;
         } else {
             self.fake_command_response(context, "RESET", None::<Option<_>>)
                 .await?;

pgdog/src/frontend/client/query_engine/start_transaction.rs

@@ -20,7 +20,7 @@ impl QueryEngine {
         context.transaction = Some(transaction_type);
 
         if self.backend.connected() {
-            self.execute(context).await?;
+            self.execute(context, None).await?;
         } else {
             let bytes_sent = if extended {
                 self.extended_transaction_reply(context, true, false)