From 8ec9de8436651df083f6d609007218479e4107db Mon Sep 17 00:00:00 2001 From: Haris Amin Date: Mon, 27 Jul 2026 12:48:17 -0400 Subject: [PATCH 01/10] feat: route multi-statement simple queries to write primary Previously, batches mixing SET with other statement types raised an error (MultiStatementMixedSet), and batches of non-SET statements were routed using only the first statement. Both cases now route the full batch to the write primary, letting the backend execute all statements correctly. Closes #395 --- pgdog/src/frontend/router/parser/error.rs | 3 -- pgdog/src/frontend/router/parser/query/mod.rs | 42 +++++++++---------- pgdog/src/frontend/router/parser/query/set.rs | 9 ++-- .../router/parser/query/test/test_set.rs | 23 +++++----- 4 files changed, 36 insertions(+), 41 deletions(-) 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:#?}", ); } From 28ac651bec6f3b0af23daa7d0d5e3721c0dbf1c1 Mon Sep 17 00:00:00 2001 From: Haris Amin Date: Mon, 27 Jul 2026 13:17:47 -0400 Subject: [PATCH 02/10] test: update integration tests for multi-statement query support Replace tests that expected MultiStatementMixedSet errors with tests that verify multi-statement batches succeed. Add test_multi_statement_select to confirm both result sets are returned for 'SELECT 1; SELECT 2'. --- .../rust/tests/integration/multi_set.rs | 34 +++++++++++++------ integration/rust/tests/sqlx/multi_set.rs | 17 +++------- 2 files changed, 29 insertions(+), 22 deletions(-) diff --git a/integration/rust/tests/integration/multi_set.rs b/integration/rust/tests/integration/multi_set.rs index b3f53af35..b9000e1e4 100644 --- a/integration/rust/tests/integration/multi_set.rs +++ b/integration/rust/tests/integration/multi_set.rs @@ -46,17 +46,31 @@ 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 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(); + assert_eq!(values, vec!["1", "2"]); } } 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"); } } From 941c9acdb10e44e7a980f54e5d25072cf3952ab4 Mon Sep 17 00:00:00 2001 From: Haris Amin Date: Tue, 28 Jul 2026 13:46:19 -0400 Subject: [PATCH 03/10] feat: add spliced_simple() to split multi-statement simple queries Uses stmt_location/stmt_len from the SQL parser (same fields used by pg_dump.rs) to extract individual statement substrings and wrap each in its own ClientRequest. Returns empty vec for single-statement or non-Query messages so callers fall through unchanged. --- pgdog/src/frontend/client/mod.rs | 5 + pgdog/src/frontend/client_request.rs | 137 ++++++++++++++++++++++++++- 2 files changed, 141 insertions(+), 1 deletion(-) diff --git a/pgdog/src/frontend/client/mod.rs b/pgdog/src/frontend/client/mod.rs index 56f9004ed..77da9c455 100644 --- a/pgdog/src/frontend/client/mod.rs +++ b/pgdog/src/frontend/client/mod.rs @@ -544,6 +544,11 @@ impl Client { // If client sent multiple requests, split them up and execute individually. let spliced = self.client_request.spliced()?; + let spliced = if spliced.is_empty() { + self.client_request.spliced_simple() + } else { + spliced + }; if spliced.is_empty() { let mut context = QueryEngineContext::new(self); query_engine.handle(&mut context).await?; diff --git a/pgdog/src/frontend/client_request.rs b/pgdog/src/frontend/client_request.rs index be49a1d3a..63fc5b7d5 100644 --- a/pgdog/src/frontend/client_request.rs +++ b/pgdog/src/frontend/client_request.rs @@ -7,11 +7,16 @@ use std::ops::{Deref, DerefMut}; use lazy_static::lazy_static; use regex::Regex; +#[cfg(not(feature = "new_parser"))] +use pg_query; +#[cfg(feature = "new_parser")] +use pg_raw_parse; + use crate::{ frontend::router::Ast, net::{ Error, Flush, Parse, ProtocolMessage, - messages::{Bind, CopyData, Protocol}, + messages::{Bind, CopyData, Protocol, Query}, }, stats::memory::MemoryUsage, }; @@ -20,6 +25,44 @@ use super::{PreparedStatements, router::Route}; pub use super::BufferedQuery; +#[cfg(feature = "new_parser")] +fn split_statements(query: &str) -> Vec<(usize, usize)> { + let Ok(parsed) = pg_raw_parse::parse(query) else { + return vec![]; + }; + let inner = parsed.into_inner(); + if inner.len() <= 1 { + return vec![]; + } + inner + .into_iter() + .map(|stmt| { + let loc = stmt.stmt_location.max(0) as usize; + let len = stmt.stmt_len as usize; + (loc, len) + }) + .collect() +} + +#[cfg(not(feature = "new_parser"))] +fn split_statements(query: &str) -> Vec<(usize, usize)> { + let Ok(parsed) = pg_query::parse(query) else { + return vec![]; + }; + let stmts = &parsed.protobuf.stmts; + if stmts.len() <= 1 { + return vec![]; + } + stmts + .iter() + .map(|stmt| { + let loc = stmt.stmt_location.max(0) as usize; + let len = stmt.stmt_len as usize; + (loc, len) + }) + .collect() +} + /// Client request, containing exactly one query. #[derive(Debug, Clone)] pub struct ClientRequest { @@ -369,6 +412,39 @@ 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. + pub fn spliced_simple(&self) -> Vec { + let Some(ProtocolMessage::Query(q)) = self.messages.iter().find(|m| m.code() == 'Q') else { + return vec![]; + }; + let query_text = q.query(); + + split_statements(query_text) + .into_iter() + .filter_map(|(offset, len)| { + let end = if len == 0 { + query_text.len() + } else { + offset + len + }; + let text = query_text.get(offset..end)?.trim(); + if text.is_empty() { + return None; + } + Some(Self { + messages: vec![ProtocolMessage::Query(Query::new(text))], + route: None, + ast: None, + last_parse: None, + }) + }) + .collect() + } } impl From for Vec { @@ -408,6 +484,65 @@ mod test { use super::*; + #[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]); + } + + #[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()); + } + + #[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()); + } + + #[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![ From 3bb00cea5551e670578349b513826cd090816684 Mon Sep 17 00:00:00 2001 From: Haris Amin Date: Tue, 28 Jul 2026 14:36:21 -0400 Subject: [PATCH 04/10] test: add multi-statement transaction batch integration test Covers a BEGIN/DELETE/INSERT/COMMIT batch sent as a single simple query message, verifying all four statements execute and the inserted row is visible after commit. --- .../rust/tests/integration/multi_set.rs | 33 +++++++++++++++++++ 1 file changed, 33 insertions(+) diff --git a/integration/rust/tests/integration/multi_set.rs b/integration/rust/tests/integration/multi_set.rs index b9000e1e4..2f8cb3657 100644 --- a/integration/rust/tests/integration/multi_set.rs +++ b/integration/rust/tests/integration/multi_set.rs @@ -74,3 +74,36 @@ async fn test_multi_statement_select() { assert_eq!(values, vec!["1", "2"]); } } + +#[tokio::test] +async fn test_multi_statement_transaction_batch() { + for conn in connections_tokio().await { + conn.batch_execute( + "CREATE TEMP TABLE IF NOT EXISTS locks ( + key TEXT PRIMARY KEY, + owner TEXT NOT NULL, + ttl TIMESTAMPTZ NOT NULL + )", + ) + .await + .unwrap(); + + conn.batch_execute( + "BEGIN;\ + DELETE FROM locks WHERE ttl < CURRENT_TIMESTAMP AT TIME ZONE 'UTC';\ + INSERT INTO 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 locks WHERE key = 'test-key'") + .await + .unwrap(); + let owner = extract_simple_query_value(&rows); + assert_eq!(owner, "test-owner"); + } +} From ea56cb6b3c40050a5311fdec54ac0c7c0a4df8b3 Mon Sep 17 00:00:00 2001 From: Haris Amin Date: Wed, 29 Jul 2026 10:09:10 -0400 Subject: [PATCH 05/10] fix: remove new_parser feature flag from statement splitting Replace the pg_raw_parse-based (feature-gated) split_statements with a pg_query-based version that works in default builds. The default parser exposes the same stmt_location/stmt_len fields, so no behaviour changes. Also simplifies split_statements to return &str slices directly instead of (offset, len) pairs, removing duplicate boundary-calculation logic in spliced_simple. --- pgdog/src/frontend/client_request.rs | 69 ++++++++-------------------- 1 file changed, 19 insertions(+), 50 deletions(-) diff --git a/pgdog/src/frontend/client_request.rs b/pgdog/src/frontend/client_request.rs index 63fc5b7d5..c84153b55 100644 --- a/pgdog/src/frontend/client_request.rs +++ b/pgdog/src/frontend/client_request.rs @@ -7,11 +7,6 @@ use std::ops::{Deref, DerefMut}; use lazy_static::lazy_static; use regex::Regex; -#[cfg(not(feature = "new_parser"))] -use pg_query; -#[cfg(feature = "new_parser")] -use pg_raw_parse; - use crate::{ frontend::router::Ast, net::{ @@ -25,40 +20,26 @@ use super::{PreparedStatements, router::Route}; pub use super::BufferedQuery; -#[cfg(feature = "new_parser")] -fn split_statements(query: &str) -> Vec<(usize, usize)> { - let Ok(parsed) = pg_raw_parse::parse(query) else { - return vec![]; - }; - let inner = parsed.into_inner(); - if inner.len() <= 1 { - return vec![]; - } - inner - .into_iter() - .map(|stmt| { - let loc = stmt.stmt_location.max(0) as usize; - let len = stmt.stmt_len as usize; - (loc, len) - }) - .collect() -} - -#[cfg(not(feature = "new_parser"))] -fn split_statements(query: &str) -> Vec<(usize, usize)> { +/// 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. +fn split_statements(query: &str) -> Vec<&str> { let Ok(parsed) = pg_query::parse(query) else { return vec![]; }; - let stmts = &parsed.protobuf.stmts; + let stmts = parsed.protobuf.stmts; if stmts.len() <= 1 { return vec![]; } stmts .iter() - .map(|stmt| { + .filter_map(|stmt| { let loc = stmt.stmt_location.max(0) as usize; - let len = stmt.stmt_len as usize; - (loc, len) + 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() } @@ -419,29 +400,17 @@ impl ClientRequest { /// 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. pub fn spliced_simple(&self) -> Vec { - let Some(ProtocolMessage::Query(q)) = self.messages.iter().find(|m| m.code() == 'Q') else { + let Some(ProtocolMessage::Query(q)) = self.messages.iter().find(|m| m.code() == 'Q') + else { return vec![]; }; - let query_text = q.query(); - - split_statements(query_text) + split_statements(q.query()) .into_iter() - .filter_map(|(offset, len)| { - let end = if len == 0 { - query_text.len() - } else { - offset + len - }; - let text = query_text.get(offset..end)?.trim(); - if text.is_empty() { - return None; - } - Some(Self { - messages: vec![ProtocolMessage::Query(Query::new(text))], - route: None, - ast: None, - last_parse: None, - }) + .map(|text| Self { + messages: vec![ProtocolMessage::Query(Query::new(text))], + route: None, + ast: None, + last_parse: None, }) .collect() } From d62a54818e461f2dd3660db2c96e62666a071fce Mon Sep 17 00:00:00 2001 From: Haris Amin Date: Wed, 29 Jul 2026 10:09:19 -0400 Subject: [PATCH 06/10] fix: suppress ReadyForQuery for intermediate statements in simple query splice MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When a multi-statement simple query is split into individual requests, all but the last must hold back the ReadyForQuery(Z) message so the client sees exactly one Z at the end of the batch — matching standard PostgreSQL simple query protocol semantics. - QueryEngineContext gains simple_query_splice flag and requests_left counter set by spliced_simple_query() builder - query.rs: suppress Z forwarding from the backend while requests remain - fake.rs: suppress synthetic Z in fake_command_response (handles SET) - start_transaction.rs: suppress synthetic Z for BEGIN - end_transaction.rs: suppress synthetic Z for COMMIT/ROLLBACK - mod.rs: track is_simple_splice to skip implicit-transaction injection that would force incorrect shard routing for split simple queries --- pgdog/src/frontend/client/mod.rs | 25 ++++++++++++++----- .../frontend/client/query_engine/context.rs | 16 ++++++++++++ .../client/query_engine/end_transaction.rs | 5 +++- .../src/frontend/client/query_engine/fake.rs | 12 ++++++--- .../src/frontend/client/query_engine/query.rs | 14 ++++++++--- .../client/query_engine/start_transaction.rs | 14 ++++++----- 6 files changed, 66 insertions(+), 20 deletions(-) diff --git a/pgdog/src/frontend/client/mod.rs b/pgdog/src/frontend/client/mod.rs index 77da9c455..b3df19591 100644 --- a/pgdog/src/frontend/client/mod.rs +++ b/pgdog/src/frontend/client/mod.rs @@ -543,11 +543,13 @@ impl Client { } // If client sent multiple requests, split them up and execute individually. - let spliced = self.client_request.spliced()?; - let spliced = if spliced.is_empty() { - self.client_request.spliced_simple() + 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 { - spliced + (extended, false) }; if spliced.is_empty() { let mut context = QueryEngineContext::new(self); @@ -556,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..e78b06663 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? }; 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); From d766017a638c7d9b053e0503d052a61b906764e2 Mon Sep 17 00:00:00 2001 From: Haris Amin Date: Wed, 29 Jul 2026 10:09:27 -0400 Subject: [PATCH 07/10] test: expand multi-statement integration test suite Add tests covering: - SELECT 1; SELECT 2 returns results in order (with dedup for sharded connections that fan out to multiple shards) - BEGIN; DELETE; INSERT; COMMIT batch commits atomically and rows are visible after commit (uses a real table instead of TEMP to survive transaction-mode pool backend switching) - INSERT; INSERT batch with a PK conflict documents per-statement auto-commit behaviour: the first INSERT commits before the second fails Switch test tables from TEMP to persistent (CREATE TABLE IF NOT EXISTS + TRUNCATE) so transaction-mode pooling can route setup and batch queries to different backends without losing visibility. --- .../rust/tests/integration/multi_set.rs | 59 +++++++++++++++++-- 1 file changed, 54 insertions(+), 5 deletions(-) diff --git a/integration/rust/tests/integration/multi_set.rs b/integration/rust/tests/integration/multi_set.rs index 2f8cb3657..2496c5e34 100644 --- a/integration/rust/tests/integration/multi_set.rs +++ b/integration/rust/tests/integration/multi_set.rs @@ -61,7 +61,7 @@ async fn test_multi_set_mixed_succeeds() { async fn test_multi_statement_select() { for conn in connections_tokio().await { let msgs = conn.simple_query("SELECT 1; SELECT 2").await.unwrap(); - let values: Vec = msgs + let mut values: Vec = msgs .iter() .filter_map(|m| { if let tokio_postgres::SimpleQueryMessage::Row(row) = m { @@ -71,6 +71,10 @@ async fn test_multi_statement_select() { } }) .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"]); } } @@ -78,8 +82,11 @@ async fn test_multi_statement_select() { #[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 TEMP TABLE IF NOT EXISTS locks ( + "CREATE TABLE IF NOT EXISTS pgdog_multi_stmt_locks ( key TEXT PRIMARY KEY, owner TEXT NOT NULL, ttl TIMESTAMPTZ NOT NULL @@ -88,10 +95,14 @@ async fn test_multi_statement_transaction_batch() { .await .unwrap(); + conn.batch_execute("TRUNCATE pgdog_multi_stmt_locks") + .await + .unwrap(); + conn.batch_execute( "BEGIN;\ - DELETE FROM locks WHERE ttl < CURRENT_TIMESTAMP AT TIME ZONE 'UTC';\ - INSERT INTO locks (key, owner, ttl) \ + 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;", @@ -100,10 +111,48 @@ async fn test_multi_statement_transaction_batch() { .unwrap(); let rows = conn - .simple_query("SELECT owner FROM locks WHERE key = 'test-key'") + .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"); + } +} From 6a98585d75b9afe7dae62e481507ff8528bb6b12 Mon Sep 17 00:00:00 2001 From: Haris Amin Date: Wed, 29 Jul 2026 13:38:29 -0400 Subject: [PATCH 08/10] fix: gate spliced_simple behind new_parser feature flag Use pg_raw_parse (new parser) for statement splitting and guard split_statements, spliced_simple, and related unit tests behind #[cfg(feature = "new_parser")]. The no-feature path returns vec![] so callers fall through to the existing single-statement path unchanged. --- pgdog/src/frontend/client_request.rs | 24 ++++++++++++++++++++---- 1 file changed, 20 insertions(+), 4 deletions(-) diff --git a/pgdog/src/frontend/client_request.rs b/pgdog/src/frontend/client_request.rs index c84153b55..dd6b7b20c 100644 --- a/pgdog/src/frontend/client_request.rs +++ b/pgdog/src/frontend/client_request.rs @@ -7,14 +7,19 @@ 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::{ Error, Flush, Parse, ProtocolMessage, - messages::{Bind, CopyData, Protocol, Query}, + messages::{Bind, CopyData, Protocol}, }, stats::memory::MemoryUsage, }; +#[cfg(feature = "new_parser")] +use crate::net::messages::Query; use super::{PreparedStatements, router::Route}; @@ -22,18 +27,19 @@ 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_query::parse(query) else { + let Ok(parsed) = pg_raw_parse::parse(query) else { return vec![]; }; - let stmts = parsed.protobuf.stmts; + let stmts = parsed.into_inner(); if stmts.len() <= 1 { return vec![]; } stmts .iter() .filter_map(|stmt| { - let loc = stmt.stmt_location.max(0) as usize; + let loc = stmt.stmt_location as usize; let end = if stmt.stmt_len == 0 { query.len() } else { @@ -399,6 +405,7 @@ impl ClientRequest { /// 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 { @@ -414,6 +421,11 @@ impl ClientRequest { }) .collect() } + + #[cfg(not(feature = "new_parser"))] + pub fn spliced_simple(&self) -> Vec { + vec![] + } } impl From for Vec { @@ -453,6 +465,7 @@ mod test { use super::*; + #[cfg(feature = "new_parser")] #[test] fn test_spliced_simple_two_selects() { let req = ClientRequest::from(vec![Query::new("SELECT 1; SELECT 2").into()]); @@ -469,12 +482,14 @@ mod test { 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![ @@ -486,6 +501,7 @@ mod test { assert!(req.spliced_simple().is_empty()); } + #[cfg(feature = "new_parser")] #[test] fn test_spliced_simple_transaction_batch() { let req = ClientRequest::from(vec![ From 5737fd2a92224300991b04fb8d4289b4455212f0 Mon Sep 17 00:00:00 2001 From: Haris Amin Date: Wed, 29 Jul 2026 13:49:25 -0400 Subject: [PATCH 09/10] test: cover suppress_rfq branch in end_not_connected and no-new_parser stub Add test_end_not_connected_suppress_rfq to exercise the branch where ReadyForQuery is suppressed for intermediate simple query splice statements. Add test_spliced_simple_returns_empty_without_new_parser to cover the cfg(not(new_parser)) spliced_simple stub. --- .../client/query_engine/end_transaction.rs | 18 ++++++++++++++++++ pgdog/src/frontend/client_request.rs | 7 +++++++ 2 files changed, 25 insertions(+) diff --git a/pgdog/src/frontend/client/query_engine/end_transaction.rs b/pgdog/src/frontend/client/query_engine/end_transaction.rs index e78b06663..2b729b536 100644 --- a/pgdog/src/frontend/client/query_engine/end_transaction.rs +++ b/pgdog/src/frontend/client/query_engine/end_transaction.rs @@ -159,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_request.rs b/pgdog/src/frontend/client_request.rs index dd6b7b20c..e0fa83a00 100644 --- a/pgdog/src/frontend/client_request.rs +++ b/pgdog/src/frontend/client_request.rs @@ -465,6 +465,13 @@ 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() { From ead6ed31911c5f55d12d44c14e44d05f13579d79 Mon Sep 17 00:00:00 2001 From: Haris Amin Date: Wed, 29 Jul 2026 15:24:00 -0400 Subject: [PATCH 10/10] test(js): update mixed SET+SELECT test to expect success MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The `MultiStatementMixedSet` error was removed — mixed batches now route to the write primary rather than error. Update the Sequelize test to assert the query succeeds instead of asserting the old error message. --- integration/js/pg_tests/test/sequelize.js | 15 +++------------ 1 file changed, 3 insertions(+), 12 deletions(-) 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);