fix: return an error for INSERT .. SELECT into sharded tables when router can't find shards of rows

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,

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,
 

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(

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

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<StatementParameters<'b>>,
@@ -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,16 +1257,29 @@ 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)
         }
     }
 }
 
+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};

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| {