Feat/multi statement query support

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

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<String> = 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");
     }
 }

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

pgdog/src/frontend/client/mod.rs

@@ -543,18 +543,36 @@ 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?;
             self.transaction = context.transaction();
         } 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();
 

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<usize>,
+    /// 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,
         }
     }
 

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

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.

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

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