feat: partial multi-query support

[levkk]

What

Support multi-statement queries sent via the simple protocol if they can be safely executed without us having to do anything fancy with transactions.

Preface

A multi-statement query looks something like this:

SELECT 1; SELECT 2;

They are sent inside a single Query message and Postgres is expected to process all commands sequentially and inside a single, implicit transaction.

This presents an issue for sharded deployments, since queries might need to go to different shards. So, what can we do about this?

Split and execute

Since we have the Postgres parser handy, we can split the query and execute each statement separately. One caveat here is we won't start a transaction inside Postgres automatically (at least, not yet), so we have to be careful about which statements we can execute safely.

SET mixed with one query

This is totally safe (and quite common):

SET statement_timeout TO 0; SELECT * FROM users;

We can split these up, execute them separately, without breaking transactional guarantees. Only one statement actually reads data, while PgDog handles SET internally. This type of query will, as of this PR, be executed normally. Yay!

Conversely, the following cannot be split safely (yet) and will continue to return an error from us:

SET statement_timeout TO 0; SELECT * FROM users; SELECT * FROM orders;

You can set as many things as you want:

SET statement_timeout TO 0;
SET lock_timeout TO 0;
SELECT * FROM pg_type;
SHOW lock_timeout;
RESET statement_timeout;
RESET lock_timeout;

This is safe to execute because only one query actually cares about data, and the rest is handled by PgDog or is inconsequential if executed outside a transaction.

DDL

CREATE TABLE x [...];
CREATE TABLE y [...];

This is quite common (e.g., in migration scripts) and is quite safe to execute without splitting at all. We send this query to all shards concurrently. No problems.

Manual transactions

This happens sometimes and now works correctly:

BEGIN; SELECT * FROM users WHERE id = 1; SELECT * FROM orders WHERE user_id = 2; COMMIT;

This is safe to split and execute because client starts a transaction and we can handle those no problem. If the transaction returns an error, the remaining statements are not executed; same behavior as regular Postgres.

Checks

Other things we check for and block now:

1. Unclosed transactions, e.g., BEGIN; SELECT 1;. Bad! Blocked.
2. Multi-statement queries outside transactions, bad, blocked!
3. Mix of transactions and not, bad, blocked!

This is only relevant for a single Query message. Normal query flow is not affected by this change.

[codecov]

Codecov Report

❌ Patch coverage is 99.37304% with 2 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
pgdog/src/frontend/client/query_engine/query.rs	92.30%	1 Missing ⚠️
...rc/frontend/router/parser/query/test/test_split.rs	99.16%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!

[sgrif]

The transaction_safe check feels a little overboard to me, but this seems fine since we intend to move to implicit transactions down the line.

With the changes to how planning works that John and I are working on, we're thinking we may end up having communication with the shards always use the extended protocol if the query parser is enabled, which will make auto-transactions much easier both here and with any other 1->N rewrites we do going forward.

[sgrif]

Love this

[sgrif]

What do you think about making this a bit more descriptive?

Suggested change

	pub(super) fn simple_pipeline_check(&self, context: &QueryEngineContext<'_>) -> bool {
	pub(super) fn in_errored_simple_pipeline(&self, context: &QueryEngineContext<'_>) -> bool {

[sgrif]

I like this structure. Long term I think it'd be nice to have more of the state bundled here (not asking for that to be changed in this PR)