Skip to content

feat: query-based sharding key lookup for sharded tables - #1274

Draft
rlittlefield wants to merge 1 commit into
pgdogdev:mainfrom
rlittlefield:sharded-tables-query-lookup
Draft

feat: query-based sharding key lookup for sharded tables#1274
rlittlefield wants to merge 1 commit into
pgdogdev:mainfrom
rlittlefield:sharded-tables-query-lookup

Conversation

@rlittlefield

@rlittlefield rlittlefield commented Jul 28, 2026

Copy link
Copy Markdown

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

Hi! This should be thought of as a potential way to do table based mapping - I'm not familiar enough with the whole codebase to determine if this is a good approach, and I had to be opinionated in a few places. Mostly this can be taken as a proof of concept, that does seem to be working fine in my local development system so far.

This implements a database-backed sharding lookup using new "lookup" and "query" parameters.
query will be the query to run, which will return a mapping between the detected sharding key and the mapped sharding key. The query takes no input parameters, and the entire result set is cached in memory for future reference.
The lookup property indicates which table to track writes for invalidation of that cache. This allows occasional appends to the mapping table. The referenced table must be omnisharded.

[[sharded_tables]]
database = "prod"
column = "tenant_id"
data_type = "varchar"
lookup = "tenants"
query = "SELECT id, COALESCE(parent_tenant_id, id) FROM tenants"

# The lookup table must be omnisharded: it's loaded from a single
# shard picked round-robin, so every shard must have the same data.
[[omnisharded_tables]]
database = "prod"
tables = ["tenants"]

Requirements

Three hard requirements on the lookup table:

  1. It must be omnisharded (validated at config load). The map is loaded from
    one shard picked round-robin, so every shard must hold identical data.

  2. It must be complete. The query must return a row for every value that
    will ever route through the sharded column, including values that translate
    to themselves (e.g. COALESCE(parent_tenant_id, id) gives root
    tenants a self-row). A value with no row fails the statement with an
    error; it never routes.

  3. Mappings are insert-only, unless a single PgDog instance fronts the cluster.
    Invalidation is per-instance: when a write to the lookup table completes,
    the instance that carried it drops and reloads its map, but a hit in another
    instance's map is served without re-verification, so that instance keeps
    translating by the old row until it observes a write itself. New rows are
    safe across instances, because a value missing from the map is never
    trusted: it triggers a verifying reload before routing, so any instance
    discovers a freshly inserted mapping on first use. Updated or deleted rows
    have no such trigger; the stale hit looks like a perfectly good answer.
    Mutating an existing mapping is therefore only safe when there's exactly
    one PgDog instance, or by treating it as an operational event (drain and
    restart the other instances, or route all traffic through one instance for
    the change).

    Even if this mapping is on a single pgdog instance, I'm not sure it would
    ever be safe to modify an existing mapping, but since the underlying table
    may have other properties that aren't part of the returned mapping, we don't
    want to prevent writes wholesale.

Why a whole-table query instead of WHERE key = $1

The query shape deviates from the original idea where a $1 query returning one
translated row per lookup. Loading the whole table with a single parameterless
two-column query, combined with the completeness requirement, buys two things a
per-key lookup can't:

  • Multi-instance correctness without a coordination channel. Because every
    legitimate value has a row, a miss is anomalous and is never trusted as-is:
    unless the map was loaded within the last second, PgDog reloads the table and
    re-checks before deciding. A mapping row inserted through one instance is
    found by any other instance on first use, since the miss itself triggers the
    verifying reload. No LISTEN/NOTIFY, no polling interval, no window where a
    new mapping is silently invisible. (A per-key design has to cache negative
    results to avoid a database round-trip on every unmapped value, and a cached
    negative goes stale the moment another instance inserts that mapping.)
  • Resharding and COPY. Translation against an in-memory map is synchronous,
    so the COPY row-sharding loop, which cannot await a per-key query, can
    translate keys, and resharding's data-copy phase places existing rows on the
    correct shard.

Summary of changes

  • pgdog-config: lookup and query fields on [[sharded_tables]], with
    validation: both set together, and the lookup table
    declared in [[omnisharded_tables]] for the same database.
  • Routing translates extracted sharding keys (constants, text bind parameters,
    integers) through the in-memory map; the translated value is hashed with the
    table's configured hasher/data_type.
  • Map loading is claim-deduplicated: one client loads, concurrent statements
    wait for the result (bounded at 5s, then error). Loads are generation-checked
    so a load that straddles an invalidation is rejected and re-read.
  • Any lookup failure (load error, no server, wait timeout, value missing from
    the table) fails the statement with a retryable ERROR instead of routing
    by the untranslated value. I'm not sure if this is the best route, but the assumption I think we make here is that it is a lot safer to ensure we don't operate on the wrong shard.
  • Invalidation is a write barrier: the parser records writes to lookup tables
    (DML, TRUNCATE, COPY ... FROM, data-modifying CTEs), and the query engine
    drops the map when the write completes, meaning statement completion, or
    COMMIT for transactions (including 2PC, after phase 2), before the client is
    told the write is done. A client disconnecting mid-write invalidates on
    teardown.
  • COPY row sharding translates keys through the map; routing a COPY ... FROM
    ensures the map is loaded before data flows. Binary-format COPY into a
    lookup-translated table errors instead of misrouting.
  • Writes to lookup tables can be paused (LookupCache::pause_writes),
    rejecting them with an error so a warmed map stays exact during resharding.

@rlittlefield
rlittlefield marked this pull request as draft July 28, 2026 15:47
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.

1 participant