feat: partial sharding support for PREPARE and EXECUTE by levkk · Pull Request #1368 · pgdogdev/pgdog

levkk · 2026-08-16T05:23:52Z

Prepared statements

Correctly store statements sent with PREPARE in the global cache. We use the whole query string with data type specifications as key, while anonymizing the statement name. This is equivalent to using Parse as cache key.

PREPARE __stmt_1(int) AS SELECT $1;

becomes

PREPARE __pgdog_template_name(int) AS SELECT $1;

It's then dynamically rewritten to

PREPARE __pgdog_1(int) AS SELECT $1

at execution time.

This is because it's kinda hard to parse the data types out of the string and we need a unique key identifying the prepared statement in the cache.

Sharding

Handle EXECUTE by extracting parameters and passing them to our query router. This works with I think most of our sharding features, i.e., direct-to-shard and cross-shard queries.

TODOs

Comment-based routing isn't supported. This is because we parse and deparse the original statement to put it in the prepared statement cache, which erases the comment.
Query rewrites, e.g., insert split, don't work yet.

Notes for reviewers

Most of the diff in client/prepared_statements module is from moving stuff around. There is some diff for handling PREPARE, that's where I would focus the review.
Simulated messages sent by Server now get processed through the server state machine; this was necessary for correct state management for handling PREPARE statements; seems to have no bad side effects.

codecov · 2026-08-16T05:28:50Z

Codecov Report

❌ Patch coverage is 94.36170% with 53 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
pgdog/src/net/messages/prepare.rs	56.25%	14 Missing ⚠️
pgdog/src/frontend/router/parser/params.rs	84.50%	11 Missing ⚠️
...gdog/src/frontend/prepared_statements/statement.rs	87.50%	6 Missing ⚠️
pgdog/src/frontend/router/parser/query/prepare.rs	85.36%	6 Missing ⚠️
...router/parser/rewrite/statement/simple_prepared.rs	93.81%	6 Missing ⚠️
.../src/frontend/client/query_engine/test/replicas.rs	75.00%	4 Missing ⚠️
...gdog/src/frontend/prepared_statements/cache_key.rs	84.61%	2 Missing ⚠️
pgdog/src/net/protocol_message.rs	60.00%	2 Missing ⚠️
...frontend/router/parser/query/test/test_prepared.rs	96.77%	1 Missing ⚠️
...rontend/router/parser/rewrite/statement/auto_id.rs	97.91%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!

meskill · 2026-08-18T10:32:41Z

+impl FromBytes for Prepare {
+    fn from_bytes(_bytes: Bytes) -> Result<Self, Error> {
+        unreachable!("Prepare must be constructed manually")
+    }
+}
+
+impl Protocol for Prepare {
+    fn code(&self) -> char {
+        'Q'
+    }
+
+    fn message(&self) -> Result<Message, Error> {
+        Ok(Message::new(self.to_bytes()).frontend())
+    }
+
+    fn streaming(&self) -> bool {
+        false
+    }
+}


is this actually needed?

meskill · 2026-08-18T10:37:07Z

+impl ToBytes for Prepare {
+    fn to_bytes(&self) -> Bytes {
+        let query = self.query();
+        let name = self.name();
+        // This is safe because the statement looks like this:
+        // PREPARE __pgdog_template_name AS [...]
+        // so the template name will always match first.
+        let query = query.replacen(PREPARE_TEMPLATE_NAME, name, 1);
+
+        Query::new(query).to_bytes()
+    }
+}


nit: I understand why we need the "fake" messages, but I don't like much that we imitate that it's actual message that has byte representation. I'd prefer if we separate this kind of message in the parent enum to make it explicit and do the conversion to the some other real ProtocolMessage to use it's implementation. I.e. instead of ToBytes implement into_query() and call to_bytes on it

meskill · 2026-08-18T11:09:47Z

+                        }
+                    }
+
+                    _ => BindParameter::new_null(),


I think we should handle TypeCast as well and not only ConstValue, so "445-45::uuid" works and the compute_shard_for_table won't trigger broadcasting.
I think there also could FuncCalls and this probably should be revisited later on routing sharding, erroring etc., but anyway null is too destructive for this and misleading

meskill · 2026-08-18T11:29:10Z

+/// A `Simple` key comes from SQL `PREPARE` and matches nothing but itself.
+/// Its declared argument types are not captured, so two of those
+/// statements are never known to be the same.
+///


that's not true anymore

meskill · 2026-08-18T11:40:25Z

+        let stmt_params = StatementParameters::Execute(&params);
+
+        // Create new parser context.
+        let mut context = context.clone();


