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
# Optional: how long a lookup query can run before the statement
# waiting on it fails, in milliseconds. Default: 5000.
sharding_lookup_timeout = 5000
# Optional: how many times routing retries after resolving lookups,
# in case the cache evicts an entry between attempts. Default: 3.
sharding_lookup_routing_attempts = 3

[[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 timeout) also fail
the statement with a retryable error. The timeout and the routing retry
bound are configurable: sharding_lookup_timeout (milliseconds, default
5000) and sharding_lookup_routing_attempts (default 3) in [general].

Metrics

Lookups are tracked per cluster through Cluster::stats(), whose
MirrorStats container is renamed to ClusterMetrics (it was only holding
mirror counts; those now live under a mirror field) and now also carries
the lookup counters, shared with the cluster's lookup cache, which records
them lock-free. Exposed on the OpenMetrics
endpoint, per cluster (user/database labels) plus a global rollup, same as
mirror stats:

sharding_lookup_cache_hits / sharding_lookup_cache_misses (counters):
cache effectiveness. A cold statement counts a miss at routing and a hit
when it routes again after the lookup resolves.
sharding_lookup_queries (counter) and
sharding_lookup_query_time_microseconds (counter): lookup queries run
for cache misses and the total time they took, including failed and
timed-out ones; the ratio is the average lookup latency.
sharding_lookup_cache_bytes / sharding_lookup_cache_entries (gauges):
cache size as computed by the weigher, and entry count.
sharding_lookup_cache_evictions (counter): entries evicted to stay
within the memory bound; invalidations and replacements don't count.

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.
Bare sharding keys, i.e. pgdog_sharding_key comments and
SET pgdog.sharding_key, translate through the lookup too, so there's no
partial state where a statement-extracted value translates but a hinted one
doesn't. The comment parser returns either a shard or a pending lookup; the
pending case flows into the same resolve-and-route-again path statements
use. Bare keys route through the common mapping, which now also requires
every table to agree on lookup_query; the cache key comes from the first
table, so hinted and extracted values share cached translations.
A shard directive on a write that only touches omnisharded tables is an
error (SQLSTATE 58000), matching the existing direct-to-shard transaction
check: the write must reach every shard, and honoring the directive lands
it on one, silently diverging the table. This covers pgdog_shard and
pgdog_sharding_key comments and their SET equivalents, with a warm or
cold cache; the cold case errors without running the lookup, since the
key may not have a mapping row yet when the INSERT is the statement
creating it. Reads keep the directive, since any shard holds the data.

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 94.42971% with 84 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
pgdog/src/frontend/router/sharding/lookup.rs	82.88%	32 Missing ⚠️
...og/src/frontend/client/query_engine/route_query.rs	33.33%	22 Missing ⚠️
.../backend/replication/logical/subscriber/context.rs	54.16%	11 Missing ⚠️
pgdog/src/backend/pool/cluster.rs	92.52%	8 Missing ⚠️
pgdog/src/stats/lookup.rs	93.75%	5 Missing ⚠️
pgdog/src/frontend/router/parser/copy.rs	98.97%	2 Missing ⚠️
...gdog/src/backend/pool/connection/mirror/handler.rs	96.42%	1 Missing ⚠️
pgdog/src/frontend/router/parser/comment/tests.rs	98.30%	1 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.

…rded writes

…d attempts

rlittlefield · 2026-07-30T16:27:56Z

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.

Just added metrics and configurable lookup timeout and lookup attempts!

sgrif · 2026-07-30T17:39:29Z

+        #[cfg(feature = "new_parser")]
+        {
+            let ast = pg_raw_parse::parse(query)
+                .map_err(|err| error(format!("\"lookup_query\" is invalid: {}", err)))?;
+
+            if ast.stmts().count() != 1 {
+                return Err(error("\"lookup_query\" must be a single statement".into()));
+            }
+
+            let mut found = HashSet::new();
+            for stmt in ast.stmts() {
+                pg_raw_parse::walk::walk(stmt, |node| {
+                    if let pg_raw_parse::Node::ParamRef(param) = node {
+                        found.insert(param.number);
+                    }
+                });
+            }
+            params = found;
+        }


The new parser version of this can be quite a bit simpler by breaking early if we ever see a parameter with a number > 1, and seeing if we got a value.

Suggested change

#[cfg(feature = "new_parser")]

{

let ast = pg_raw_parse::parse(query)

.map_err(|err| error(format!("\"lookup_query\" is invalid: {}", err)))?;

if ast.stmts().count() != 1 {

return Err(error("\"lookup_query\" must be a single statement".into()));

}

let mut found = HashSet::new();

for stmt in ast.stmts() {

pg_raw_parse::walk::walk(stmt, |node| {

if let pg_raw_parse::Node::ParamRef(param) = node {

found.insert(param.number);

}

});

}

params = found;

}

{

let ast = pg_raw_parse::parse(query)

.map_err(|err| error(format!("\"lookup_query\" is invalid: {}", err)))?;

let stmt = ast.stmts().exactly_one().map_err(|_| error("\"lookup_query\" must be a single statement".into()))?;

if pg_raw_parse::walk::walk_manual(stmt, |node| match node {

Node::ParamRef(param) if param.number != 1 => ControlFlow::Break(()),

_ => Recurse::yes(),

}).is_some() {

return Err(error(

"\"lookup_query\" must reference exactly one parameter, \"$1\"".into(),

));

}

}

(That's going to need some extra imports and formatting)

I would also split the whole function in two rather than a single statement, but we're also considering the old parser to be legacy at this point so it's completely fine to just not support this feature without using the new parser

levkk · 2026-07-30T17:46:10Z

    Lazy::new(|| Regex::new(r#"pgdog_shard: *([0-9]+)"#).unwrap());
 pub(super) static SHARDING_KEY: Lazy<Regex> = Lazy::new(|| {
-    Regex::new(r#"pgdog_sharding_key: *(?:"([^"]*)"|'([^']*)'|([0-9a-zA-Z_-]+))"#).unwrap()
+    Regex::new(r#"pgdog_sharding_key: *(?:"([^"]*)"|'([^']*)'|([0-9a-zA-Z-]+))"#).unwrap()


Just want to double check this isn't a breaking change (since I see tests were changed too).

We use uxid style ids, so we had org_1234 as the sharding key, but I didn't realize that there was a quoting option. So in an earlier commit I added the _ to the list, but I think that rather than add support for _, I'm just going to put the quotes around it on my app side. I think the net difference in this PR should be no actual changes to this Regex::new call.

levkk · 2026-07-30T17:47:26Z

+    /// How long a sharding key lookup query can run before the
+    /// statement waiting on it fails. Zero (a bare test cluster
+    /// built without config) falls back to the config default.
+    pub fn sharding_lookup_timeout(&self) -> Duration {


Defaults should propagate from pgdog-config/general.rs.

levkk · 2026-07-30T17:49:47Z


            attempts += 1;
-            if attempts > 3 {
+            if attempts > cluster.sharding_lookup_routing_attempts() {


I'm a bit hazy on this one: why do we have attempts at all? A lookup failure indicates that 1) sharding map is out of date or 2) the database isn't reachable.

I guess it could reduce the failure rate in case of 2). Just thinking out loud, this piece kinda stood out to me since we don't do retries anywhere else in the query engine, so why is this place special?

levkk · 2026-07-30T17:50:58Z

+            "counter",
+        );
+        let mut query_time = Series::new(
+            "sharding_lookup_query_time_microseconds",


We use ms with floating point 3 point precision everywhere else.

sgrif · 2026-07-30T18:03:19Z

+
+            params = ast
+                .protobuf
+                .nodes()


I'm not sure this will actually work, pg_query::protobuf::ParseResult::nodes does not actually iterate over all nodes as you might expect from the name.

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

rlittlefield added 3 commits July 29, 2026 19:06

feat: route sharding key directives through lookups, error on omnisha…

c0c2acf

…rded writes

fix: remove new _ because quoted strings work fine

bfed0a4

feat: lookup metrics on cluster stats, configurable lookup timeout an…

aac7637

…d attempts

rlittlefield force-pushed the sharded-tables-lookup-lazy branch from 9617b9c to aac7637 Compare July 30, 2026 15:39

sgrif reviewed Jul 30, 2026

View reviewed changes

levkk reviewed Jul 30, 2026

View reviewed changes

sgrif reviewed Jul 30, 2026

View reviewed changes

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

Metrics

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!

rlittlefield commented Jul 30, 2026

Uh oh!

sgrif Jul 30, 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!

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!

4 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

sgrif Jul 30, 2026 •

edited

Loading