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

[mscrivo]

Problem

EXECUTE of a server-side prepared statement targeting an omnisharded table returned the sum of row counts from all shards instead of the count from one shard:

PREPARE upd AS UPDATE sharded_omni SET value = $1;
EXECUTE upd('x');  -- UPDATE 6 on a 2-shard cluster with 3 rows (should be UPDATE 3)

Plain UPDATE/DELETE/INSERT on omnisharded tables already report the correct count (#752), but ExecuteStmt fell through to the DDL catch-all, which broadcasts with Route::write(Shard::All) and no omnisharded flag — so MultiShard aggregated the CommandComplete counts.

Fix

• Route PREPARE explicitly: still broadcast to all shards, but extract the statement behind the name and store it in the regular prepared statements cache (Parse::named(name, statement) + PreparedStatements::insert()), the same way the prepared_statements = "full" rewriter does. First PREPARE wins, matching the server, where a duplicate PREPARE fails.
• Route EXECUTE explicitly: resolve the name via PreparedStatements::parse() and inspect the statement through the statement cache, so each unique statement is parsed once, not per execution. If it's a write that only touches omnisharded tables, the route carries the omnisharded flag and MultiShard deduplicates results across shards instead of aggregating them.
• Reads behind EXECUTE are not flagged: EXECUTE always routes as a write (the statement only exists on the connections PREPARE reached), and an omnisharded write requires full shard coverage — flagging reads would reject shard directives on statements that can't diverge the shards.
• The router reaches the client's prepared statements via a new optional RouterContext::with_prepared_statements().

This also extends the omni-write-in-direct-to-shard-transaction guard (#1086) and the omni-write-with-directive guard to EXECUTE of omnisharded writes, since both key off the same route flag.

Testing

• Routing tests in parser/query/test/test_execute.rs (fail without the fix): omnisharded UPDATE/DELETE/INSERT behind EXECUTE carry the flag; SELECT/VALUES/sharded/unknown names don't; shard directives are allowed on reads and rejected on omnisharded writes.
• Verified end-to-end against a 2-shard cluster: EXECUTE upd('x') returns UPDATE 3 instead of UPDATE 6, and the write still reaches every shard. Same for DELETE.

🤖 Generated with Claude Code

[codecov]

Codecov Report

❌ Patch coverage is 93.64162% with 11 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
pgdog/src/frontend/router/parser/query/execute.rs	88.40%	8 Missing ⚠️
...og/src/frontend/client/query_engine/route_query.rs	25.00%	3 Missing ⚠️

📢 Thoughts on this report? Let us know!

[sgrif]

This touches the parser, which we are in the middle of a major rewrite of. This isn't likely to get reviewed until after that is done (we are hoping next week, but take that with a huge grain of salt)

[levkk]

I'm not sure this needs its own dedicated cache. You could extract the statement from PREPARE and store it in the global cache, e.g.:

PREPARE __sqlx_1 AS SELECT $1;

when extracted becomes:

Name	Statement
__sqlx_1	SELECT $1

You can create a Parse from this and insert it into the regular cache, e.g.:

let parse = Parse::named("__sqlx_1", "SELECT $1");
prepared_statements.insert(&mut parse);

When capturing EXECUTE, you can now fetch the original statement by calling:

let name = /* get name from EXECUTE using parser */;
prepared_statements.parse(name);

We do this already in the rewriter here:

pgdog/pgdog/src/frontend/router/parser/rewrite/statement/simple_prepared.rs

Line 77 in bafec81

let mut parse = Parse::named(

[mscrivo]

thanks Lev! I'm reworking this now since I found some other issues as well. Sorry for the churn

[mscrivo]

Done

[levkk]

This should be using the statement cache, so we don't have to parse the statement on each execution.

[mscrivo]

Done

[levkk]

It's a good start, but I think the correct solution here is to handle EXECUTE in sharded databases overall. While this works for omni tables, we should handle this for sharded tables, too.

[mscrivo]

It's a good start, but I think the correct solution here is to handle EXECUTE in sharded databases overall. While this works for omni tables, we should handle this for sharded tables, too.

@levkk I'll leave it with you and @sgrif at this point. I don't have enough context with the codebase or experience with rust, so I'm in over my head at this point :) It does demonstrate the outcome I'd like to see, so if you all close this PR and rework it entirely, that'd be fine with me as well.

[levkk]

It's a good start, but I think the correct solution here is to handle EXECUTE in sharded databases overall. While this works for omni tables, we should handle this for sharded tables, too.

@levkk I'll leave it with you and @sgrif at this point. I don't have enough context with the codebase or experience with rust, so I'm in over my head at this point :) It does demonstrate the outcome I'd like to see, so if you all close this PR and rework it entirely, that'd be fine with me as well.

I'm not opposed to merging this (with maybe some minor changes) to unblock you guys! Let me know how soon you need this and we'll prioritize. We can always add support for sharded EXECUTE as a follow up.

[mscrivo]

It's a good start, but I think the correct solution here is to handle EXECUTE in sharded databases overall. While this works for omni tables, we should handle this for sharded tables, too.

@levkk I'll leave it with you and @sgrif at this point. I don't have enough context with the codebase or experience with rust, so I'm in over my head at this point :) It does demonstrate the outcome I'd like to see, so if you all close this PR and rework it entirely, that'd be fine with me as well.

I'm not opposed to merging this (with maybe some minor changes) to unblock you guys! Let me know how soon you need this and we'll prioritize. We can always add support for sharded EXECUTE as a follow up.

as soon as you're able I guess. We're not in any particular rush, but the problem on our end is that there's still a rather large bucket of test failures we're dealing with and so it makes it difficult to triage and categorize them. ie are they from this issue, another fixed or still open issue, etc. So the more we knock out, the easier that process becomes.