From 8aeadf75db66cc7aa1d57dc06d4e5bc910599d32 Mon Sep 17 00:00:00 2001 From: jkaczman Date: Mon, 24 Aug 2026 13:32:06 -0400 Subject: [PATCH 1/2] refactor(query_engine_context): remove rw result from context field Starting to restructure RewriteResult to become a "Query Planner"; As part of our goal to support more "multi-query" queries such as subselects and CTEs, as well as supporting better EXPLAIN output, we need an intermediate step between parsing and execution that creates a proper "Query Plan". After going through the code and poking it a few different ways, we determined that the best place to live is what is currently called RewriteResult, and what we expect to eventually call `QueryPlanner`. Our long term goal is to have this take the query engine context as an argument, and return the abstract steps required, those steps being the query to be run (with subqueries replaced with placeholders), and the route it is going to be run on. As we're promoting this to a "thing that does the thing", it no longer makes conceptual sense to have it live on QueryEngineContext, as it's not context. Additionally, we intend to eventually have this start to hold onto borrowed data, and adding an additional lifetime to QueryEngineContext will affect much more code than what is necessary for this refactor. --- .../frontend/client/query_engine/context.rs | 5 ----- .../client/query_engine/end_transaction.rs | 2 +- pgdog/src/frontend/client/query_engine/mod.rs | 16 ++++++++-------- .../src/frontend/client/query_engine/query.rs | 6 ++++-- .../frontend/client/query_engine/rewrite.rs | 19 +++++++++++-------- .../client/query_engine/route_query.rs | 3 ++- pgdog/src/frontend/client/query_engine/set.rs | 4 ++-- .../client/query_engine/start_transaction.rs | 2 +- 8 files changed, 29 insertions(+), 28 deletions(-) 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..f1cc761e 100644 --- a/pgdog/src/frontend/client/query_engine/mod.rs +++ b/pgdog/src/frontend/client/query_engine/mod.rs @@ -45,6 +45,7 @@ pub use context::QueryEngineContext; use notify_buffer::NotifyBuffer; use two_pc::TwoPc; pub use two_pc::phase::TwoPcPhase; +use crate::frontend::router::parser::rewrite::statement::plan::RewriteResult; /// Implements the entire client/server message exchange. /// State here is preserved between requests. @@ -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..1c82a866 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,7 @@ 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 +78,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..81f00b88 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. @@ -22,8 +22,8 @@ impl QueryEngine { /// Parse client request and rewrite it, if necessary. pub(super) fn parse_and_rewrite( &mut self, - context: &mut QueryEngineContext<'_>, - ) -> Result { + context: & mut QueryEngineContext<'_>, + ) -> 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,13 @@ 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(true) + Ok(None) + } } } 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) From a071dcaa09a0450142803d39ba095346de83a06d Mon Sep 17 00:00:00 2001 From: jkaczman Date: Mon, 24 Aug 2026 13:48:54 -0400 Subject: [PATCH 2/2] Cargo fmt --- pgdog/src/frontend/client/query_engine/mod.rs | 2 +- pgdog/src/frontend/client/query_engine/query.rs | 6 +++++- pgdog/src/frontend/client/query_engine/rewrite.rs | 7 +++---- 3 files changed, 9 insertions(+), 6 deletions(-) diff --git a/pgdog/src/frontend/client/query_engine/mod.rs b/pgdog/src/frontend/client/query_engine/mod.rs index f1cc761e..bc2c1d35 100644 --- a/pgdog/src/frontend/client/query_engine/mod.rs +++ b/pgdog/src/frontend/client/query_engine/mod.rs @@ -40,12 +40,12 @@ 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; use two_pc::TwoPc; pub use two_pc::phase::TwoPcPhase; -use crate::frontend::router::parser::rewrite::statement::plan::RewriteResult; /// Implements the entire client/server message exchange. /// State here is preserved between requests. diff --git a/pgdog/src/frontend/client/query_engine/query.rs b/pgdog/src/frontend/client/query_engine/query.rs index 1c82a866..7158402d 100644 --- a/pgdog/src/frontend/client/query_engine/query.rs +++ b/pgdog/src/frontend/client/query_engine/query.rs @@ -60,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, query_planner)).await; + let result = safe_timeout( + query_timeout, + self.client_server_exchange(context, query_planner), + ) + .await; match result { Ok(response) => response?, diff --git a/pgdog/src/frontend/client/query_engine/rewrite.rs b/pgdog/src/frontend/client/query_engine/rewrite.rs index 81f00b88..43c2f306 100644 --- a/pgdog/src/frontend/client/query_engine/rewrite.rs +++ b/pgdog/src/frontend/client/query_engine/rewrite.rs @@ -22,7 +22,7 @@ impl QueryEngine { /// Parse client request and rewrite it, if necessary. pub(super) fn parse_and_rewrite( &mut self, - context: & mut QueryEngineContext<'_>, + context: &mut QueryEngineContext<'_>, ) -> Result, Error> { let use_parser = self .backend @@ -45,8 +45,7 @@ impl QueryEngine { context.client_request.ast = Some(ast); Ok(Some(rewrite_result)) } else { - - Ok(None) - } + Ok(None) + } } }