Feat/multi statement query support by hamin · Pull Request #1275 · pgdogdev/pgdog

hamin · 2026-07-28T16:11:50Z

Context and motivation

At my company we're using PgDog as part of a PostgreSQL version upgrade migration strategy. We have a polyglot environment with services across different languages, ORMs, query builders, and driver versions. Some have hard dependencies we can't control (in our case, Kong) that generate raw multi-statement simple queries we couldn't inspect or modify. The goal was safety in the read path without having to audit every query generated across every language and library. We're not using sharding, so the intention was to keep changes minimal there.

What changed

Removed MultiStatementMixedSet — mixed batches no longer error
Router routes multi-statement non-SET batches to the write primary instead of only using the first statement for routing
Added ClientRequest::spliced_simple() behind #[cfg(feature = "new_parser")] that splits a Q message into per-statement requests using stmt_location/stmt_len from pg_raw_parse
Wired spliced_simple() into the client loop after spliced(), so simple-query batches go through the same splice processing path as extended-protocol pipelines when the new parser is enabled
Added simple_query_splice context flag and suppress_rfq logic in four places (query.rs, fake.rs, start_transaction.rs, end_transaction.rs) to hold back intermediate ReadyForQuery messages so the client sees exactly one Z at the end, matching simple query protocol semantics
All-SET batches continue to be intercepted by PgDog's parameter tracking; single-statement queries and parse failures fall through unchanged

Default builds (without new_parser): multi-statement batches route to the write primary as a whole. spliced_simple() returns empty vec and callers fall through.

new_parser builds: full per-statement splitting with SET interception and correct ReadyForQuery handling.

On the complexity

The suppress_rfq work across four handlers is more invasive than I originally wanted. It was necessary because SET, BEGIN, and COMMIT all have code paths that generate synthetic ReadyForQuery responses without going through the backend, so each one needed the same fix. The motivation for splitting rather than forwarding verbatim is that libraries we can't control frequently send batches like SET statement_timeout TO '30s'; SELECT ... as part of connection setup. If PgDog forwards the whole Q verbatim it never sees the SET, and when that backend connection gets returned to the pool the next client inherits stale parameter state silently.

If you feel this is too much complexity for what's being solved, or there are implications for the broader codebase I'm not seeing, happy to scale this back to just the routing fix.

On @sgrif's review comments

Old parser / new_parser feature flag: spliced_simple(), split_statements(), and the related unit tests are now all gated behind #[cfg(feature = "new_parser")] using pg_raw_parse. Default builds fall through unchanged.
&str -> Vec<&str> signature: done, split_statements returns slices directly now
Error propagation in split_statements: parse failures fall through to the single-statement path, which forwards to the backend and lets Postgres return the error to the client. Feels safer than surfacing a proxy-level parse error for something the backend might handle fine
Should router error if it sees more than one statement: with the latest changes this code path no longer sees multi-statement batches directly when new_parser is enabled. spliced_simple() handles the split upstream before routing runs so each statement routes individually. The stmts.len() > 1 guard in the router is now just a fallback for parse failures, session mode, and default builds
Implicit transactional behavior: this was the right concern to raise. Added tests documenting the behavior. PgDog splits the batch so each statement gets its own auto-commit rather than an implicit transaction wrapping the whole thing. test_multi_statement_implicit_transaction_behavior documents this as a known divergence from standard PostgreSQL behavior

Testing

I don't have permissions to run CI on this repo so I iterated locally. I ran the full Rust integration suite (cargo nextest run --profile integration) and the Go, Python, and SQL suites directly against a locally built binary. It took a bit of work to get local integration testing set up properly but wanted to make sure I wasn't breaking anything as I iterated. The failures I saw were all pre-existing environment issues (sharded pool config, unique_id node config, bundler version mismatch) and unrelated to these changes. The 9 new multi-statement tests all pass and I didn't see regressions in the suites I was able to run.

Note: the splitting behavior and new tests only activate under new_parser. The default build integration tests exercise the routing-only path.

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 pgdogdev#395

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

CLAassistant · 2026-07-28T16:12:06Z

All committers have signed the CLA.

hamin · 2026-07-28T16:20:50Z

@levkk :)

First of thank you for pgcat and pgdog!
cc would love for you take look when have the opportunity! Also pinged you on Discord FYI

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.

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.

sgrif

I don't think this handles the implicit transactional behavior of a multi-statement query. Can you add some integration tests for that?

sgrif · 2026-07-28T20:20:57Z

 pub use super::BufferedQuery;