that's looks fishy - we clone to mutate and then after mutation we drop this. I don't know maybe it's covered somehow or not harmful, but that is very brittle. I noticed only because the derive(Clone) appeared on the Context

meskill · 2026-08-18T13:49:24Z

+                // Will use parameters for replacing args, not materialize values.
+                self.extended = true;


won't this cause any issues with the actual extended protocol? yeah, like using PREPARE inside Parse 😄

Oh god I hope Postgres triggers an error when that happens....

meskill · 2026-08-18T13:53:50Z

+    #[error("prepare statement can only be DML")]
+    PrepareNotDml,


meskill · 2026-08-18T13:57:20Z

-                if self.contains(name) {
+            ProtocolMessage::PrepareFromClient(prepare) => {
+                use crate::net::{CommandComplete, ReadyForQuery};
+                if self.contains(prepare.name()) {


I think we need to use check_prepared at some extend - the ttl functionality and parses deduplication is not working for prepared rn

meskill · 2026-08-18T14:01:56Z

        if let NodeMut::SelectStmt(mut select) = stmt.stmt_mut() {
            self.rewrite_aggregates(&mut select, mem, &mut plan, self.db_schema)?;
            self.limit_offset(&select, &mut plan);
        }


well, should we do something specific for the prepare/execute here?

Ah yeah, probably doesn't work currently.

meskill · 2026-08-18T14:06:50Z

                // INVARIANT: omni dedup in multi_shard relies on this being process-unique;
                // never substitute a non-unique value here.
-                return Ok(message.backend(self.id));
+                break message.backend(self.id);


why this change? that seems like simulated prepared will now be processed and setting sync_prepared=true that will cause more work on check-in

Because we need CommandComplete and ReadyForQuery to reset the server state back to Idle when we prepare a statement. Otherwise, it gets stuck in reading data. Preparing a statement using Query basically executes a separate query before the client's request, so the state management is a bit more complex.

We just push whatever message we "simulate" through the state manager though and it seems to work.

#1368 introduced support for `PREPARE`, but handling of it was duplicating an internal state tracker which could prevent the server connection from being returned into the pool after the transaction was done. Also removed some dead code.

levkk requested a review from meskill August 16, 2026 20:02

levkk marked this pull request as ready for review August 16, 2026 20:03

levkk changed the title ~~feat: sharding support for PREPARE and EXECUTE~~ feat: partial sharding support for PREPARE and EXECUTE Aug 17, 2026

This was referenced Aug 17, 2026

Cache sharing for PREPARED statements #1350

Closed

fix: don't sum row counts across shards for EXECUTE on omnisharded tables #1179

Closed

levkk added 19 commits August 17, 2026 11:45

feat: save simple prepared in global cache

4f1b98c

fix invariant

5fa99b6

save

f0df3ca

stmt parameter

53e0237

sharding support for prepare

edf0627

todo

bf6675f

support for prepared

fc209d7

use correct key

5cdec64

fix tests

c16ec72

tests

44d9020

rename

01a9579

fmt

312a6cb

aha!

99526be

fix test

e0d26c6

save

6a1db8a

move test

f01d710

process simulated messages through state manager

de027de

tests

6b8cbe2

Drop: Clone shouldnt mutate global state

ebcc0c7

levkk force-pushed the levkk-prepared-simple-routing branch from e2e9044 to ebcc0c7 Compare August 17, 2026 18:48

levkk added 3 commits August 17, 2026 14:57

woops

158f761

support unique id and id injection in prepare

0038a5e

test

42634de

levkk merged commit 0c8cf53 into main Aug 17, 2026
7 checks passed

levkk deleted the levkk-prepared-simple-routing branch August 17, 2026 23:09

meskill reviewed Aug 18, 2026

View reviewed changes

This was referenced Aug 18, 2026

PREPARE doesn't work with prepared statements TTL #1382

Open

PREPARE and OFFSET doesn't work for cross-shard queries #1383

Open

fix: PREPARE could prevent connection re-use #1385

Merged

levkk mentioned this pull request Aug 18, 2026

docs: document sharding support for EXECUTE (and PREPARE) pgdogdev/docs#110

Open

		// Will use parameters for replacing args, not materialize values.
		self.extended = true;

Uh oh!

Conversation

levkk commented Aug 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Prepared statements

Sharding

TODOs

Notes for reviewers

Uh oh!

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

Uh oh!

Codecov Report

Uh oh!

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!

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!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

levkk commented Aug 16, 2026 •

edited

Loading

codecov Bot commented Aug 16, 2026 •

edited

Loading