fix: don't sum row counts across shards for EXECUTE on omnisharded tables

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;
         }

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());

pgdog/src/frontend/prepared_statements/mod.rs

@@ -32,6 +32,10 @@ fn str_mem(s: &str) -> usize {
 pub struct PreparedStatements {
     pub(super) global: Arc<RwLock<GlobalCache>>,
     pub(super) local: HashMap<String, String>,
+    /// 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<String, String>,
     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<Parse> {
         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();

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()
     }

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)]

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<Command, Error> {
+        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 <name>` 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 <rows>` 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<Command, Error> {
+        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()
+    }
+}

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(),

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;

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())