db-based sharding key lookup_query mapping by rlittlefield · Pull Request #1279 · pgdogdev/pgdog

rlittlefield · 2026-07-28T20:24:37Z

Query-based sharding key lookup for `[[sharded_tables]]`

Implements the lazy per-key lookup design: a sharded table config can declare a
lookup query that translates an extracted sharding key value into the value
that actually gets hashed. Query returns a row: the result is hashed, and the
translation is cached, keyed by the value. No row: the statement fails with a
retryable error, and nothing is cached.

[general]
# Optional: memory bound for cached translations, per cluster.
# Default: 64MB.
sharding_lookup_cache_size = 67108864

[[sharded_tables]]
database = "prod"
column = "tenant_id"
data_type = "varchar"
lookup_query = "SELECT COALESCE(parent_tenant_id, id) FROM tenants WHERE id = $1"

Semantics

When the router extracts a sharding key for a lookup-configured table, the
value is translated through the cache first; the result is hashed with the
table's configured hasher and data_type to pick the shard.
On a cache miss, the statement waits while the lookup query runs (one query,
one shard, round-robin), then routes with the translated value. The value is
bound with the extended protocol, so it's never spliced into SQL text.
A value with no row fails the statement with an error instead of routing by
the untranslated value, which could place data on the wrong shard. The main
way this happens in practice is traffic for a key whose row hasn't been
inserted yet; because absence isn't cached, the first statement after the
row appears translates correctly. The query should therefore return a row
for every valid value, including values that translate to themselves
(e.g. COALESCE(parent_tenant_id, id) gives root tenants a self-row).
The cache is scoped to the cluster, following the SchemaCache pattern:
it's recreated, and therefore emptied, on config reload. Since only
translations are cached and rows are looked up on first use, a reload is
only needed to pick up changes to existing rows.
The cache is memory-bound, backed by moka and sized by a new
sharding_lookup_cache_size setting in [general] (bytes, default 64MB,
per cluster). The weigher counts key and value bytes plus an approximation
of per-entry overhead. Entries are keyed by the sharded table (schema,
name, column) and the value, so identical queries on different tables
can't collide. An evicted translation is simply looked up again on its
next use; the resolve loops retry bounded in case an entry is evicted
between routing attempts.
Lookup failures (query error, no server available, a 5s timeout) also fail
the statement with a retryable error.

Validation

lookup_query is validated with the SQL parser at config load: it must be
syntactically valid, a single statement, and reference exactly one parameter,
$1. The validation runs in the pgdog crate, matched to whichever parser
the build uses; linking pg_raw_parse into pgdog-config unconditionally
turned out to corrupt pg_query builds through duplicate libpg_query C
symbols. The table the query reads must be omnisharded, i.e. hold the same
data on all shards, since the query runs on one shard picked round-robin;
this part is documented rather than validated, because the config doesn't
name the table.

Notes

"No row" is an error rather than routing by the original value. Routing an
untranslated value silently splits a tenant family across shards the moment
a lookup races a row insert, and a loud, retryable error seemed strictly
better; happy to flip it (or make it a per-table option) if you prefer the
identity behavior. In my use-case we would have the ability to ensure we
can insert the row before we start using the new sharding key id.
The lookup runs through two new building blocks that are generally useful
on their own: ServerRequest, which lets execute, execute_checked and
fetch_all take either simple-protocol queries or an extended-protocol
batch (anonymous Parse + Bind + Execute + Sync, parameters bound by the
server), and Cluster::fetch_all_round_robin, which runs a parameterized
query on one shard picked round-robin.
Binary-format values (bind parameters, COPY tuples, replicated WAL data)
are decoded to their text form and translated like text values, reusing the
existing binary decoding in sharding::Value.
COPY row sharding and logical replication consult the lookups too, so
resharding a lookup-translated table places rows correctly. CopyParser::shard
is async and resolves missing keys inline, in row order, through the same
resolver routed queries use. Replicated statements resolve their lookups in
StreamContext the same way. Both text and binary COPY formats are supported, so
the default resharding_copy_format works unchanged. During data sync,
lookup queries run on the source cluster, which has the complete mapping
while the destination is still syncing, and fill the destination's cache;
table sync order therefore doesn't matter. During the streaming phase the
destination answers lookups, which works because data sync completed and
the stream itself keeps the mapping current; a mapping row and rows keyed
by it written in the same source transaction would fail the stream, which
the insert-the-mapping-first application pattern avoids.
A hot key that legitimately has no row runs the lookup query on every
statement before erroring, since absence isn't cached. Complete mappings
(a row per valid value) avoid this entirely.

