From eb76b315987088159a979d7fadb06a7ac80e81a3 Mon Sep 17 00:00:00 2001 From: Michael Scrivo Date: Mon, 10 Aug 2026 11:56:57 -0400 Subject: [PATCH 1/2] fix: don't sum row counts across shards for EXECUTE on omnisharded tables EXECUTE of a server-side prepared statement was routed through the DDL catch-all, which broadcasts to all shards without the omnisharded flag. For statements that only touch omnisharded tables, the CommandComplete row counts from each shard were summed, so a DELETE/UPDATE of N rows reported N * shards rows to the client. Route PREPARE explicitly and store the statement behind the name in the prepared statements cache. EXECUTE resolves the name and parses the statement through the statement cache (once per unique statement, not per execution). Writes that only touch omnisharded tables carry the omnisharded flag, so cross-shard results are deduplicated, not aggregated. Reads are not flagged: EXECUTE always routes as a write, and an omnisharded write requires full shard coverage, which would reject shard directives on statements that can't diverge the shards. Co-Authored-By: Claude Fable 5 --- .../client/query_engine/route_query.rs | 6 +- pgdog/src/frontend/router/context.rs | 15 ++- .../frontend/router/parser/query/execute.rs | 127 ++++++++++++++++++ pgdog/src/frontend/router/parser/query/mod.rs | 5 + 4 files changed, 150 insertions(+), 3 deletions(-) create mode 100644 pgdog/src/frontend/router/parser/query/execute.rs diff --git a/pgdog/src/frontend/client/query_engine/route_query.rs b/pgdog/src/frontend/client/query_engine/route_query.rs index 0bb668bcd..b7874f16e 100644 --- a/pgdog/src/frontend/client/query_engine/route_query.rs +++ b/pgdog/src/frontend/client/query_engine/route_query.rs @@ -93,7 +93,8 @@ impl QueryEngine { context.params, context.transaction, context.sticky, - )?; + )? + .with_prepared_statements(context.prepared_statements); let mut result = self.router.query(router_context).map(|_| ()); // Resolve sharding key lookups that missed the cache and route @@ -115,7 +116,8 @@ impl QueryEngine { context.transaction, context.sticky, )? - .with_resolved_lookups(resolved); + .with_resolved_lookups(resolved) + .with_prepared_statements(context.prepared_statements); result = self.router.query(router_context).map(|_| ()); // Defensive: can't happen unless routing stops diff --git a/pgdog/src/frontend/router/context.rs b/pgdog/src/frontend/router/context.rs index df99253ef..6a8f10dba 100644 --- a/pgdog/src/frontend/router/context.rs +++ b/pgdog/src/frontend/router/context.rs @@ -2,7 +2,7 @@ use super::{Error, ParameterHints}; use crate::{ backend::{Cluster, Schema}, frontend::{ - BufferedQuery, ClientRequest, + BufferedQuery, ClientRequest, PreparedStatements, client::{Sticky, TransactionType}, router::Ast, router::sharding::ResolvedLookups, @@ -42,6 +42,9 @@ pub struct RouterContext<'a> { /// reads these before the lookup cache, so a second routing pass /// after resolving lookups can't miss. pub resolved_lookups: ResolvedLookups, + /// Client's prepared statements, used to route `EXECUTE` + /// based on the statement behind the name. + pub prepared_statements: Option<&'a mut PreparedStatements>, } impl<'a> RouterContext<'a> { @@ -71,6 +74,7 @@ impl<'a> RouterContext<'a> { schema: cluster.schema(), client_request: buffer, resolved_lookups: ResolvedLookups::default(), + prepared_statements: None, }) } @@ -80,6 +84,15 @@ impl<'a> RouterContext<'a> { self } + /// Give the router access to the client's prepared statements. + pub fn with_prepared_statements( + mut self, + prepared_statements: &'a mut PreparedStatements, + ) -> Self { + self.prepared_statements = Some(prepared_statements); + self + } + pub fn in_transaction(&self) -> bool { self.transaction.is_some() } diff --git a/pgdog/src/frontend/router/parser/query/execute.rs b/pgdog/src/frontend/router/parser/query/execute.rs new file mode 100644 index 000000000..1269c4888 --- /dev/null +++ b/pgdog/src/frontend/router/parser/query/execute.rs @@ -0,0 +1,127 @@ +//! Routing for SQL-level `PREPARE` and `EXECUTE` statements. + +use tracing::warn; + +use crate::frontend::BufferedQuery; +use crate::net::Parse; + +use super::*; + +impl QueryParser { + /// Route a SQL-level `PREPARE` statement. + /// + /// It's broadcast to all shards. The statement behind the name is + /// stored in the prepared statements cache, so `EXECUTE` can be + /// routed based on it. + pub(super) fn prepare_statement( + stmt: &nodes::PrepareStmt, + context: &mut QueryParserContext, + ) -> Result { + if let Some(name) = stmt.name() + && let Some(prepared_statements) = + context.router_context.prepared_statements.as_deref_mut() + // First PREPARE wins: a duplicate fails on the server, + // which keeps the original statement. + && prepared_statements.name(name).is_none() + { + match pg_raw_parse::deparse(stmt.query()) { + Ok(query) => { + let mut parse = Parse::named(name, query.as_str()); + prepared_statements.insert(&mut parse); + } + Err(err) => { + warn!("failed to record PREPARE statement: {}", err); + } + } + } + + context + .shards_calculator + .push(ShardWithPriority::new_table(Shard::All)); + + Ok(Command::Query(Route::write( + context.shards_calculator.shard(), + ))) + } + + /// Route `EXECUTE ` of a server-side prepared statement. + /// + /// `PREPARE` is broadcast to all shards, so `EXECUTE` is broadcast as + /// well. If the statement behind the name is a write that only touches + /// omnisharded tables, mark the route, so results are deduplicated + /// across shards instead of aggregated, e.g. `UPDATE ` reports + /// the row count from one shard, not the sum of all of them. + pub(super) fn execute_prepared( + stmt: &nodes::ExecuteStmt, + context: &mut QueryParserContext, + ) -> Result { + let omnisharded = Self::executed_statement_omnisharded(stmt, context); + + let shard = if omnisharded { + ShardWithPriority::new_table_omni(Shard::All) + } else { + ShardWithPriority::new_table(Shard::All) + }; + context.shards_calculator.push(shard); + + Ok(Command::Query( + Route::write(context.shards_calculator.shard()).with_omnisharded(omnisharded), + )) + } + + /// Check if the statement behind an `EXECUTE` name is a write that + /// only touches omnisharded tables. + /// + /// `PREPARE` accepts SELECT, INSERT, UPDATE, DELETE, MERGE and VALUES. + /// Only writes are flagged: `EXECUTE` always routes as a write, and the + /// omnisharded flag on a write requires full shard coverage, which + /// would reject shard directives on read-only statements. MERGE is + /// left out conservatively; its row counts keep aggregating. + fn executed_statement_omnisharded( + stmt: &nodes::ExecuteStmt, + context: &mut QueryParserContext, + ) -> bool { + if context.sharding_schema.tables.omnishards().is_empty() { + return false; + } + + let Some(name) = stmt.name() else { + return false; + }; + let Some(prepared_statements) = context.router_context.prepared_statements.as_deref_mut() + else { + return false; + }; + let Some(parse) = prepared_statements.parse(name) else { + return false; + }; + + // The statement cache parses each unique statement once, + // not on every EXECUTE. + let ast_context = AstContext { + sharding_schema: context.sharding_schema.clone(), + db_schema: context.router_context.schema.clone(), + user: context.router_context.cluster.user(), + search_path: context.router_context.parameter_hints.search_path, + }; + let Ok(ast) = Cache::get().query( + &BufferedQuery::Prepared(parse), + &ast_context, + prepared_statements, + ) else { + return false; + }; + + let Some(root) = ast.ast.stmts().next() else { + return false; + }; + if !matches!( + root, + Node::InsertStmt(_) | Node::UpdateStmt(_) | Node::DeleteStmt(_) + ) { + return false; + } + + StatementParser::new(root, None, &context.sharding_schema, None).is_all_omnisharded() + } +} diff --git a/pgdog/src/frontend/router/parser/query/mod.rs b/pgdog/src/frontend/router/parser/query/mod.rs index 5f480bdfc..ca43345b8 100644 --- a/pgdog/src/frontend/router/parser/query/mod.rs +++ b/pgdog/src/frontend/router/parser/query/mod.rs @@ -23,6 +23,7 @@ use super::{ }; mod ddl; mod delete; +mod execute; mod explain; mod plugins; mod select; @@ -407,6 +408,10 @@ impl QueryParser { Node::ExplainStmt(stmt) => self.explain(&statement, stmt, context), + Node::PrepareStmt(stmt) => Self::prepare_statement(stmt, context), + + Node::ExecuteStmt(stmt) => Self::execute_prepared(stmt, context), + Node::DiscardStmt { .. } => { return Ok(Command::Discard { extended: !context.query()?.simple(), From 1ad0939d67fb9d8e67006803821b8aae5b9683dd Mon Sep 17 00:00:00 2001 From: Michael Scrivo Date: Mon, 10 Aug 2026 11:56:57 -0400 Subject: [PATCH 2/2] test: cover EXECUTE routing for omnisharded and regular statements Covers writes vs reads behind the name, VALUES, shard directives on both (allowed on reads, rejected on omnisharded writes), sharded tables and unknown statement names. Co-Authored-By: Claude Fable 5 --- .../frontend/router/parser/query/test/mod.rs | 1 + .../router/parser/query/test/setup.rs | 3 +- .../router/parser/query/test/test_execute.rs | 160 ++++++++++++++++++ 3 files changed, 163 insertions(+), 1 deletion(-) create mode 100644 pgdog/src/frontend/router/parser/query/test/test_execute.rs diff --git a/pgdog/src/frontend/router/parser/query/test/mod.rs b/pgdog/src/frontend/router/parser/query/test/mod.rs index e9e744c85..b25276cff 100644 --- a/pgdog/src/frontend/router/parser/query/test/mod.rs +++ b/pgdog/src/frontend/router/parser/query/test/mod.rs @@ -26,6 +26,7 @@ pub mod test_comments; pub mod test_ddl; pub mod test_delete; pub mod test_dml; +pub mod test_execute; pub mod test_explain; pub mod test_functions; pub mod test_insert; diff --git a/pgdog/src/frontend/router/parser/query/test/setup.rs b/pgdog/src/frontend/router/parser/query/test/setup.rs index eac948163..44cd02b87 100644 --- a/pgdog/src/frontend/router/parser/query/test/setup.rs +++ b/pgdog/src/frontend/router/parser/query/test/setup.rs @@ -217,7 +217,8 @@ impl QueryParserTest { self.sticky, ) .unwrap() - .with_resolved_lookups(self.resolved_lookups.clone()); + .with_resolved_lookups(self.resolved_lookups.clone()) + .with_prepared_statements(&mut self.prepared); let command = self.parser.parse(router_ctx)?; Ok(command.clone()) diff --git a/pgdog/src/frontend/router/parser/query/test/test_execute.rs b/pgdog/src/frontend/router/parser/query/test/test_execute.rs new file mode 100644 index 000000000..a8639a963 --- /dev/null +++ b/pgdog/src/frontend/router/parser/query/test/test_execute.rs @@ -0,0 +1,160 @@ +//! Routing tests for SQL-level `PREPARE`/`EXECUTE` statements. +//! +//! `EXECUTE` must be routed based on the statement behind the name. If that +//! statement is a write that only touches omnisharded tables, the results +//! are identical on all shards, so the response (e.g. `UPDATE `) must +//! be deduplicated across shards instead of aggregated. + +use crate::frontend::router::parser::{Error, Shard}; + +use super::setup::{QueryParserTest, *}; + +#[test] +fn test_execute_omni_update_is_omnisharded() { + let mut test = QueryParserTest::new(); + test.execute(vec![ + Query::new("PREPARE upd AS UPDATE sharded_omni SET value = $1").into(), + ]); + + let command = test.execute(vec![Query::new("EXECUTE upd('x')").into()]); + + let route = command.route(); + assert!(route.is_write()); + assert_eq!(route.shard(), &Shard::All); + assert!( + route.is_omnisharded(), + "EXECUTE of an omnisharded UPDATE must carry the omnisharded flag, got {:?}", + route + ); +} + +#[test] +fn test_execute_omni_delete_is_omnisharded() { + let mut test = QueryParserTest::new(); + test.execute(vec![ + Query::new("PREPARE del AS DELETE FROM sharded_omni WHERE id = $1").into(), + ]); + + let command = test.execute(vec![Query::new("EXECUTE del(1)").into()]); + + let route = command.route(); + assert!(route.is_write()); + assert_eq!(route.shard(), &Shard::All); + assert!( + route.is_omnisharded(), + "EXECUTE of an omnisharded DELETE must carry the omnisharded flag, got {:?}", + route + ); +} + +#[test] +fn test_execute_omni_insert_is_omnisharded() { + let mut test = QueryParserTest::new(); + test.execute(vec![ + Query::new("PREPARE ins AS INSERT INTO sharded_omni (id, value) VALUES ($1, $2)").into(), + ]); + + let command = test.execute(vec![Query::new("EXECUTE ins(1, 'a')").into()]); + + let route = command.route(); + assert_eq!(route.shard(), &Shard::All); + assert!(route.is_omnisharded()); +} + +/// Reads are not flagged: `EXECUTE` always routes as a write, and an +/// omnisharded write requires full shard coverage, which would reject +/// shard directives on statements that can't diverge the shards. +#[test] +fn test_execute_omni_select_not_omnisharded() { + let mut test = QueryParserTest::new(); + test.execute(vec![ + Query::new("PREPARE sel AS SELECT * FROM sharded_omni WHERE id = $1").into(), + ]); + + let command = test.execute(vec![Query::new("EXECUTE sel(1)").into()]); + + let route = command.route(); + assert_eq!(route.shard(), &Shard::All); + assert!(!route.is_omnisharded()); +} + +#[test] +fn test_execute_values_not_omnisharded() { + let mut test = QueryParserTest::new(); + test.execute(vec![Query::new("PREPARE vals AS VALUES (1), (2)").into()]); + + let command = test.execute(vec![Query::new("EXECUTE vals").into()]); + + assert!(!command.route().is_omnisharded()); +} + +/// A shard directive on `EXECUTE` of a read-only statement is allowed; +/// the statement can't diverge the shards. +#[test] +fn test_execute_omni_select_with_shard_directive() { + let mut test = QueryParserTest::new(); + test.execute(vec![ + Query::new("PREPARE sel AS SELECT * FROM sharded_omni WHERE id = $1").into(), + ]); + + let command = test.execute(vec![ + Query::new("/* pgdog_shard: 0 */ EXECUTE sel(1)").into(), + ]); + + assert_eq!(command.route().shard(), &Shard::Direct(0)); +} + +/// A shard directive on `EXECUTE` of an omnisharded write is rejected, +/// like on the equivalent direct statement: reaching only one shard +/// would silently diverge the table. +#[test] +fn test_execute_omni_write_with_shard_directive_rejected() { + let mut test = QueryParserTest::new(); + test.execute(vec![ + Query::new("PREPARE upd AS UPDATE sharded_omni SET value = $1").into(), + ]); + + let result = test.try_execute(vec![ + Query::new("/* pgdog_shard: 0 */ EXECUTE upd('x')").into(), + ]); + + assert!(matches!(result, Err(Error::OmniWriteWithDirective))); +} + +#[test] +fn test_prepare_routes_to_all_shards() { + let mut test = QueryParserTest::new(); + let command = test.execute(vec![ + Query::new("PREPARE upd AS UPDATE sharded_omni SET value = $1").into(), + ]); + + let route = command.route(); + assert!(route.is_write()); + assert_eq!(route.shard(), &Shard::All); +} + +#[test] +fn test_execute_sharded_table_not_omnisharded() { + let mut test = QueryParserTest::new(); + test.execute(vec![ + Query::new("PREPARE upd AS UPDATE sharded SET value = $1").into(), + ]); + + let command = test.execute(vec![Query::new("EXECUTE upd('x')").into()]); + + let route = command.route(); + assert!(route.is_write()); + assert_eq!(route.shard(), &Shard::All); + assert!(!route.is_omnisharded()); +} + +#[test] +fn test_execute_unknown_statement_not_omnisharded() { + let mut test = QueryParserTest::new(); + let command = test.execute(vec![Query::new("EXECUTE not_prepared(1)").into()]); + + let route = command.route(); + assert!(route.is_write()); + assert_eq!(route.shard(), &Shard::All); + assert!(!route.is_omnisharded()); +}