+#[cfg(feature = "new_parser")]
+fn split_statements(query: &str) -> Vec<(usize, usize)> {


Should the signature of this be &str -> Vec<&str>?

sgrif · 2026-07-28T20:21:40Z

+#[cfg(feature = "new_parser")]
+fn split_statements(query: &str) -> Vec<(usize, usize)> {
+    let Ok(parsed) = pg_raw_parse::parse(query) else {
+        return vec![];


Why are we ignoring errors here? Do you think we should propagate them?

I updated the PR description for more context and my decision on this. Let me know how u feel

sgrif · 2026-07-28T20:23:45Z

        }

-        //
-        // Get the root AST node.


Now that we are relying on there only being one query per command, should we return an error if we got more than one?

I addressed this in the PR description...let me know what u think

sgrif · 2026-07-28T20:24:53Z


                let stmts = &statement.parse_result().protobuf.stmts;

-                // Handle multi-statement SET commands (e.g. "SET x TO 1; SET y TO 2").


The old parser is considered legacy code, new features don't need to support it.

ah awesome didn't know that. This makes perfect sense. Will remove this!

sgrif · 2026-07-28T20:25:27Z

+        .collect()
+}
+
+#[cfg(not(feature = "new_parser"))]


New features don't need to support the old parser.

hamin · 2026-07-29T13:36:54Z

I don't think this handles the implicit transactional behavior of a multi-statement query. Can you add some integration tests for that?

@sgrif thanks for your review and the is the biggest point you mentioned. I was hoping to see tests run on CI to see if I was failing anything but managed to finally run them on my machine and yeah saw that this definitely broke things...taking a different approach but will ping you all soon

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.

…ry splice 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

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.

hamin · 2026-07-29T15:12:22Z

OK @sgrif , thank u for catching the implicit transaction behavior! Force me to make sure integration tests pass locally before making further changes. I think I took @levkk 's initial comments in a Discord thread a bit too literally trying to share some code eventhough his comment was more about approach. I gave @levkk a bit more context in the Discord channel but I updated the PR description with what motivated this PR to begin with.

@sgrif @levkk would love a review/comments/feedback. FYI if you all feel this is way off base or conflicts with the future decisions that I'm completely unaware please let me know. Totally dogfooding this because we have a big PG upgrade migration that we really want to test out and use pgdog with to introduce it to our toolset not just for this but other purposes too (maybe sharding too but unsure of that currently)

sgrif · 2026-07-29T15:34:05Z

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


Sorry, I should have been more clear. This shouldn't use pg_query with the old parser, we can just feature flag the change behind the new parser for now.

Gottya, anything else?

Ok, updated spliced_simple(), split_statements(), and the related unit tests to all be gated behind #[cfg(feature = "new_parser")] using pg_raw_parse. On a default build spliced_simple() returns empty vec and callers fall through to the existing path.

codecov · 2026-07-29T17:13:37Z

Codecov Report

❌ Patch coverage is 90.79755% with 15 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
pgdog/src/frontend/client/query_engine/context.rs	16.66%	10 Missing ⚠️
pgdog/src/frontend/client_request.rs	97.43%	2 Missing ⚠️
pgdog/src/frontend/client/mod.rs	91.66%	1 Missing ⚠️
pgdog/src/frontend/client/query_engine/fake.rs	85.71%	1 Missing ⚠️
.../frontend/client/query_engine/start_transaction.rs	85.71%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!

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.

…r 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.

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.

hamin added 2 commits July 27, 2026 12:48

test: update integration tests for multi-statement query support

28ac651

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

hamin added 2 commits July 28, 2026 14:37

test: add multi-statement transaction batch integration test

3bb00ce

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.

sgrif requested changes Jul 28, 2026

View reviewed changes

hamin added 3 commits July 29, 2026 10:09

hamin force-pushed the feat/multi-statement-query-support branch from d766017 to 3bb00ce Compare July 29, 2026 14:10

sgrif reviewed Jul 29, 2026

View reviewed changes

hamin requested a review from sgrif July 29, 2026 17:46

hamin added 2 commits July 29, 2026 13:49

test(js): update mixed SET+SELECT test to expect success

ead6ed3

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.

hamin force-pushed the feat/multi-statement-query-support branch from 4569096 to ead6ed3 Compare July 29, 2026 19:25


		let stmts = &statement.parse_result().protobuf.stmts;

		// Handle multi-statement SET commands (e.g. "SET x TO 1; SET y TO 2").

Uh oh!

Conversation

hamin commented Jul 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Context and motivation

What changed

On the complexity

On @sgrif's review comments

Testing

Uh oh!

CLAassistant commented Jul 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

hamin commented Jul 28, 2026

Uh oh!

sgrif left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

hamin commented Jul 29, 2026

Uh oh!

hamin commented Jul 29, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

codecov Bot commented Jul 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

hamin commented Jul 28, 2026 •

edited

Loading

CLAassistant commented Jul 28, 2026 •

edited

Loading

codecov Bot commented Jul 29, 2026 •

edited

Loading