Query Router

Component 10 of 11 (Distribution & High Availability, 2 of 2) in the strata-database-engine ecosystem. Repo: github.com/juanmmm21/distributed-query-router

A sharding and query routing layer, written in Python: hash or range partitioning of a table across shards, writes always routed to the current shard leader, reads fanned out to replicas according to a requested consistency level, and partial per-shard results merged back into a single answer — with distributed SUM/COUNT/AVG/MIN/MAX reconstructed correctly from partitions of very different sizes, and any node that fails or times out during a fan-out marked explicitly rather than silently dropped.

What it is and what problem it solves

A single database node — even a replicated one, like the cluster raft-replication-log implements — has a hard ceiling: a table that does not fit in one node's storage, or a write workload that saturates one node's throughput, cannot be solved by adding more replicas of that same node. The standard fix is horizontal partitioning ("sharding"): split the table's rows across several independent shards, each one a small replicated cluster on its own, and put a routing layer in front that knows, for any given query, which shard(s) it needs and how to talk to them.

That routing layer has to answer three questions correctly, and get all three wrong in ways that are easy to hide behind a demo that never tests failure:

1. Where does this write go? Exactly one node — the current leader of the shard that owns the write's key — may accept it. Routing a write to the wrong node either loses it silently or, worse, creates two nodes that both think they hold the authoritative copy.
2. Where can this read go? Depends on how fresh the caller needs the answer to be: the leader only (guaranteed latest), or any replica (faster, spreads load, but may be slightly stale).
3. How do partial answers from several shards become one answer? A COUNT(*) is just a sum of per-shard counts. An AVG(amount) is not the average of per-shard averages — that silently produces the wrong number the moment two shards hold different row counts, which in a real sharded table is the normal case, not the exception.

And underneath all three: what happens when a shard's leader changes while the router is mid-query, or when one of several shards being fanned out to simply does not answer in time? A router that quietly returns whatever came back, without saying so, is worse than one that fails loudly — a caller that doesn't know a result is incomplete cannot decide whether that matters for its use case.

Role in strata-database-engine

raft-replication-log                cost-based-query-optimizer
        │                                       │
        │ "current leader per shard"            │ physical plan run on
        │ (no direct import --                  │ each node's own
        │  LeadershipOracle Protocol)            │ vectorized-execution-engine
        ▼                                       ▼
┌───────────────────────────────────────────────────────┐
│               distributed-query-router (this repo)      │
│  sharding (hash/range) · write routing to leader          │
│  read fan-out by consistency level · aggregate merge        │
│  explicit partial-result marking on node failure/timeout      │
└──────────────────────────────┬────────────────────────────────┘
                                │ MergedResult, as if from one logical node
                                ▼
                             nanosql

This repo imports neither raft-replication-log nor cost-based-query-optimizer / vectorized-execution-engine — the ecosystem rule is that cross-repo integration only happens inside nanosql. Instead it redefines, locally and independently, the minimal shape it needs from each: a LeadershipOracle protocol shaped exactly like raft_replication_log.cluster.Cluster.current_leader(), and an opaque PartialResult type standing in for whatever a real node's physical plan execution would return. See Data format / interface exposed to the ecosystem for the exact contracts.

Goal / skills demonstrated

