Skip to content

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

Draft
mscrivo wants to merge 2 commits into
pgdogdev:mainfrom
affinity:fix/omni-execute-row-counts
Draft

fix: don't sum row counts across shards for EXECUTE on omnisharded tables#1179
mscrivo wants to merge 2 commits into
pgdogdev:mainfrom
affinity:fix/omni-execute-row-counts

Conversation

@mscrivo

@mscrivo mscrivo commented Jul 10, 2026

Copy link
Copy Markdown
Contributor

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 record the statement behind the name in the client's PreparedStatements state (new simple map, cleared with close_all()).
  • Route EXECUTE explicitly: resolve the name — first from the recorded SQL-level PREPAREs, falling back to the global cache for names rewritten by prepared_statements = "full" — and set the omnisharded flag on the route when the underlying statement only touches omnisharded tables. MultiShard then deduplicates results across shards instead of aggregating them.
  • The router gets access to 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) to EXECUTE of omnisharded writes, since that check keys off the same route flag.

Testing

  • New routing tests in parser/query/test/test_execute.rs (fail without the fix): EXECUTE of omnisharded UPDATE/DELETE carries the omnisharded flag; sharded/unknown statements don't; global-name resolution works.
  • Verified end-to-end against a 2-shard cluster: EXECUTE upd('x') now returns UPDATE 3 instead of UPDATE 6, and the write still reaches every shard. Same for DELETE.
  • cargo nextest parser suite, cargo clippy, cargo check --features new_parser all pass.

🤖 Generated with Claude Code

mscrivo and others added 2 commits July 10, 2026 16:46
…bles

EXECUTE of a server-side prepared statement was routed through the DDL
catch-all, which broadcasts to all shards without the omnisharded flag.
For statements that only touch omnisharded tables, the CommandComplete
row counts from each shard were summed, so a DELETE/UPDATE of N rows
reported N * shards rows to the client.

Route PREPARE explicitly and record the statement behind the name in
the client's prepared statements state. EXECUTE resolves the name (or
a globally cached name rewritten by prepared_statements = "full") and
carries the omnisharded flag, so cross-shard results are deduplicated,
not aggregated.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@codecov

codecov Bot commented Jul 10, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 94.23077% with 9 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
pgdog/src/frontend/router/parser/query/execute.rs 90.76% 6 Missing ⚠️
pgdog/src/frontend/prepared_statements/mod.rs 84.61% 2 Missing ⚠️
...og/src/frontend/client/query_engine/route_query.rs 50.00% 1 Missing ⚠️

📢 Thoughts on this report? Let us know!

@sgrif sgrif self-assigned this Jul 10, 2026
@sgrif

sgrif commented Jul 10, 2026

Copy link
Copy Markdown
Contributor

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)

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants