diff --git a/pgdog/src/frontend/client/query_engine/context.rs b/pgdog/src/frontend/client/query_engine/context.rs index cd24a1f3..d1cf2c10 100644 --- a/pgdog/src/frontend/client/query_engine/context.rs +++ b/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, /// 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, } diff --git a/pgdog/src/frontend/client/query_engine/end_transaction.rs b/pgdog/src/frontend/client/query_engine/end_transaction.rs index a6929b24..d98f409d 100644 --- a/pgdog/src/frontend/client/query_engine/end_transaction.rs +++ b/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(()) diff --git a/pgdog/src/frontend/client/query_engine/mod.rs b/pgdog/src/frontend/client/query_engine/mod.rs index 484ae538..bc2c1d35 100644 --- a/pgdog/src/frontend/client/query_engine/mod.rs +++ b/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?, } diff --git a/pgdog/src/frontend/client/query_engine/query.rs b/pgdog/src/frontend/client/query_engine/query.rs index e133706a..7158402d 100644 --- a/pgdog/src/frontend/client/query_engine/query.rs +++ b/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, ) -> 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, ) -> 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?; diff --git a/pgdog/src/frontend/client/query_engine/rewrite.rs b/pgdog/src/frontend/client/query_engine/rewrite.rs index 0fe9d3cb..43c2f306 100644 --- a/pgdog/src/frontend/client/query_engine/rewrite.rs +++ b/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,7 +23,7 @@ impl QueryEngine { pub(super) fn parse_and_rewrite( &mut self, context: &mut QueryEngineContext<'_>, - ) -> Result { + ) -> Result, Error> { let use_parser = self .backend .cluster() @@ -31,7 +31,7 @@ impl QueryEngine { .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) } } diff --git a/pgdog/src/frontend/client/query_engine/route_query.rs b/pgdog/src/frontend/client/query_engine/route_query.rs index 2e54a9a0..7df79afa 100644 --- a/pgdog/src/frontend/client/query_engine/route_query.rs +++ b/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 { // 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)?; } diff --git a/pgdog/src/frontend/client/query_engine/set.rs b/pgdog/src/frontend/client/query_engine/set.rs index dff7cfde..0b4f2be7 100644 --- a/pgdog/src/frontend/client/query_engine/set.rs +++ b/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::>) .await?; diff --git a/pgdog/src/frontend/client/query_engine/start_transaction.rs b/pgdog/src/frontend/client/query_engine/start_transaction.rs index 675611a7..74d3b703 100644 --- a/pgdog/src/frontend/client/query_engine/start_transaction.rs +++ b/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)