feat: partial multi-query support

integration/js/pg_tests/test/sequelize.js

@@ -200,22 +200,12 @@ 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 returns query result", 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}`,
-      );
-    }
-
-    const [rows] = await seq.query("SELECT 1 AS val");
+    const [rows] = await seq.query(
+      "SET statement_timeout TO '10s'; SELECT 1 AS val",
+    );
     assert.strictEqual(rows[0].val, 1);
 
     await seq.close();

integration/rust/tests/integration/multi_set.rs

@@ -46,17 +46,10 @@ async fn test_multi_set_with_timezone_interval() {
 }
 
 #[tokio::test]
-async fn test_multi_set_mixed_returns_error() {
+async fn test_multi_set_mixed_works() {
     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();
     }
 }

integration/rust/tests/sqlx/multi_set.rs

@@ -76,15 +76,9 @@ async fn test_multi_set_mixed_returns_error() {
     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")

pgdog/src/frontend/client/mod.rs

@@ -24,7 +24,6 @@ use crate::backend::{
 use crate::config::convert::user_from_params;
 use crate::config::{self, AuthType, ConfigAndUsers, config};
 use crate::frontend::ClientComms;
-use crate::frontend::client::query_engine::{QueryEngine, QueryEngineContext, QueryEngineResult};
 use crate::net::messages::{
     Authentication, BackendKeyData, ErrorResponse, FromBytes, FrontendPid, Message, Password,
     Protocol, ProtocolVersion, ReadyForQuery, ToBytes,
@@ -39,6 +38,7 @@ pub mod sticky;
 pub mod timeouts;
 pub mod transaction_type;
 
+use query_engine::QueryEngine;
 pub(crate) use sticky::Sticky;
 pub use transaction_type::TransactionType;
 
@@ -557,6 +557,8 @@ impl Client {
         query_engine: &mut QueryEngine,
         message: Message,
     ) -> Result<(), Error> {
+        use query_engine::QueryEngineContext;
+
         let mut context = QueryEngineContext::new(self);
         query_engine
             .process_server_message(&mut context, message)
@@ -587,22 +589,25 @@ impl Client {
 
     /// Handle client messages.
     async fn client_messages(&mut self, query_engine: &mut QueryEngine) -> Result<(), Error> {
+        use query_engine::{Pipeline, QueryEngineContext, QueryEngineResult};
         self.check_maintenance_mode(query_engine).await;
 
         match query_engine
             .handle(&mut QueryEngineContext::new(self))
             .await?
         {
             QueryEngineResult::Done(transaction) => self.transaction = transaction,
-            QueryEngineResult::Split(requests) => {
+            QueryEngineResult::Split { requests, extended } => {
                 let mut requests = requests.into_iter();
-                self.transaction.get_or_insert(TransactionType::Implicit);
+                if extended {
+                    self.transaction.get_or_insert(TransactionType::Implicit);
+                }
 
                 while let Some(mut request) = requests.next() {
                     match query_engine
                         .handle(
                             &mut QueryEngineContext::new(self)
-                                .extended_pipeline(&mut request, requests.len()),
+                                .pipelined(&mut request, Pipeline::new(requests.len(), extended)),
                         )
                         .await?
                     {

pgdog/src/frontend/client/query_engine/context.rs

@@ -8,6 +8,8 @@ use crate::{
     net::{FrontendPid, Parameters, Stream},
 };
 
+use super::split::Pipeline;
+
 /// Context passed to the query engine to execute a query.
 pub struct QueryEngineContext<'a> {
     /// Client ID running the query.
@@ -19,7 +21,7 @@ pub struct QueryEngineContext<'a> {
     /// Request.
     pub(super) client_request: &'a mut ClientRequest,
     /// How many requests are left to execute in an extended pipeline.
-    pub(super) extended_pipeline_requests_left: usize,
+    pub(super) pipeline: Pipeline,
     /// Client's socket to send responses to.
     pub(super) stream: &'a mut Stream,
     /// Client in transaction?
@@ -59,7 +61,7 @@ impl<'a> QueryEngineContext<'a> {
             cross_shard_disabled: None,
             memory_stats,
             admin: client.admin,
-            extended_pipeline_requests_left: 0,
+            pipeline: Pipeline::None,
             rollback: false,
             sticky: client.sticky,
             rewrite_result: None,
@@ -70,13 +72,9 @@ impl<'a> QueryEngineContext<'a> {
 
     /// The request is an extended protocol pipeline
     /// with a counter of how many requests are left to process.
-    pub(crate) fn extended_pipeline(
-        mut self,
-        req: &'a mut ClientRequest,
-        request_left: usize,
-    ) -> Self {
+    pub(crate) fn pipelined(mut self, req: &'a mut ClientRequest, pipeline: Pipeline) -> Self {
         self.client_request = req;
-        self.extended_pipeline_requests_left = request_left;
+        self.pipeline = pipeline;
         self
     }
 
@@ -93,7 +91,7 @@ impl<'a> QueryEngineContext<'a> {
             cross_shard_disabled: None,
             memory_stats: MemoryStats::default(),
             admin: false,
-            extended_pipeline_requests_left: 0,
+            pipeline: Pipeline::None,
             rollback: false,
             sticky: Sticky::new(),
             rewrite_result: None,

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());
+
+            if context.pipeline.is_done() || !context.pipeline.is_simple() {
+                messages.push(ReadyForQuery::idle().message());
+            }
 
             context.stream.send_many(&messages).await?
         };

pgdog/src/frontend/client/query_engine/fake.rs

@@ -84,10 +84,16 @@ impl QueryEngine {
                     } else {
                         0
                     }) + context.stream.send(&CommandComplete::new(command)).await?
-                        + context
-                            .stream
-                            .send(&ReadyForQuery::in_transaction(context.in_transaction()))
-                            .await?
+                        + if context.pipeline.is_simple() && !context.pipeline.is_done() {
+                            // Don't send ReadyForQuery for intermediate queries in a simple query
+                            // pipeline.
+                            0
+                        } else {
+                            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/mod.rs

@@ -47,6 +47,7 @@ pub(crate) use advisory_lock::AdvisoryLocks;
 pub use context::QueryEngineContext;
 use notify_buffer::NotifyBuffer;
 pub(crate) use result::QueryEngineResult;
+pub(crate) use split::Pipeline;
 use two_pc::TwoPc;
 pub use two_pc::phase::TwoPcPhase;
 
@@ -119,15 +120,15 @@ impl QueryEngine {
         &mut self,
         context: &mut QueryEngineContext<'_>,
     ) -> Result<QueryEngineResult, Error> {
-        if let Some(result) = Self::check_extended_request_split(context.client_request)? {
+        if let Some(result) = Self::check_extended_pipeline_rewrite(context.client_request)? {
             return Ok(result);
         }
 
         self.stats
             .received(context.client_request.total_message_len());
         self.set_state(State::Active); // Client is active.
 
-        if self.extended_in_sync_check(context) {
+        if self.in_extended_pipeline_error(context) {
             return Ok(QueryEngineResult::Done(context.transaction()));
         }
 
@@ -267,14 +268,7 @@ impl QueryEngine {
             Command::Copy(_) => self.execute(context).await?,
             Command::Deallocate => self.deallocate(context).await?,
             Command::Discard { extended } => self.discard(context, *extended).await?,
-            Command::Split(_) => {
-                use crate::frontend::router::parser::Error as ParserError;
-                self.error_response(
-                    context,
-                    ErrorResponse::from_err(&Error::Parser(ParserError::MultiStatementMixedSet)),
-                )
-                .await?;
-            }
+            Command::Split(queries) => return Ok(Self::build_simple_split(queries)),
         }
 
         self.hooks.after_execution(context)?;

pgdog/src/frontend/client/query_engine/query.rs

@@ -29,6 +29,12 @@ impl QueryEngine {
             return Ok(());
         }
 
+        // Skip statements inside a simple pipeline
+        // if we are inside an errored-out transaction.
+        if self.in_simple_pipeline_error(context) {
+            return Ok(());
+        }
+
         // Check if we need to do 2pc automatically
         // for single-statement writes.
         self.two_pc_check(context);
@@ -238,12 +244,20 @@ impl QueryEngine {
         // Do this before flushing, because flushing can take time.
         self.cleanup_backend(context)?;
 
-        trace!("{:#?} >>> {:?}", message, context.stream.peer_addr());
-
-        if flush {
-            context.stream.send_flush(&message).await?;
-        } else {
-            context.stream.send(&message).await?;
+        // Pipelined requests only return
+        // one ReadyForQuery message.
+        let drop_message = message.code() == 'Z'
+            && !context.pipeline.is_done()
+            && context.pipeline.is_simple()
+            && !context.in_error(); // On error, pipeline is done executing.
+        if !drop_message {
+            trace!("{:#?} >>> {:?}", message, context.stream.peer_addr());
+
+            if flush {
+                context.stream.send_flush(&message).await?;
+            } else {
+                context.stream.send(&message).await?;
+            }
         }
 
         if code == 'Z' {
@@ -291,7 +305,7 @@ impl QueryEngine {
 
             // Release the connection back into the pool before flushing data to client.
             // Flushing can take a minute and we don't want to block the connection from being reused.
-            if !self.backend.session_mode() && context.extended_pipeline_requests_left == 0 {
+            if !self.backend.session_mode() && context.pipeline.is_done() {
                 self.backend.disconnect();
             }
 

pgdog/src/frontend/client/query_engine/result.rs

@@ -6,5 +6,8 @@ pub(crate) enum QueryEngineResult {
     Done(Option<TransactionType>),
     /// Query engine requests the request to be resubmitted
     /// as a series of separate requests.
-    Split(Vec<ClientRequest>),
+    Split {
+        requests: Vec<ClientRequest>,
+        extended: bool,
+    },
 }

pgdog/src/frontend/client/query_engine/split.rs

@@ -1,16 +1,58 @@
+use itertools::Itertools;
+
 use super::*;
-use crate::frontend::ClientRequest;
+use crate::{frontend::ClientRequest, net::Query};
+
+/// Query engine pipeline state.
+pub(crate) enum Pipeline {
+    Extended { requests_left: usize },
+    Simple { requests_left: usize },
+    None,
+}
+
+impl Pipeline {
+    /// Create new pipeline.
+    pub(crate) fn new(requests_left: usize, extended: bool) -> Self {
+        if extended {
+            Self::Extended { requests_left }
+        } else {
+            Self::Simple { requests_left }
+        }
+    }
+
+    /// How many requests left in the pipeline.
+    pub(super) fn requests_left(&self) -> usize {
+        match self {
+            Self::Extended { requests_left } => *requests_left,
+            Self::Simple { requests_left } => *requests_left,
+            Self::None => 0,
+        }
+    }
+
+    /// Is the pipeline finished executing?
+    pub(super) fn is_done(&self) -> bool {
+        self.requests_left() == 0
+    }
+
+    /// Is the pipeline consists of simple queries only?
+    pub(super) fn is_simple(&self) -> bool {
+        matches!(self, Self::Simple { .. })
+    }
+}
 
 impl QueryEngine {
     /// Check if the request needs splitting and perform the split if necessary.
     ///
     /// Caller is expected to abort the request and return the result back to the caller
     /// for resubmission.
-    pub(super) fn check_extended_request_split(
+    pub(super) fn check_extended_pipeline_rewrite(
         request: &ClientRequest,
     ) -> Result<Option<QueryEngineResult>, Error> {
         if request.is_multi_exec() {
-            Ok(Some(QueryEngineResult::Split(request.split_extended()?)))
+            Ok(Some(QueryEngineResult::Split {
+                requests: request.split_extended()?,
+                extended: true,
+            }))
         } else {
             Ok(None)
         }
@@ -24,9 +66,29 @@ impl QueryEngine {
     /// If we see a [`crate::net::Sync`]-only request, we execute it to restore servers
     /// back to normal state.
     ///
-    pub(super) fn extended_in_sync_check(&self, context: &QueryEngineContext<'_>) -> bool {
+    pub(super) fn in_extended_pipeline_error(&self, context: &QueryEngineContext<'_>) -> bool {
         self.backend.out_of_sync()
             && !context.client_request.is_sync_only()
-            && context.extended_pipeline_requests_left > 0
+            && !context.pipeline.is_done()
+    }
+
+    /// Return true if we should ignore this query because
+    /// the simple query pipeline is in an error state, i.e., inside a failed
+    /// transaction.
+    pub(super) fn in_simple_pipeline_error(&self, context: &QueryEngineContext<'_>) -> bool {
+        context.pipeline.is_simple() && context.in_error()
+    }
+
+    /// Build a multi-query split.
+    pub(super) fn build_simple_split(queries: &[String]) -> QueryEngineResult {
+        let requests = queries
+            .iter()
+            .map(|query| ClientRequest::from(vec![Query::new(query).into()]))
+            .collect_vec();
+
+        QueryEngineResult::Split {
+            requests,
+            extended: false,
+        }
     }
 }