• Two interchangeable, independently testable sharding schemes (HashShardingScheme, RangeShardingScheme) behind one ShardingScheme protocol, with the classic trade-off between them (hash: uniform load, no range pruning; range: range pruning, possible hotspots) documented rather than glossed over.
• A deterministic sharding hash (hashlib.blake2b, never Python's seed-randomized hash()) — the same lesson lsm-tree-engine's bloom filter already had to learn, applied here to a different structure.
• A mutable shard topology (ShardDirectory) that is never cached downstream of the router, so a live shard relocation takes effect on the very next call with no invalidation step — verified directly, including a relocation triggered as a side effect during an in-flight fan-out.
• Correct reconstruction of AVG from per-shard partial state: the algebraic (Gray et al.) decomposition into (sum, count) pairs merged separately and divided only once, at the very end — never the average of per-shard averages, which is wrong the instant shard sizes differ.
• A MergedResult/FanoutResponse design where "some node didn't answer" is a first-class, always-propagated fact (is_partial, failures) rather than a value silently missing from a return type that looks identical whether everything succeeded or not.
• A stateless router: DistributedRouter holds no state of its own between calls (only injected, externally-owned dependencies), which is what makes it safe to retry an entire query from scratch after the router process itself dies mid-fan-out — verified directly, including for an aggregated query, confirming a retry never double-counts.
• A property-based test suite (20 seeds) that builds random sharded datasets, randomly fails a subset of shards per fan-out, and checks the merged result against an independent reference model (plain Python sum/len/min/max over the surviving rows) — never reusing merge.py's own logic to check itself.

How it works

Sharding

ShardingScheme.shard_for_key(key) -> ShardId is the single decision point every other module builds on. HashShardingScheme(num_shards) hashes the key with blake2b and reduces modulo num_shards; RangeShardingScheme(boundaries, shard_ids) uses bisect_right over a sorted tuple of upper bounds, so a shard's range is always known statically and query pruning by key range is possible. Both are configured explicitly by the caller — there is no implicit default scheme, matching the ecosystem's preference for explicit, testable configuration over hidden defaults.

ShardDirectory is a separate, mutable piece: it tracks which physical nodes currently replicate each shard (ShardReplicaSet), independent of who among them is the leader. relocate_shard() swaps a shard's replica set in place. Every routing and fan-out call reads the directory fresh — nothing downstream holds a cached copy — so a relocation is visible on the very next call, including calls already in flight for other shards within the same fan-out.

Write routing

routing.resolve_write_target() looks up the shard for a key, asks the injected LeadershipOracle who its current leader is, and rejects the request (NoLeaderError) unless that leader is both non-None and still a member of the shard's current replica set — the second check matters specifically during a shard relocation, where the oracle can briefly still report a leader that the topology has already retired. execute_write() then sends the write to exactly that node and returns its acknowledgement or propagates the node's failure untouched: a write is never silently retried against a different node, because that would violate "exactly the leader accepts writes."

Read fan-out and consistency levels

A ReadRequest carries a ConsistencyLevel: STRONG requires the shard's current leader (fails outright if unknown, just like a write); BOUNDED_STALENESS prefers the leader but falls back to any replica if none is known, trading strict freshness for availability only when freshness is not free; EVENTUAL picks any replica, for maximum read-load spreading. fanout.fan_out_read() resolves the set of shards the request touches — all of them, or only those reachable from ReadRequest.shard_keys if the caller already knows which keys it needs — and, for each, selects a target and calls the injected NodeExecutor. A NodeTimeoutError or NodeUnavailableError from any single shard is caught per shard, recorded as an explicit NodeFailure, and does not abort the rest of the fan-out — a FanoutResponse.is_partial flag is True if and only if at least one shard failed.

Aggregate merge

merge.merge_partial_results() has two paths. Without aggregates, it is a concatenation of each shard's raw rows — a SELECT over a sharded table is a union of its partitions. With aggregates, each shard's PartialResult.grouped_aggregates carries one AggregateAccumulator per output column per GROUP BY group; the merge combines same-column accumulators pairwise (merge_accumulators) and only calls finalize_accumulator once, at the very end, per group. SUM/COUNT/MIN/MAX are straightforward (add, add, min, max); AVG is the one that needs care — its accumulator carries sum and count separately (identical shape to the SUM/COUNT accumulators) and only divides once both are fully merged, which is what makes the result correct regardless of how unevenly the underlying rows are distributed across shards. MergedResult.is_partial/.failures are propagated from the FanoutResponse completely unmodified — the merge step has no mechanism to "forget" that some of its input was incomplete.

Architecture

src/distributed_query_router/
├── __init__.py
├── models.py         # NodeId/ShardId/ShardKey aliases, ConsistencyLevel, AggregateFunction,
│                      # AggregateAccumulator, PartialResult, FanoutResponse, MergedResult, ...
├── errors.py           # RouterError hierarchy (NoLeaderError, NodeTimeoutError, ...)
├── protocols.py         # LeadershipOracle / ShardingScheme / NodeExecutor Protocols,
│                        # plus InMemoryLeadershipOracle / SimulatedNodeExecutor defaults
├── sharding.py            # HashShardingScheme, RangeShardingScheme, ShardDirectory
├── routing.py               # resolve_write_target, execute_write
├── fanout.py                  # select_read_target, fan_out_read
├── merge.py                    # merge_accumulators, finalize_accumulator, merge_partial_results
├── pipeline.py                   # DistributedRouter: stateless facade over the four modules above
└── __main__.py                    # demo / relocate-demo / failure-demo / benchmark CLI

Per-domain logic (routing.py, fanout.py, merge.py) is factored as free functions taking their dependencies (ShardDirectory, LeadershipOracle, NodeExecutor) as explicit arguments rather than methods on a stateful object — the same separation raft-replication-log uses for election.py/replication.py/snapshot.py, and for the same reason: it keeps each piece unit-testable in isolation, without needing to construct a full DistributedRouter. pipeline.DistributedRouter is a thin convenience facade over those functions for callers (and the CLI) that want one object to hold onto.

Requirements and installation

• Python >=3.11, no runtime dependencies.

git clone https://github.com/juanmmm21/distributed-query-router.git
cd distributed-query-router
python3 -m venv .venv && source .venv/bin/activate
pip install -e ".[dev]"  # or: pip install -e . pytest mypy ruff

Usage

CLI

# write routing to the shard leader + SUM/COUNT/AVG merged across 3 shards
python -m distributed_query_router demo

# a shard's leader changes between two writes; the second write is routed
# to the new leader with no explicit cache invalidation
python -m distributed_query_router relocate-demo

# one shard's leader times out mid-fan-out; the merged result comes back
# marked is_partial=True with the exact shard/node/reason recorded
python -m distributed_query_router failure-demo

# fan-out + merge throughput for 1/4/16/64 shards
python -m distributed_query_router benchmark --sizes 1 4 16 64 --queries 2000

Programmatic usage

import random

from distributed_query_router.models import (
    AggregateFunction, AggregateSpec, ConsistencyLevel, MergeSpec,
    ReadRequest, ShardReplicaSet, WriteRequest,
)
from distributed_query_router.pipeline import DistributedRouter
from distributed_query_router.protocols import InMemoryLeadershipOracle, SimulatedNodeExecutor
from distributed_query_router.sharding import HashShardingScheme, ShardDirectory

scheme = HashShardingScheme(num_shards=4)
directory = ShardDirectory({
    shard_id: ShardReplicaSet(shard_id, (f"{shard_id}-a", f"{shard_id}-b"))
    for shard_id in scheme.all_shard_ids()
})
oracle = InMemoryLeadershipOracle({sid: f"{sid}-a" for sid in scheme.all_shard_ids()})
executor = SimulatedNodeExecutor(read_responses={...}, rng=random.Random(1))

router = DistributedRouter(directory, oracle, scheme, executor, random.Random(1))

ack = router.execute_write(WriteRequest(shard_key="customer-42", payload=b"..."))
# ack.shard_id, ack.node_id -- which shard and which leader accepted it

result = router.execute_read(ReadRequest(
    payload=b"SELECT COUNT(*) FROM orders",
    consistency=ConsistencyLevel.EVENTUAL,
    merge_spec=MergeSpec(aggregates=(AggregateSpec(AggregateFunction.COUNT, "n"),)),
))
if result.is_partial:
    ...  # decide whether an incomplete count is acceptable for this caller

InMemoryLeadershipOracle/SimulatedNodeExecutor are the default in-memory implementations used by tests and the CLI demos — a real assembly inside nanosql supplies its own LeadershipOracle backed by an actual raft_replication_log.cluster.Cluster per shard, and its own NodeExecutor that dispatches to a real node's vectorized-execution-engine pipeline instead of returning a pre-loaded value.

Data format / interface exposed to the ecosystem

This module never imports raft-replication-log or cost-based-query-optimizer / vectorized-execution-engine — the cross-repo integration rule for this ecosystem is that it only happens inside nanosql. Instead it redefines the minimal shapes it needs:

class LeadershipOracle(Protocol):
    def current_leader(self, shard_id: ShardId) -> NodeId | None: ...

shaped identically to raft_replication_log.cluster.Cluster.current_leader(); a real assembly passes one Cluster instance per shard wrapped in an adapter implementing this Protocol.

@dataclass(frozen=True, slots=True)
class PartialResult:
    shard_id: ShardId
    node_id: NodeId
    rows: tuple[Row, ...] = ()
    grouped_aggregates: Mapping[GroupKey, Mapping[str, AggregateAccumulator]] = field(default_factory=dict)

PartialResult is opaque from this module's point of view — it never re-executes or reinterprets a node's query, only fuses what the node already computed. In a real assembly, a NodeExecutor implementation would run the physical plan produced by cost-based-query-optimizer through vectorized-execution-engine on the target node and translate its output (raw rows, or HashAggregateOperator output) into this shape. The two things nanosql consumes from this module are WriteAck (which shard, which leader, for a write) and MergedResult (rows, is_partial, failures) for a read — both presented as if they came from a single logical node.

Development

pytest
ruff check .
ruff format --check .
mypy --strict src/

The test suite (42 tests) covers: sharding schemes and the mutable shard directory in isolation (hash-scheme coverage and determinism across processes, range-scheme boundary edge cases, relocation taking effect immediately) — 10 tests; write routing (leader resolution, rejection of a leader reported outside the current topology, failure propagation without silent retry, and a 200-key sweep asserting the "never route to a non-leader" invariant) — 6 tests; read fan-out (target selection for all three consistency levels, shard-key pruning, per-shard timeout/failure recorded without aborting the rest of the fan-out, missing-leader recorded with no fabricated node id) — 9 tests; aggregate merge (SUM/COUNT/AVG/MIN/MAX merge and finalize, the weighted-vs-naive AVG case explicitly, GROUP BY with multiple groups, and is_partial propagating through a successful merge of incomplete data) — 11 tests; full end-to-end scenarios matching the DoD's two named cases (a shard relocated between two writes, a relocation triggered as a side effect mid-fan-out, and a router "crash" — an untyped exception distinct from the domain's own NodeTimeoutError/NodeUnavailableError — proven safe to retry for both a plain and an aggregated query) — 4 tests; and a property-based suite (20 seeds, each building a random sharded dataset of up to 200 rows across 2–6 shards with a random subset of shards failing, plus a second suite of 20 seeds × 300 randomized leader-change steps) checking the merged result against an independent reference model (plain Python sum/len/min/max, never reusing merge.py) and the "never route to a non-leader" invariant — 2 tests generating several thousand assertions between them.

