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(), 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()); +}