diff --git a/pgdog/src/backend/pool/cluster.rs b/pgdog/src/backend/pool/cluster.rs index f11729c48..673a851d2 100644 --- a/pgdog/src/backend/pool/cluster.rs +++ b/pgdog/src/backend/pool/cluster.rs @@ -958,6 +958,10 @@ mod test { self.rw_strategy = rw_strategy; } + pub fn set_sharded_tables(&mut self, sharded_tables: ShardedTables) { + self.sharded_tables = sharded_tables; + } + pub fn set_prefer_primary(&mut self, prefer_primary: bool) { self.prefer_primary = prefer_primary; } diff --git a/pgdog/src/frontend/client/query_engine/route_query.rs b/pgdog/src/frontend/client/query_engine/route_query.rs index abd8e90e1..d563e0004 100644 --- a/pgdog/src/frontend/client/query_engine/route_query.rs +++ b/pgdog/src/frontend/client/query_engine/route_query.rs @@ -91,7 +91,8 @@ impl QueryEngine { context.params, context.transaction, context.sticky, - )?; + )? + .with_prepared_statements(context.prepared_statements); match self.router.query(router_context) { Ok(command) => { context.client_request.route = Some(command.route().clone()); diff --git a/pgdog/src/frontend/prepared_statements/mod.rs b/pgdog/src/frontend/prepared_statements/mod.rs index 0fc270c9a..c4c101aec 100644 --- a/pgdog/src/frontend/prepared_statements/mod.rs +++ b/pgdog/src/frontend/prepared_statements/mod.rs @@ -32,6 +32,10 @@ fn str_mem(s: &str) -> usize { pub struct PreparedStatements { pub(super) global: Arc>, pub(super) local: HashMap, + /// SQL-level `PREPARE` statements sent over the simple protocol, + /// name -> statement text. Tracked so `EXECUTE` can be routed + /// based on the statement behind the name. + pub(super) simple: HashMap, pub(super) level: PreparedStatementsLevel, pub(super) memory_used: usize, } @@ -41,6 +45,7 @@ impl Default for PreparedStatements { Self { global: Arc::new(RwLock::new(GlobalCache::default())), local: HashMap::default(), + simple: HashMap::default(), level: PreparedStatementsLevel::Extended, memory_used: 0, } @@ -110,6 +115,21 @@ impl PreparedStatements { self.local.get(name) } + /// Record a SQL-level `PREPARE` statement. + pub fn insert_simple(&mut self, name: &str, query: &str) { + if let Some(old_query) = self.simple.insert(name.to_owned(), query.to_owned()) { + self.memory_used = self.memory_used.saturating_sub(str_mem(&old_query)); + self.memory_used += str_mem(query); + } else { + self.memory_used += str_mem(name) + str_mem(query); + } + } + + /// Get the query behind a SQL-level `PREPARE` statement name. + pub fn simple_query(&self, name: &str) -> Option<&str> { + self.simple.get(name).map(|query| query.as_str()) + } + /// Get globally-prepared statement by local name. pub fn parse(&self, name: &str) -> Option { self.local @@ -154,6 +174,7 @@ impl PreparedStatements { } self.local.clear(); + self.simple.clear(); self.memory_used = 0; } @@ -230,6 +251,28 @@ mod test { assert!(statements.local.is_empty()); } + #[test] + fn test_insert_simple_overwrite() { + let mut statements = PreparedStatements::default(); + + statements.insert_simple("upd", "UPDATE one SET x = 1"); + let memory_used = statements.memory_used(); + + statements.insert_simple("upd", "UPDATE two SET x = 22"); + assert_eq!( + statements.simple_query("upd"), + Some("UPDATE two SET x = 22") + ); + assert_eq!( + statements.memory_used(), + memory_used + "UPDATE two SET x = 22".len() - "UPDATE one SET x = 1".len() + ); + + statements.close_all(); + assert_eq!(statements.simple_query("upd"), None); + assert_eq!(statements.memory_used(), 0); + } + #[test] fn test_counted_only_once_per_client() { let mut statements = PreparedStatements::default(); diff --git a/pgdog/src/frontend/router/context.rs b/pgdog/src/frontend/router/context.rs index 3e9e0c922..9db650ee1 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, }, @@ -37,6 +37,9 @@ pub struct RouterContext<'a> { pub schema: Schema, /// Original client request. pub client_request: &'a ClientRequest, + /// 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> { @@ -65,9 +68,19 @@ impl<'a> RouterContext<'a> { ast: buffer.ast.clone(), schema: cluster.schema(), client_request: buffer, + prepared_statements: None, }) } + /// 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/prepare.rs b/pgdog/src/frontend/router/parser/prepare.rs index 910404d46..609407f7d 100644 --- a/pgdog/src/frontend/router/parser/prepare.rs +++ b/pgdog/src/frontend/router/parser/prepare.rs @@ -26,6 +26,14 @@ impl Prepare { statement, }) } + + pub fn name(&self) -> &str { + &self.name + } + + pub fn statement(&self) -> &str { + &self.statement + } } #[cfg(test)] 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..25e0e67c6 --- /dev/null +++ b/pgdog/src/frontend/router/parser/query/execute.rs @@ -0,0 +1,126 @@ +//! Routing for SQL-level `PREPARE` and `EXECUTE` statements. + +#[cfg(not(feature = "new_parser"))] +use pg_query::parse; + +use crate::frontend::PreparedStatements; + +use super::*; + +impl QueryParser { + /// Route a SQL-level `PREPARE` statement. + /// + /// It's broadcast to all shards. The statement behind the name is + /// recorded, so `EXECUTE` can be routed based on it. + pub(super) fn prepare_statement( + stmt: &PrepareStmt, + context: &mut QueryParserContext, + ) -> Result { + if let Ok(prepare) = Prepare::from_stmt(stmt, context.sharding_schema.query_parser_engine) + && let Some(prepared_statements) = + context.router_context.prepared_statements.as_deref_mut() + { + prepared_statements.insert_simple(prepare.name(), prepare.statement()); + } + + 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 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: &ExecuteStmt, + context: &mut QueryParserContext, + ) -> Result { + let omnisharded = Self::prepared_statement_omnisharded(&stmt.name, 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 query behind a prepared statement name only touches + /// omnisharded tables. + fn prepared_statement_omnisharded(name: &str, context: &QueryParserContext) -> bool { + let schema = &context.sharding_schema; + if schema.tables.omnishards().is_empty() { + return false; + } + + // SQL-level PREPARE, recorded when it was routed. + if let Some(prepared_statements) = context.router_context.prepared_statements.as_deref() + && let Some(query) = prepared_statements.simple_query(name) + { + return Self::statement_omnisharded(query, schema); + } + + // Globally cached statement, e.g. renamed by the `prepared_statements = "full"` + // rewrite. Copy the query so the cache isn't locked while we parse it. + let query = { + let global = PreparedStatements::global(); + let guard = global.read(); + match guard.query(name) { + Some(query) => query.to_string(), + None => return false, + } + }; + + Self::statement_omnisharded(&query, schema) + } + + #[cfg(not(feature = "new_parser"))] + fn statement_omnisharded(query: &str, schema: &ShardingSchema) -> bool { + let Ok(ast) = parse(query) else { + return false; + }; + let Some(root) = ast + .protobuf + .stmts + .first() + .and_then(|stmt| stmt.stmt.as_ref()) + .and_then(|stmt| stmt.node.as_ref()) + else { + return false; + }; + + let mut parser = match root { + NodeEnum::SelectStmt(stmt) => StatementParser::from_select(stmt, None, schema, None), + NodeEnum::UpdateStmt(stmt) => StatementParser::from_update(stmt, None, schema, None), + NodeEnum::DeleteStmt(stmt) => StatementParser::from_delete(stmt, None, schema, None), + NodeEnum::InsertStmt(stmt) => StatementParser::from_insert(stmt, None, schema, None), + _ => return false, + }; + + parser.is_all_omnisharded() + } + + #[cfg(feature = "new_parser")] + fn statement_omnisharded(query: &str, schema: &ShardingSchema) -> bool { + let Ok(ast) = pg_raw_parse::parse(query) else { + return false; + }; + let Some(node) = ast.stmts().next() else { + return false; + }; + + StatementParser::from_select(node, None, 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 63f12ff8c..a7d823ba8 100644 --- a/pgdog/src/frontend/router/parser/query/mod.rs +++ b/pgdog/src/frontend/router/parser/query/mod.rs @@ -26,6 +26,7 @@ use super::{ }; mod ddl; mod delete; +mod execute; mod explain; mod plugins; mod select; @@ -383,6 +384,10 @@ impl QueryParser { Some(NodeEnum::ExplainStmt(ref stmt)) => self.explain(&statement, stmt, context), + Some(NodeEnum::PrepareStmt(ref stmt)) => Self::prepare_statement(stmt, context), + + Some(NodeEnum::ExecuteStmt(ref stmt)) => Self::execute_prepared(stmt, context), + Some(NodeEnum::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 511db73d6..bfe559651 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 3bce17175..1a154bd61 100644 --- a/pgdog/src/frontend/router/parser/query/test/setup.rs +++ b/pgdog/src/frontend/router/parser/query/test/setup.rs @@ -95,6 +95,12 @@ impl QueryParserTest { self } + /// Replace the sharded/omnisharded tables on the cluster. + pub(crate) fn with_sharded_tables(mut self, tables: crate::backend::ShardedTables) -> Self { + self.cluster.set_sharded_tables(tables); + self + } + /// Route reads to the primary by default on the cluster. pub(crate) fn with_prefer_primary(mut self, prefer_primary: bool) -> Self { self.cluster.set_prefer_primary(prefer_primary); @@ -176,7 +182,8 @@ impl QueryParserTest { self.transaction, self.sticky, ) - .unwrap(); + .unwrap() + .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..71655d7c1 --- /dev/null +++ b/pgdog/src/frontend/router/parser/query/test/test_execute.rs @@ -0,0 +1,171 @@ +//! Routing tests for SQL-level `PREPARE`/`EXECUTE` statements. +//! +//! `EXECUTE` must be routed based on the statement behind the name. If that +//! statement 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::PreparedStatements; +use crate::frontend::router::parser::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_select_is_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_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()); +} + +#[test] +fn test_execute_without_omnisharded_tables_not_omnisharded() { + let mut test = QueryParserTest::new().with_sharded_tables(crate::backend::ShardedTables::new( + vec![], + vec![], + false, + Default::default(), + )); + 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()]); + + assert!(!command.route().is_omnisharded()); +} + +#[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()); +} + +/// Statements that resolve but can't be routed by table are not omnisharded. +#[test] +fn test_execute_unroutable_statements_not_omnisharded() { + for query in [ + "NOT VALID SQL", + "-- just a comment", + "CREATE TABLE t (id BIGINT)", + ] { + let parse = Parse::new_anonymous(query); + let name = PreparedStatements::global().write().insert_anyway(&parse); + + let mut test = QueryParserTest::new(); + let command = test.execute(vec![Query::new(format!("EXECUTE {}(1)", name)).into()]); + + assert!( + !command.route().is_omnisharded(), + "{:?} must not be omnisharded", + query + ); + } +} + +/// With `prepared_statements = "full"`, `EXECUTE` names are rewritten to +/// globally unique names before routing. Those resolve through the global +/// prepared statements cache instead. +#[test] +fn test_execute_global_name_resolves_through_global_cache() { + let parse = Parse::new_anonymous("DELETE FROM sharded_omni WHERE id = $1"); + let name = PreparedStatements::global().write().insert_anyway(&parse); + + let mut test = QueryParserTest::new(); + let command = test.execute(vec![Query::new(format!("EXECUTE {}(1)", name)).into()]); + + let route = command.route(); + assert!(route.is_write()); + assert_eq!(route.shard(), &Shard::All); + assert!(route.is_omnisharded()); +}