No bug was found by the test suite before publishing this time — unlike some of this ecosystem's other components (see raft-replication-log's or cost-based-query-optimizer's README for two real examples), every invariant checked here passed on its first real run against the implementation. Stated plainly rather than implied, since the ecosystem's convention is to report this honestly either way.

Benchmarks

Measured with the benchmark subcommand on the reference development machine (fan-out + merge of a single-row COUNT(*)-style aggregate per shard, STRONG consistency, 2000 queries per size):

Throughput scales down roughly linearly with shard count, as expected: each additional shard adds one more NodeExecutor.execute_read call (and one more accumulator merge) to every single fan-out, with no batching across shards in this simulated in-memory setting. As with raft-replication-log's benchmark, these numbers characterize the shape of the routing/merge overhead itself (Python function-call and dict-merge cost per shard) rather than a number comparable to a real deployment, whose latency would be dominated by actual network round trips to each node instead.

Troubleshooting

• NoLeaderError when calling execute_write/a STRONG read: the injected LeadershipOracle returned None, or returned a node no longer present in that shard's current ShardReplicaSet (the second case happens specifically during a relocation window). Update the oracle (a real assembly's oracle reflects the underlying Cluster's actual state) or use BOUNDED_STALENESS/EVENTUAL if strict freshness is not required for that read.
• ShardNotFoundError: the ShardId a ShardingScheme produced is not present in the ShardDirectory passed to the router — usually means the directory was built against a different num_shards/boundaries than the scheme currently in use. Both must be kept in sync by whoever owns the topology.
• A MergedResult comes back with is_partial=True and you expected a clean result: check MergedResult.failures — each entry names the shard, the node that was targeted (or None if no leader could even be selected), and whether it was a timeout or an unavailable failure. This is by design never silent; a caller that wants to retry should re-issue the whole query, since DistributedRouter holds no partial progress to resume from (see tests/test_scenarios.py::test_router_crash_mid_query_is_safe_to_retry).
• AVG looks wrong when computed by hand from per-shard averages: that is expected to disagree — see Aggregate merge. AVG must be reconstructed from (sum, count) pairs, never averaged directly across shards of different sizes; merge.finalize_accumulator does this correctly, a naive mean(per_shard_averages) will not.
• InvalidShardingSchemeError from RangeShardingScheme: boundaries must be strictly sorted and len(shard_ids) == len(boundaries) + 1 — a range scheme needs exactly one more shard than it has internal boundaries (the last shard covers everything above the final boundary).

License

MIT — see LICENSE.