From 6df7724b60a3e886b787d909f94a39bcc8f29beb Mon Sep 17 00:00:00 2001 From: jkaczman Date: Fri, 6 Nov 2026 10:29:33 -0800 Subject: [PATCH] fix: return an error for INSERT .. SELECT into sharded tables when router can't find shards of rows --- .../client/query_engine/route_query.rs | 9 +++ pgdog/src/frontend/router/parser/error.rs | 3 + pgdog/src/frontend/router/parser/query/mod.rs | 8 ++- .../router/parser/query/test/test_insert.rs | 70 +++++++++++++++++-- pgdog/src/frontend/router/parser/statement.rs | 46 +++++++++++- pgdog/src/frontend/router/sharding/tables.rs | 9 +++ 6 files changed, 137 insertions(+), 8 deletions(-) diff --git a/pgdog/src/frontend/client/query_engine/route_query.rs b/pgdog/src/frontend/client/query_engine/route_query.rs index 8591a0805..b3d877991 100644 --- a/pgdog/src/frontend/client/query_engine/route_query.rs +++ b/pgdog/src/frontend/client/query_engine/route_query.rs @@ -199,6 +199,15 @@ impl QueryEngine { return Ok(false); } + Err(RouterError::Parser(err @ ParserError::InsertSelectWithoutShardingKey)) => { + if let Some(transaction) = context.transaction.as_mut() { + transaction.set_error(); + } + self.error_response(context, client_request, ErrorResponse::unsupported(&err)) + .await?; + + return Ok(false); + } Err(RouterError::Parser(ParserError::UnmappedShardKey(shard_key))) => { self.error_response( context, diff --git a/pgdog/src/frontend/router/parser/error.rs b/pgdog/src/frontend/router/parser/error.rs index 9dd6e757b..6b35dd4d7 100644 --- a/pgdog/src/frontend/router/parser/error.rs +++ b/pgdog/src/frontend/router/parser/error.rs @@ -103,6 +103,9 @@ pub(crate) enum Error { #[error("unmapped sharding key was specified")] UnmappedShardKey(String), + #[error("INSERT .. SELECT into sharded table is not supported without explicit sharding key")] + InsertSelectWithoutShardingKey, + #[error("execute requires prepared statements to be set to full")] ExecuteRequiresFull, diff --git a/pgdog/src/frontend/router/parser/query/mod.rs b/pgdog/src/frontend/router/parser/query/mod.rs index 24d60501d..b3ee8605e 100644 --- a/pgdog/src/frontend/router/parser/query/mod.rs +++ b/pgdog/src/frontend/router/parser/query/mod.rs @@ -667,10 +667,16 @@ impl QueryParser { user: context.router_context.cluster.user(), search_path: context.router_context.parameter_hints.search_path, }; + let insert_rewritten = context + .router_context + .ast + .is_some_and(|ast| !ast.rewrite_plan.subqueries.is_empty()); + let mut parser = StatementParser::new(stmt, context.router_context.bind, context.sharding_schema) .with_schema_lookup(schema_lookup) - .with_explain(self.recorder_mut().is_some()); + .with_explain(self.recorder_mut().is_some()) + .with_insert_rewritten(insert_rewritten); parser.set_resolved_lookups(&context.router_context.resolved_lookups); let is_sharded = parser.is_sharded( diff --git a/pgdog/src/frontend/router/parser/query/test/test_insert.rs b/pgdog/src/frontend/router/parser/query/test/test_insert.rs index 7517d046c..81680ee5e 100644 --- a/pgdog/src/frontend/router/parser/query/test/test_insert.rs +++ b/pgdog/src/frontend/router/parser/query/test/test_insert.rs @@ -1,4 +1,6 @@ -use crate::frontend::router::parser::Shard; +use std::ops::Deref; + +use crate::frontend::router::parser::{Error, Shard}; use crate::net::messages::bind::Parameter; use super::setup::*; @@ -67,17 +69,48 @@ fn test_insert_multi_row() { #[test] fn test_insert_select() { + // Without subquery rewrite being on, we don't know the sharding keys, so we must error (as opposed to incorrectly routing) + for sql in [ + "INSERT INTO sharded (id, value) SELECT id, value FROM other_table WHERE id = 1", + "INSERT INTO sharded (id) SELECT g FROM generate_series(1, 10) g", + "INSERT INTO sharded (id) SELECT generate_series(1, 3)", + "INSERT INTO sharded (id, value) SELECT * FROM json_to_recordset('[]') AS t(id bigint, value text)", + "INSERT INTO sharded (id, value) SELECT * FROM (VALUES (1, 'a'), (2, 'b')) v(id, value)", + "INSERT INTO sharded (id, value) SELECT 1, 'a' UNION ALL SELECT 2, 'b'", + "INSERT INTO sharded (id, value) SELECT * FROM unnest(ARRAY[1, 2], ARRAY['a', 'b'])", + "INSERT INTO sharded SELECT * FROM other_table", + ] { + let mut test = QueryParserTest::new(); + let error = test.try_execute(vec![Query::new(sql).into()]).unwrap_err(); + + assert!( + matches!(error, Error::InsertSelectWithoutShardingKey), + "{sql}: {error:?}" + ); + } + + // Without the sharding key column specified, we assume that the default for the column generates the keys on the shard that will perform the INSERT, thus, round-robin is + // acceptable and stays as-is (no error) + for sql in [ + "INSERT INTO sharded (value) SELECT v FROM generate_series(1, 10) v", + "INSERT INTO sharded (value) SELECT 'a'", + ] { + let mut test = QueryParserTest::new(); + let command = test.execute(vec![Query::new(sql).into()]); + assert!(command.route().shard().is_direct(), "{sql}"); + } + + // Allow round-robin here too since we don't have the Bind param yet let mut test = QueryParserTest::new(); - let command = test.execute(vec![ - Query::new( - "INSERT INTO sharded (id, value) SELECT id, value FROM other_table WHERE id = 1", + Parse::new_anonymous( + "INSERT INTO sharded (id, value) SELECT id, value FROM sharded WHERE id = $1", ) .into(), + Describe::new_statement("").into(), + Sync.into(), ]); - assert!(command.route().is_write()); - assert!(command.route().shard().is_direct()); } #[test] @@ -91,3 +124,28 @@ fn test_insert_default_values() { assert!(command.route().is_write()); assert!(command.route().shard().is_direct()); } + +#[test] +fn test_insert_select_with_subquery_rewrite() { + let mut config = crate::config::config().deref().clone(); + config.config.rewrite.subqueries = pgdog_config::RewriteMode::Rewrite; + crate::config::set(config.clone()).unwrap(); + + // subquery rewrite evaluates the SELECT before routing (thus, router doesn't need to error) + let mut test = QueryParserTest::new_with_config(&config); + let command = test.execute(vec![ + Query::new("INSERT INTO sharded (id, value) SELECT g, 'a' FROM generate_series(1, 10) g") + .into(), + ]); + assert!(command.route().is_write()); + + // subquery rewrite gap still errors (as UNION is not handled yet) + let mut test = QueryParserTest::new_with_config(&config); + let error = test + .try_execute(vec![ + Query::new("INSERT INTO sharded (id, value) SELECT 1, 'a' UNION ALL SELECT 2, 'b'") + .into(), + ]) + .unwrap_err(); + assert!(matches!(error, Error::InsertSelectWithoutShardingKey)); +} diff --git a/pgdog/src/frontend/router/parser/statement.rs b/pgdog/src/frontend/router/parser/statement.rs index 31c9eba60..4e533256d 100644 --- a/pgdog/src/frontend/router/parser/statement.rs +++ b/pgdog/src/frontend/router/parser/statement.rs @@ -287,6 +287,9 @@ pub(crate) struct StatementParser<'a, 'b> { /// Added for the ability to try to `shard()` the top-level statement (ignoring any subqueries within) /// to see if there's a sharding key or not, to determine solely the top-level's routing ignore_subqueries: bool, + + // If a rewrite runs the INSERT instead of this route (e.g., the subquery rewrite evaluates the SELECT first), we don't need the shards of the rows. + insert_rewritten: bool, } impl<'a, 'b: 'a> StatementParser<'a, 'b> { @@ -307,9 +310,15 @@ impl<'a, 'b: 'a> StatementParser<'a, 'b> { pending_lookups: Vec::new(), resolved_lookups: None, ignore_subqueries: false, + insert_rewritten: false, } } + pub(crate) fn with_insert_rewritten(mut self, insert_rewritten: bool) -> Self { + self.insert_rewritten = insert_rewritten; + self + } + pub(crate) fn new_that_ignore_subqueries( stmt: Node<'a>, bind: Option>, @@ -1128,6 +1137,28 @@ impl<'a, 'b: 'a> StatementParser<'a, 'b> { Ok(shard) } + /// If we're about to send a query round-robin, however, the SELECT provides the sharding key column, + /// meaning that we don't have a way to know where the query should actually go, then we error. + /// + /// This function performs those checks to see if it qualifies for the error. + fn select_sends_keys_without_shards( + &self, + stmt: &'a nodes::InsertStmt, + sharded_table: &ShardedTable, + ) -> bool { + let Node::SelectStmt(select) = stmt.select_stmt() else { + return false; + }; + + let waiting_for_bind = self.bind.is_none() && contains_param(Node::InsertStmt(stmt)); + + self.schema.shards > 1 + && select.values_lists().is_empty() + && sharded_table.insert_sends_key_column(stmt) + && !self.insert_rewritten + && !waiting_for_bind + } + /// Get column names from the INSERT statement, or look them up from schema if not specified. fn get_insert_columns( &self, @@ -1226,9 +1257,14 @@ impl<'a, 'b: 'a> StatementParser<'a, 'b> { // Round-robin fallback: if table is sharded but no sharding key found, // pick a shard at random + let tables = Tables::new(self.schema); if let Some(table) = ctx.table - && Tables::new(self.schema).sharded(table).is_some() + && let Some(sharded_table) = tables.sharded(table) { + if self.select_sends_keys_without_shards(stmt, sharded_table) { + return Err(Error::InsertSelectWithoutShardingKey); + } + Ok(Some(Shard::Direct(round_robin::next(self.schema.shards)))) } else { Ok(None) @@ -1236,6 +1272,14 @@ impl<'a, 'b: 'a> StatementParser<'a, 'b> { } } +fn contains_param(node: Node<'_>) -> bool { + walk::walk_manual(node, |node| match node { + Node::ParamRef(_) => ControlFlow::Break(()), + _ => Recurse::yes(), + }) + .is_some() +} + #[cfg(test)] mod test { use crate::frontend::router::sharding::{Mapping, ShardedTable}; diff --git a/pgdog/src/frontend/router/sharding/tables.rs b/pgdog/src/frontend/router/sharding/tables.rs index 9282c3622..175a1564f 100644 --- a/pgdog/src/frontend/router/sharding/tables.rs +++ b/pgdog/src/frontend/router/sharding/tables.rs @@ -1,3 +1,4 @@ +use pg_raw_parse::{Node, nodes}; use pgdog_config::{DataType, Hasher, LookupResult}; use pgdog_vector::Vector; @@ -27,6 +28,14 @@ pub(crate) struct ShardedTable { } impl ShardedTable { + // Returns true if the INSERT sends the sharding key column (as opposed to relying on DEFAULT) + pub(crate) fn insert_sends_key_column(&self, insert: &nodes::InsertStmt) -> bool { + insert.cols().is_empty() + || insert.cols().into_iter().any(|column| { + matches!(column, Node::ResTarget(target) if target.name() == Some(self.column.as_str())) + }) + } + /// Whether `table` matches this table's configured name (or prefix). pub(crate) fn name_matches(&self, table: &str) -> bool { self.name.as_deref().is_some_and(|name| {