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

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

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

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<Command, Error> {
+        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 <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 is a write that 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: &nodes::ExecuteStmt,
+        context: &mut QueryParserContext,
+    ) -> Result<Command, Error> {
+        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()
+    }
+}

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

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

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

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