CLAassistant · 2026-07-28T20:24:45Z

All committers have signed the CLA.

levkk · 2026-07-28T23:55:48Z

+
+    fn validate_sharded_tables(&self) -> Result<(), Error> {
+        for table in &self.config.sharded_tables {
+            if let Some(query) = &table.lookup_query


nit: we should use pg_parse_raw::parse here:

will validate that the query is syntactically valid

doesn't rely on string interpolation

levkk · 2026-07-28T23:57:19Z

+        // so lookup queries run on the source, which has the complete
+        // mapping, and fill the destination's cache, which the row
+        // sharder reads.
+        let pending = self.copy.pending_lookups();


nit: copy/pasta.

levkk · 2026-07-29T00:00:36Z

+    /// Execute a query with text-format parameters using the extended
+    /// protocol and return all rows. Parameters are bound by the server,
+    /// so values don't need to be escaped.
+    pub async fn fetch_all_params<T: From<DataRow>>(


refactor: we should change fetch_all and consequently execute_checked and execute to accept a generic type that resolves to either a Vec<Query> or Vec<"extended protocol request">. The protocol handling below is brittle and needs to re-use existing logic.

levkk · 2026-07-29T00:02:07Z

+            if !pending.is_empty() {
+                match lookup::resolve(cluster, pending).await {
+                    Ok(()) => {
+                        let router_context = RouterContext::new(


nit: I'd probably construct a new router here instead of relying on internals and calling take_pending_lookups() twice.

codecov · 2026-07-29T00:05:08Z

Codecov Report

❌ Patch coverage is 92.23301% with 72 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
pgdog/src/frontend/router/sharding/lookup.rs	69.81%	32 Missing ⚠️
...og/src/frontend/client/query_engine/route_query.rs	26.66%	22 Missing ⚠️
.../backend/replication/logical/subscriber/context.rs	54.16%	11 Missing ⚠️
pgdog/src/backend/pool/cluster.rs	80.00%	3 Missing ⚠️
pgdog/src/frontend/router/parser/copy.rs	98.97%	2 Missing ⚠️
pgdog/src/frontend/router/parser/statement.rs	99.49%	1 Missing ⚠️
pgdog/src/frontend/router/sharding/value.rs	96.96%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!

levkk · 2026-07-29T00:09:21Z

    system_catalogs: SystemCatalogsBehavior,
+    /// Cached sharding key lookup query results. Scoped here so the
+    /// cache is recreated, i.e. emptied, on config reload.
+    lookup_cache: LookupCache,


I lost track of this one. We want to double check that it lives on the Cluster, i.e., we get a new version of this when we call reload_from_existing or reload.

levkk · 2026-07-29T00:10:05Z

Looks pretty good! I need to read it some more, but architecturally it's sound. Just a few nits around code maintenability / quality, but nothing major/blocking I think.

levkk · 2026-07-29T00:11:24Z

+        // can't be resolved fails the statement: routing by the
+        // untranslated value could put it on the wrong shard.
+        if result.is_ok() {
+            let pending = self.router.take_pending_lookups();


These state modifications are hard to track. I think we introduced this pattern and should not do it anymore. Maybe we can return these as part of the Route? That one is created for each request, so no concerns about leaking state between calls.

levkk · 2026-07-29T00:11:52Z

    pub fn parse(&mut self, context: RouterContext) -> Result<Command, Error> {
+        // Pending lookups are per-statement; don't leak them into
+        // the next statement if this parser is reused.
+        self.pending_lookups.clear();


Same comment re: Route as above. We shouldn't leak this between calls like this in the first place.

levkk · 2026-07-29T00:12:27Z

            }
        };

+        self.pending_lookups.extend(pending_lookups);


Maybe move this to QueryParserContext? A new one is created with each request, so data between calls won't leak.

levkk · 2026-07-29T00:14:00Z

+                            // The terminator must ship after every data row;
+                            // hold it back while rows are deferred.
+                            if !self.deferred.is_empty() {
+                                self.deferred.push(DeferredRow {


Maybe we make this method async instead? I get the defer pattern, but it sure makes things hard to follow. I'm pretty sure (i.e., not sure at all) that this method only gets called in async contexts anyway.

levkk · 2026-07-29T00:16:23Z

+    fn translate_sharding_key(&mut self, table: &ShardedTable, value: &str) -> Option<Arc<str>> {
+        let query = table.lookup_query.as_deref()?;
+
+        match self.schema.tables.lookup_cache().get(query, value) {


This should really be a lookup by ShardedTable (name, schema, database is the fully-qualified lookup). It's possible (but unlikely) to have two sharded databases with the same lookup query, and different lookup tables, so this could route things incorrectly.

levkk · 2026-07-29T00:19:39Z

+/// on the source cluster, which has the complete mapping while the
+/// destination is still syncing, and fills the destination's cache,
+/// which the row sharder reads.
+pub(crate) async fn resolve_with(


Need to think about this one. Nothing wrong with it, but a lot of that code is duplicated, and I'm thinking we should either move it or add a function to Cluster to execute a query against a shard using round robin.

levkk · 2026-07-29T00:25:55Z

+#[derive(Debug, Default, Clone)]
+pub struct LookupCache {
+    // Keyed by lookup query.
+    cache: Arc<DashMap<String, Translations>>,


Another thing that came up lately is we really should have memory-bound/unit-bound caching for this kind of stuff. Maybe use https://github.com/moka-rs/moka? We otherwise have the lru crate for a simpler LRU cache, which we could also use here, albeit it won't be sharded, so maybe that's why moka would be a better choice.

levkk · 2026-07-29T21:29:00Z

I think we just need to make all hardcoded timeouts / attempts configurable and this might be it. I need to double check the parser changes a couple more times, that one is always hard to review in GitHub because of how large that file is, but if the tests are passing, I suspect it's a-ok.

levkk · 2026-07-29T21:30:53Z

Oh and before I forget. Metrics! We need to track lookup queries: cache hits and misses and actual cache size + evictions. That one is gonna be an important metric of performance. For cache misses, we should track the latency of lookups.

For those, I would follow the cluster-global metrics pattern

pgdog/pgdog/src/backend/pool/cluster.rs

Line 506 in b2a61dd

pub fn stats(&self) -> Arc<Mutex<MirrorStats>> {

... which should be renamed to ClusterMetrics of course. Currently used for mirroring stats only, but this would expand it.

rlittlefield added 4 commits July 28, 2026 12:06

feat: lazy per-key sharding key lookup for sharded tables

2766e74

feat: consult sharding key lookups in COPY and logical replication

a352448

feat: translate binary sharding keys through lookups

864a3ce

fix: hold end-of-copy markers behind deferred lookup rows

4decdc2

feat: resolve data sync lookups on the source cluster

12b3f14

levkk reviewed Jul 28, 2026

View reviewed changes

levkk reviewed Jul 29, 2026

View reviewed changes

rlittlefield added 2 commits July 29, 2026 11:29

refactor: address review feedback on sharding key lookups

7b4bc6f

feat: configurable sharding key lookup cache size

3745a8a

rlittlefield changed the title ~~Sharded tables lookup lazy~~ db-based sharding key lookup_query mapping Jul 29, 2026

Uh oh!

Conversation

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

Uh oh!

Query-based sharding key lookup for [[sharded_tables]]

Semantics

Validation

Notes

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!

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!

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!

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

Uh oh!

Choose a reason for hiding this comment

Uh oh!

levkk 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!

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!

levkk commented Jul 29, 2026

Uh oh!

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

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

rlittlefield commented Jul 28, 2026 •

edited

Loading

Query-based sharding key lookup for `[[sharded_tables]]`

CLAassistant commented Jul 28, 2026 •

edited

Loading

codecov Bot commented Jul 29, 2026 •

edited

Loading

levkk Jul 29, 2026 •

edited

Loading

levkk commented Jul 29, 2026 •

edited

Loading