diff --git a/integration/js/pg_tests/test/sequelize.js b/integration/js/pg_tests/test/sequelize.js index 1e00fd2a3..ca73c3490 100644 --- a/integration/js/pg_tests/test/sequelize.js +++ b/integration/js/pg_tests/test/sequelize.js @@ -200,20 +200,11 @@ describe("Sequelize multi-statement SET", function () { await seq.close(); }); - it("mixed SET and non-SET returns error", async function () { + it("mixed SET and non-SET succeeds and applies SET", async function () { const seq = createSequelize(); - try { - await seq.query("SET statement_timeout TO '10s'; SELECT 1"); - assert.fail("expected error for mixed SET + SELECT"); - } catch (err) { - assert.ok( - err.message.includes( - "multi-statement queries cannot mix SET with other commands", - ), - `unexpected error: ${err.message}`, - ); - } + // Should succeed: PgDog intercepts the SET and forwards SELECT to the backend. + await seq.query("SET statement_timeout TO '10s'; SELECT 1"); const [rows] = await seq.query("SELECT 1 AS val"); assert.strictEqual(rows[0].val, 1); diff --git a/integration/rust/tests/integration/multi_set.rs b/integration/rust/tests/integration/multi_set.rs index b3f53af35..2496c5e34 100644 --- a/integration/rust/tests/integration/multi_set.rs +++ b/integration/rust/tests/integration/multi_set.rs @@ -46,17 +46,113 @@ async fn test_multi_set_with_timezone_interval() { } #[tokio::test] -async fn test_multi_set_mixed_returns_error() { +async fn test_multi_set_mixed_succeeds() { for conn in connections_tokio().await { - let err = conn - .batch_execute("SET statement_timeout TO '10s'; SELECT 1") + conn.batch_execute("SET statement_timeout TO '10s'; SELECT 1") .await - .unwrap_err(); - let db_err = err.as_db_error().expect("Expected a DbError"); - let msg = db_err.message(); - assert!( - msg.contains("multi-statement queries cannot mix SET with other commands"), - "unexpected error: {msg}", - ); + .unwrap(); + + let rows = conn.simple_query("SHOW statement_timeout").await.unwrap(); + assert_eq!(extract_simple_query_value(&rows), "10s"); + } +} + +#[tokio::test] +async fn test_multi_statement_select() { + for conn in connections_tokio().await { + let msgs = conn.simple_query("SELECT 1; SELECT 2").await.unwrap(); + let mut values: Vec = msgs + .iter() + .filter_map(|m| { + if let tokio_postgres::SimpleQueryMessage::Row(row) = m { + row.get(0).map(|s| s.to_string()) + } else { + None + } + }) + .collect(); + // On a sharded connection, each scalar SELECT fans out to all shards, so + // "SELECT 1; SELECT 2" may yield ["1","1","2","2"]. Dedup consecutive + // duplicates to assert ordering: all SELECT 1 results before SELECT 2. + values.dedup(); + assert_eq!(values, vec!["1", "2"]); + } +} + +#[tokio::test] +async fn test_multi_statement_transaction_batch() { + for conn in connections_tokio().await { + // Use a real table (not TEMP) — transaction-mode pooling may route the setup + // query to a different backend than the BEGIN/COMMIT batch, so TEMP tables + // (which are session-scoped) would not be visible. + conn.batch_execute( + "CREATE TABLE IF NOT EXISTS pgdog_multi_stmt_locks ( + key TEXT PRIMARY KEY, + owner TEXT NOT NULL, + ttl TIMESTAMPTZ NOT NULL + )", + ) + .await + .unwrap(); + + conn.batch_execute("TRUNCATE pgdog_multi_stmt_locks") + .await + .unwrap(); + + conn.batch_execute( + "BEGIN;\ + DELETE FROM pgdog_multi_stmt_locks WHERE ttl < CURRENT_TIMESTAMP AT TIME ZONE 'UTC';\ + INSERT INTO pgdog_multi_stmt_locks (key, owner, ttl) \ + VALUES ('test-key', 'test-owner', NOW() + INTERVAL '1 hour') \ + ON CONFLICT DO NOTHING;\ + COMMIT;", + ) + .await + .unwrap(); + + let rows = conn + .simple_query("SELECT owner FROM pgdog_multi_stmt_locks WHERE key = 'test-key'") + .await + .unwrap(); + let owner = extract_simple_query_value(&rows); + assert_eq!(owner, "test-owner"); + } +} + +/// PostgreSQL normally wraps a multi-statement simple query in an implicit transaction, +/// so a failure in any statement rolls back all preceding ones. PgDog splits the batch +/// into individual requests, meaning each statement runs with its own auto-commit. +/// This test documents that behavior: the first INSERT succeeds even when the second fails. +#[tokio::test] +async fn test_multi_statement_implicit_transaction_behavior() { + for conn in connections_tokio().await { + // Use a real table (not TEMP) — transaction-mode pooling may route the setup + // query to a different backend than the INSERT batch, so TEMP tables + // (which are session-scoped) would not be visible. + conn.batch_execute( + "CREATE TABLE IF NOT EXISTS pgdog_multi_stmt_implicit (id INT PRIMARY KEY)", + ) + .await + .unwrap(); + + conn.batch_execute("TRUNCATE pgdog_multi_stmt_implicit") + .await + .unwrap(); + + // Second INSERT violates the primary key. In standard PostgreSQL simple query + // protocol both would roll back; with per-statement routing the first commits. + let _ = conn + .batch_execute( + "INSERT INTO pgdog_multi_stmt_implicit VALUES (1); INSERT INTO pgdog_multi_stmt_implicit VALUES (1)", + ) + .await; + + let rows = conn + .simple_query("SELECT count(*) FROM pgdog_multi_stmt_implicit WHERE id = 1") + .await + .unwrap(); + let count = extract_simple_query_value(&rows); + // The first INSERT committed; the second failed independently. + assert_eq!(count, "1"); } } diff --git a/integration/rust/tests/sqlx/multi_set.rs b/integration/rust/tests/sqlx/multi_set.rs index b268a6e7b..e3fe0ee06 100644 --- a/integration/rust/tests/sqlx/multi_set.rs +++ b/integration/rust/tests/sqlx/multi_set.rs @@ -72,25 +72,18 @@ async fn test_multi_set_in_transaction() { } #[tokio::test] -async fn test_multi_set_mixed_returns_error() { +async fn test_multi_set_mixed_succeeds() { for pool in connections_sqlx().await { let mut conn = pool.acquire().await.unwrap(); - let err = conn - .execute("SET statement_timeout TO '10s'; SELECT 1") + conn.execute("SET statement_timeout TO '10s'; SELECT 1") .await - .unwrap_err(); - assert!( - err.to_string() - .contains("multi-statement queries cannot mix SET with other commands"), - "unexpected error: {err}", - ); + .unwrap(); - // Connection should still be usable after the error. - let val: String = sqlx::query_scalar("SHOW server_version") + let timeout: String = sqlx::query_scalar("SHOW statement_timeout") .fetch_one(&mut *conn) .await .unwrap(); - assert!(!val.is_empty()); + assert_eq!(timeout, "10s"); } } diff --git a/pgdog/src/frontend/client/mod.rs b/pgdog/src/frontend/client/mod.rs index 56f9004ed..b3df19591 100644 --- a/pgdog/src/frontend/client/mod.rs +++ b/pgdog/src/frontend/client/mod.rs @@ -543,7 +543,14 @@ impl Client { } // If client sent multiple requests, split them up and execute individually. - let spliced = self.client_request.spliced()?; + let extended = self.client_request.spliced()?; + let (spliced, is_simple_splice) = if extended.is_empty() { + let simple = self.client_request.spliced_simple(); + let is_simple = !simple.is_empty(); + (simple, is_simple) + } else { + (extended, false) + }; if spliced.is_empty() { let mut context = QueryEngineContext::new(self); query_engine.handle(&mut context).await?; @@ -551,10 +558,21 @@ impl Client { } else { let total = spliced.len(); let mut reqs = spliced.into_iter().enumerate(); - self.transaction.get_or_insert(TransactionType::Implicit); + // For extended protocol pipelining, mark the client as being in an implicit + // transaction so the backend connection is kept across statements in the pipeline. + // For simple query splicing, each statement routes independently — don't force + // Implicit here; the requests_left mechanism in cleanup_backend handles connection + // reuse without affecting shard routing. + if !is_simple_splice { + self.transaction.get_or_insert(TransactionType::Implicit); + } while let Some((num, mut req)) = reqs.next() { debug!("processing spliced request {}/{}", num + 1, total); - let mut context = QueryEngineContext::new(self).spliced(&mut req, reqs.len()); + let mut context = if is_simple_splice { + QueryEngineContext::new(self).spliced_simple_query(&mut req, reqs.len()) + } else { + QueryEngineContext::new(self).spliced(&mut req, reqs.len()) + }; query_engine.handle(&mut context).await?; self.transaction = context.transaction(); diff --git a/pgdog/src/frontend/client/query_engine/context.rs b/pgdog/src/frontend/client/query_engine/context.rs index 118272fec..8fce9171d 100644 --- a/pgdog/src/frontend/client/query_engine/context.rs +++ b/pgdog/src/frontend/client/query_engine/context.rs @@ -43,6 +43,9 @@ pub struct QueryEngineContext<'a> { pub(super) query_log_stdout: bool, /// Maximum query message size before a warning is logged. pub(super) query_size_limit: Option, + /// Spliced simple query: suppress intermediate ReadyForQuery messages so the + /// client (which sent one Q) sees exactly one ReadyForQuery at the end. + pub(super) simple_query_splice: bool, } impl<'a> QueryEngineContext<'a> { @@ -66,6 +69,7 @@ impl<'a> QueryEngineContext<'a> { rewrite_result: None, query_log_stdout: client.query_log_stdout, query_size_limit: client.query_size_limit, + simple_query_splice: false, } } @@ -75,6 +79,17 @@ impl<'a> QueryEngineContext<'a> { self } + pub fn spliced_simple_query( + mut self, + req: &'a mut ClientRequest, + requests_left: usize, + ) -> Self { + self.client_request = req; + self.requests_left = requests_left; + self.simple_query_splice = true; + self + } + /// Create context from mirror. pub fn new_mirror(mirror: &'a mut Mirror, buffer: &'a mut ClientRequest) -> Self { Self { @@ -94,6 +109,7 @@ impl<'a> QueryEngineContext<'a> { rewrite_result: None, query_log_stdout: false, query_size_limit: None, + simple_query_splice: false, } } diff --git a/pgdog/src/frontend/client/query_engine/end_transaction.rs b/pgdog/src/frontend/client/query_engine/end_transaction.rs index affcc9053..2b729b536 100644 --- a/pgdog/src/frontend/client/query_engine/end_transaction.rs +++ b/pgdog/src/frontend/client/query_engine/end_transaction.rs @@ -24,7 +24,10 @@ impl QueryEngine { vec![] }; messages.push(cmd.message()?); - messages.push(ReadyForQuery::idle().message()?); + let suppress_rfq = context.simple_query_splice && context.requests_left > 0; + if !suppress_rfq { + messages.push(ReadyForQuery::idle().message()?); + } context.stream.send_many(&messages).await? }; @@ -156,4 +159,22 @@ mod tests { context.transaction ); } + + #[tokio::test] + async fn test_end_not_connected_suppress_rfq() { + load_test(); + + let mut client = + crate::frontend::Client::new_test(Stream::dev_null(), Parameters::default()); + client.transaction = Some(TransactionType::ReadWrite); + + let mut engine = QueryEngine::from_client(&client).unwrap(); + let mut context = QueryEngineContext::new(&mut client); + context.simple_query_splice = true; + context.requests_left = 1; + + let result = engine.end_not_connected(&mut context, false, false).await; + assert!(result.is_ok()); + assert_eq!(context.transaction, None); + } } diff --git a/pgdog/src/frontend/client/query_engine/fake.rs b/pgdog/src/frontend/client/query_engine/fake.rs index 02feb772b..c3b7ebd51 100644 --- a/pgdog/src/frontend/client/query_engine/fake.rs +++ b/pgdog/src/frontend/client/query_engine/fake.rs @@ -61,16 +61,22 @@ impl QueryEngine { .await? } ProtocolMessage::Query(_) => { - (if let Some(row) = data_row.as_ref() { + let suppress_rfq = + context.simple_query_splice && context.requests_left > 0; + let cc = (if let Some(row) = data_row.as_ref() { context.stream.send(&row_description).await? + context.stream.send(row).await? } else { 0 - }) + context.stream.send(&CommandComplete::new(command)).await? - + context + }) + context.stream.send(&CommandComplete::new(command)).await?; + if suppress_rfq { + cc + } else { + cc + context .stream .send(&ReadyForQuery::in_transaction(context.in_transaction())) .await? + } } // TODO(lev): Elixir closes the statement it just asked us to prepare. // That's very memory-conscious of it, and we appreciate it. diff --git a/pgdog/src/frontend/client/query_engine/query.rs b/pgdog/src/frontend/client/query_engine/query.rs index 5156ad891..a96d82c51 100644 --- a/pgdog/src/frontend/client/query_engine/query.rs +++ b/pgdog/src/frontend/client/query_engine/query.rs @@ -243,10 +243,16 @@ impl QueryEngine { trace!("{:#?} >>> {:?}", message, context.stream.peer_addr()); - if flush { - context.stream.send_flush(&message).await?; - } else { - context.stream.send(&message).await?; + // For spliced simple queries the client sent one Q and expects one ReadyForQuery. + // Suppress intermediate ones; only the final statement's ReadyForQuery is forwarded. + let suppress_rfq = code == 'Z' && context.simple_query_splice && context.requests_left > 0; + + if !suppress_rfq { + if flush { + context.stream.send_flush(&message).await?; + } else { + context.stream.send(&message).await?; + } } if code == 'Z' { diff --git a/pgdog/src/frontend/client/query_engine/start_transaction.rs b/pgdog/src/frontend/client/query_engine/start_transaction.rs index b448897e3..5dcd584e7 100644 --- a/pgdog/src/frontend/client/query_engine/start_transaction.rs +++ b/pgdog/src/frontend/client/query_engine/start_transaction.rs @@ -26,13 +26,15 @@ impl QueryEngine { self.extended_transaction_reply(context, true, false) .await? } else { - context - .stream - .send_many(&[ - CommandComplete::new_begin().message()?, + let suppress_rfq = + context.simple_query_splice && context.requests_left > 0; + let mut msgs = vec![CommandComplete::new_begin().message()?]; + if !suppress_rfq { + msgs.push( ReadyForQuery::in_transaction(context.in_transaction()).message()?, - ]) - .await? + ); + } + context.stream.send_many(&msgs).await? }; self.stats.sent(bytes_sent); diff --git a/pgdog/src/frontend/client_request.rs b/pgdog/src/frontend/client_request.rs index be49a1d3a..e0fa83a00 100644 --- a/pgdog/src/frontend/client_request.rs +++ b/pgdog/src/frontend/client_request.rs @@ -7,6 +7,9 @@ use std::ops::{Deref, DerefMut}; use lazy_static::lazy_static; use regex::Regex; +#[cfg(feature = "new_parser")] +use pg_raw_parse; + use crate::{ frontend::router::Ast, net::{ @@ -15,11 +18,38 @@ use crate::{ }, stats::memory::MemoryUsage, }; +#[cfg(feature = "new_parser")] +use crate::net::messages::Query; use super::{PreparedStatements, router::Route}; pub use super::BufferedQuery; +/// Split a multi-statement simple query string into individual statement slices. +/// Returns an empty vec if parsing fails or the query contains only one statement. +#[cfg(feature = "new_parser")] +fn split_statements(query: &str) -> Vec<&str> { + let Ok(parsed) = pg_raw_parse::parse(query) else { + return vec![]; + }; + let stmts = parsed.into_inner(); + if stmts.len() <= 1 { + return vec![]; + } + stmts + .iter() + .filter_map(|stmt| { + let loc = stmt.stmt_location as usize; + let end = if stmt.stmt_len == 0 { + query.len() + } else { + loc + stmt.stmt_len as usize + }; + query.get(loc..end).map(str::trim).filter(|s| !s.is_empty()) + }) + .collect() +} + /// Client request, containing exactly one query. #[derive(Debug, Clone)] pub struct ClientRequest { @@ -369,6 +399,33 @@ impl ClientRequest { Ok(requests) } + + /// Split a multi-statement simple Query into individual single-statement requests. + /// + /// Analogous to `spliced()` for extended protocol. Returns empty vec when + /// the request is not a simple Query, contains only one statement, or parsing fails. + /// In all those cases the caller falls through to the normal query engine path. + #[cfg(feature = "new_parser")] + pub fn spliced_simple(&self) -> Vec { + let Some(ProtocolMessage::Query(q)) = self.messages.iter().find(|m| m.code() == 'Q') + else { + return vec![]; + }; + split_statements(q.query()) + .into_iter() + .map(|text| Self { + messages: vec![ProtocolMessage::Query(Query::new(text))], + route: None, + ast: None, + last_parse: None, + }) + .collect() + } + + #[cfg(not(feature = "new_parser"))] + pub fn spliced_simple(&self) -> Vec { + vec![] + } } impl From for Vec { @@ -408,6 +465,76 @@ mod test { use super::*; + #[cfg(not(feature = "new_parser"))] + #[test] + fn test_spliced_simple_returns_empty_without_new_parser() { + let req = ClientRequest::from(vec![Query::new("SELECT 1; SELECT 2").into()]); + assert!(req.spliced_simple().is_empty()); + } + + #[cfg(feature = "new_parser")] + #[test] + fn test_spliced_simple_two_selects() { + let req = ClientRequest::from(vec![Query::new("SELECT 1; SELECT 2").into()]); + let split = req.spliced_simple(); + assert_eq!(split.len(), 2, "expected 2 split requests"); + let texts: Vec<&str> = split + .iter() + .map(|r| match r.messages.first().unwrap() { + ProtocolMessage::Query(q) => q.query(), + _ => panic!("expected Query"), + }) + .collect(); + assert!(texts[0].contains("SELECT 1"), "first: {}", texts[0]); + assert!(texts[1].contains("SELECT 2"), "second: {}", texts[1]); + } + + #[cfg(feature = "new_parser")] + #[test] + fn test_spliced_simple_single_statement_no_split() { + let req = ClientRequest::from(vec![Query::new("SELECT 1").into()]); + assert!(req.spliced_simple().is_empty()); + } + + #[cfg(feature = "new_parser")] + #[test] + fn test_spliced_simple_non_query_no_split() { + let req = ClientRequest::from(vec![ + Parse::named("s", "SELECT $1").into(), + Bind::new_statement("s").into(), + Execute::new().into(), + Sync::new().into(), + ]); + assert!(req.spliced_simple().is_empty()); + } + + #[cfg(feature = "new_parser")] + #[test] + fn test_spliced_simple_transaction_batch() { + let req = ClientRequest::from(vec![ + Query::new( + "BEGIN;\ + DELETE FROM locks WHERE ttl < CURRENT_TIMESTAMP AT TIME ZONE 'UTC';\ + INSERT INTO locks (key, owner, ttl) VALUES ('k', 'o', NOW() + INTERVAL '1 min') ON CONFLICT DO NOTHING;\ + COMMIT;" + ).into(), + ]); + let split = req.spliced_simple(); + assert_eq!( + split.len(), + 4, + "expected BEGIN/DELETE/INSERT/COMMIT as 4 requests" + ); + let codes: Vec = split + .iter() + .map(|r| r.messages.first().unwrap().code()) + .collect(); + assert!( + codes.iter().all(|&c| c == 'Q'), + "all should be simple Query messages" + ); + } + #[test] fn test_request_splice() { let messages = vec![ diff --git a/pgdog/src/frontend/router/parser/error.rs b/pgdog/src/frontend/router/parser/error.rs index e9f4bb80a..88947899f 100644 --- a/pgdog/src/frontend/router/parser/error.rs +++ b/pgdog/src/frontend/router/parser/error.rs @@ -101,7 +101,4 @@ pub enum Error { #[error("sharded databases require the query parser to be enabled")] QueryParserRequired, - - #[error("multi-statement queries cannot mix SET with other commands")] - MultiStatementMixedSet, } diff --git a/pgdog/src/frontend/router/parser/query/mod.rs b/pgdog/src/frontend/router/parser/query/mod.rs index 587ec8d5c..beffbc310 100644 --- a/pgdog/src/frontend/router/parser/query/mod.rs +++ b/pgdog/src/frontend/router/parser/query/mod.rs @@ -290,19 +290,18 @@ impl QueryParser { .run()?; } - // Handle multi-statement SET commands (e.g. "SET x TO 1; SET y TO 2"). - if stmts.len() > 1 - && let Some(command) = self.try_multi_set(&**stmts, context)? - { - return Ok(command); + // Handle multi-statement queries. + if stmts.len() > 1 { + // All-SET batch: intercepted for parameter tracking. + if let Some(command) = self.try_multi_set(&**stmts, context)? { + return Ok(command); + } + // Mixed or non-SET batch: forward the full query to the write primary. + return Ok(Command::Query(Route::write( + context.shards_calculator.shard(), + ))); } - // - // Get the root AST node. - // - // We don't expect clients to send multiple queries. If they do - // only the first one is used for routing. - // let root = stmts.first(); let Some(root) = root else { @@ -555,19 +554,18 @@ impl QueryParser { let stmts = &statement.parse_result().protobuf.stmts; - // Handle multi-statement SET commands (e.g. "SET x TO 1; SET y TO 2"). - if stmts.len() > 1 - && let Some(command) = self.try_multi_set(stmts, context)? - { - return Ok(command); + // Handle multi-statement queries. + if stmts.len() > 1 { + // All-SET batch: intercepted for parameter tracking. + if let Some(command) = self.try_multi_set(stmts, context)? { + return Ok(command); + } + // Mixed or non-SET batch: forward the full query to the write primary. + return Ok(Command::Query(Route::write( + context.shards_calculator.shard(), + ))); } - // - // Get the root AST node. - // - // We don't expect clients to send multiple queries. If they do - // only the first one is used for routing. - // let root = stmts.first(); let root = if let Some(root) = root { diff --git a/pgdog/src/frontend/router/parser/query/set.rs b/pgdog/src/frontend/router/parser/query/set.rs index 95ed30bea..34f97f7df 100644 --- a/pgdog/src/frontend/router/parser/query/set.rs +++ b/pgdog/src/frontend/router/parser/query/set.rs @@ -118,8 +118,7 @@ impl QueryParser { /// Try to handle multi-statement queries containing SET commands. /// /// - All SETs → returns `Ok(Some(Command::Set { .. }))` - /// - No SETs → returns `Ok(None)`, caller falls through to default parsing - /// - Mix of SET + non-SET → returns `Err(MultiStatementMixedSet)` + /// - No SETs or mixed SET + non-SET → returns `Ok(None)`, caller routes to write primary /// /// In session mode, returns `Ok(Some(Command::Query(..)))` immediately so that /// all multi-statement queries are forwarded to the server verbatim. @@ -151,10 +150,8 @@ impl QueryParser { }) .collect::, _>>()?; - if params.is_empty() { + if params.is_empty() || has_other { Ok(None) - } else if has_other { - Err(Error::MultiStatementMixedSet) } else { Ok(Some(Command::Set { params, @@ -204,7 +201,7 @@ impl QueryParser { } if has_set && has_other { - return Err(Error::MultiStatementMixedSet); + return Ok(None); } } diff --git a/pgdog/src/frontend/router/parser/query/test/test_set.rs b/pgdog/src/frontend/router/parser/query/test/test_set.rs index d8094e4cc..b4cc206dd 100644 --- a/pgdog/src/frontend/router/parser/query/test/test_set.rs +++ b/pgdog/src/frontend/router/parser/query/test/test_set.rs @@ -29,15 +29,15 @@ fn test_mixed_set_passthrough_in_session_mode() { } #[test] -fn test_mixed_set_rejected_in_transaction_mode() { +fn test_mixed_set_routes_to_write_in_transaction_mode() { let mut test = QueryParserTest::new(); - let result = test.try_execute(vec![ + let command = test.execute(vec![ Query::new("SET DateStyle='ISO'; show transaction_isolation").into(), ]); assert!( - result.is_err(), - "expected error for mixed SET in transaction mode, got {result:#?}", + matches!(command, Command::Query(ref r) if r.is_write()), + "expected write Command::Query for mixed SET in transaction mode, got {command:#?}", ); } @@ -121,23 +121,26 @@ fn test_set_multi_statement_mixed_local() { } #[test] -fn test_set_multi_statement_mixed_returns_error() { +fn test_set_multi_statement_mixed_routes_to_write() { let mut test = QueryParserTest::new(); - let result = test.try_execute(vec![ + let command = test.execute(vec![ Query::new("SET statement_timeout TO 1; SELECT 1").into(), ]); - assert!(result.is_err()); + assert!( + matches!(command, Command::Query(ref r) if r.is_write()), + "expected write Command::Query for mixed SET+SELECT, got {command:#?}", + ); } #[test] -fn test_multi_statement_no_set_falls_through() { +fn test_multi_statement_no_set_routes_to_write() { let mut test = QueryParserTest::new(); let command = test.execute(vec![Query::new("SELECT 1; SELECT 2").into()]); assert!( - matches!(command, Command::Query(_)), - "multi-statement without SET should fall through, got {command:#?}", + matches!(command, Command::Query(ref r) if r.is_write()), + "multi-statement should route to write primary, got {command:#?}", ); }