SQLAlchemy Horizontal Sharding & Partition-Aware Routing

This walkthrough wires SQLAlchemy 2.x’s ShardedSession to route every ORM statement to the shard that owns its tenant — one of the ORM integration & partition-aware routing patterns within partitioning implementation patterns & routing. You will define one engine per shard, key the three chooser callbacks on tenant_id, control single-shard versus fan-out execution, and add a before_execute guard that rejects queries missing the partition key.

Prerequisites

Step 1 — Define one engine per shard

The horizontal sharding extension consumes a dictionary mapping shard identifiers to Engine objects. Keep the shard ids short, stable strings — they are recorded in the identity key of every loaded object, so renaming a shard id invalidates cached identity state. Inside each shard, the orders table is HASH-partitioned on tenant_id so that even a multi-tenant shard prunes to one partition per query.

from sqlalchemy import create_engine

SHARDS = {
    "shard_a": create_engine(
        "postgresql+psycopg://app_rw@pg-shard-a:5432/orders_db",
        pool_size=10, max_overflow=5, pool_pre_ping=True,
    ),
    "shard_b": create_engine(
        "postgresql+psycopg://app_rw@pg-shard-b:5432/orders_db",
        pool_size=10, max_overflow=5, pool_pre_ping=True,
    ),
    "shard_c": create_engine(
        "postgresql+psycopg://app_rw@pg-shard-c:5432/orders_db",
        pool_size=10, max_overflow=5, pool_pre_ping=True,
    ),
}

Operational note: Connection maths changes under sharding. Each application process now holds pool_size + max_overflow connections per shard, so 20 workers against 3 shards at pool_size=10, max_overflow=5 can open 900 server connections. Size pools per shard, not per app, and front busy shards with PgBouncer before raising max_connections.

SRE tip: pool_pre_ping=True costs one lightweight round-trip per checkout but prevents a failed-over shard from serving dead connections out of the pool. During a single-shard failover, only queries routed to that shard error — dashboards that aggregate errors across all shards will understate the blast radius, so label error metrics by shard id from day one.

Step 2 — Key shard_chooser and identity_chooser on tenant_id

Two of the three chooser callbacks handle object-level routing: shard_chooser decides where a new instance is flushed, and identity_chooser decides which shards to search when loading by primary key. Both must resolve deterministically from tenant_id. An explicit assignment map beats arithmetic on the shard count — it lets you move one tenant at a time, unlike modulo schemes that remap almost everything when a shard is added (the failure consistent hash routing is designed to avoid).

TENANT_SHARD_MAP = {  # loaded from config; refreshed on SIGHUP
    "acme": "shard_a", "globex": "shard_a",
    "initech": "shard_b", "umbrella": "shard_c",
}

def shard_for_tenant(tenant_id: str) -> str:
    return TENANT_SHARD_MAP[tenant_id]

def shard_chooser(mapper, instance, clause=None):
    """Called at flush time for new instances."""
    if instance is not None:
        return shard_for_tenant(instance.tenant_id)
    raise RuntimeError("cannot infer shard: no instance and no tenant hint")

def identity_chooser(mapper, primary_key, *, lazy_loaded_from,
                     execution_options, bind_arguments, **kw):
    """Called for get()-style loads; primary_key is (tenant_id, order_no)."""
    tenant_id = primary_key[0]
    return [shard_for_tenant(tenant_id)]

Operational note: shard_chooser fires during flush, so tenant_id must be populated before session.add(). Enforce this in the model constructor or a @validates hook; an instance flushed with tenant_id=None raises inside the chooser and aborts the whole flush batch, including unrelated pending objects.

DBA tip: Make tenant_id the first column of the composite primary key (PRIMARY KEY (tenant_id, order_no)) on the partitioned table. PostgreSQL requires the partition key inside every unique constraint, and identity_chooser gets the tenant for free from primary_key[0] instead of a fan-out search across all shards.

Step 3 — Route statements with execute_chooser

execute_chooser receives every select(), update() and delete() executed through the session and must return the list of shard ids to run it on. Inspect the statement’s WHERE clause for a tenant_id = :param comparison; if none is found, return every shard — a deliberate fan-out that Step 5 turns into a hard error for tables where it is never legitimate. This is application-level sharding logic expressed as a single choke point instead of scattered per-query code.

from sqlalchemy import sql
from sqlalchemy.sql import visitors

def _tenant_ids_in_where(statement):
    found = []
    for elem in visitors.iterate(statement):
        if (isinstance(elem, sql.elements.BinaryExpression)
                and elem.operator is sql.operators.eq
                and getattr(elem.left, "name", None) == "tenant_id"
                and isinstance(elem.right, sql.elements.BindParameter)):
            found.append(elem.right.effective_value)
    return found

def execute_chooser(context):  # context: ORMExecuteState
    tenants = _tenant_ids_in_where(context.statement)
    if tenants:
        return sorted({shard_for_tenant(t) for t in tenants})
    return list(SHARDS)  # fan-out fallback

Operational note: Fan-out SELECTs run serially, shard by shard, and results are concatenated — three shards at 40 ms each is 120 ms, and ORDER BY ... LIMIT is only correct per shard, not globally. Treat any fan-out on a request path as a defect: either the query is missing its tenant filter, or the access pattern belongs on a reporting replica.

SRE tip: Increment a sharded_query_fanout_total counter (labelled by table) inside the fallback branch. A sudden rise after a deploy is the earliest signal that someone shipped a query without a tenant filter — you will see it minutes before shard CPU graphs move.

Step 4 — Create the ShardedSession and control shard targeting

Wire the three choosers into a sessionmaker. Day-to-day usage is unchanged ORM code; the choosers make routing invisible. When you need to override them — administrative scripts, backfills, or queries whose tenant is known but not in the WHERE clause — force a shard explicitly with the set_shard_id() loader option rather than filtering tricks.

from sqlalchemy import select
from sqlalchemy.orm import sessionmaker
from sqlalchemy.ext.horizontal_shard import ShardedSession, set_shard_id

Session = sessionmaker(
    class_=ShardedSession,
    shards=SHARDS,
    shard_chooser=shard_chooser,
    identity_chooser=identity_chooser,
    execute_chooser=execute_chooser,
)

with Session() as session:
    # Routed write: shard_chooser sends this to shard_a ("acme")
    session.add(Order(tenant_id="acme", order_no=1001, total=99.50))
    session.commit()

    # Routed read: execute_chooser matches tenant_id = 'acme'
    stmt = select(Order).where(Order.tenant_id == "acme")
    orders = session.scalars(stmt).all()

    # Forced single-shard read: bypasses execute_chooser entirely
    stmt = select(Order).where(Order.total > 500).options(set_shard_id("shard_b"))
    big_orders = session.scalars(stmt).all()

Operational note: Objects loaded from different shards coexist in one identity map, keyed by (class, primary_key, shard_id). A commit flushes each object back to the shard it came from — but a single Session.commit() across objects from multiple shards is not atomic. Keep one logical tenant per unit of work; treat multi-shard writes as sagas with explicit compensation.

DBA tip: For backfills that must touch every shard, iterate SHARDS explicitly with set_shard_id() per batch instead of relying on fan-out. You get per-shard progress checkpoints, and a failure on shard_c does not force re-running shard_a.

Step 5 — Reject unfiltered queries with a before_execute guard

Routing solves where a query runs; it does not stop a query that forgot its tenant filter from running everywhere. A Core-level before_execute listener on each shard engine is the last line of defence: any SELECT, UPDATE, or DELETE touching a sharded table without a tenant_id comparison is rejected before it reaches PostgreSQL.

from sqlalchemy import event
from sqlalchemy.sql import Select, Update, Delete

SHARDED_TABLES = {"orders", "order_items"}

class MissingPartitionKey(Exception):
    pass

def enforce_tenant_filter(conn, clauseelement, multiparams, params,
                          execution_options):
    if not isinstance(clauseelement, (Select, Update, Delete)):
        return
    tables = {t.name for t in clauseelement.get_final_froms()} \
        if isinstance(clauseelement, Select) else {clauseelement.table.name}
    if tables & SHARDED_TABLES and not _tenant_ids_in_where(clauseelement):
        raise MissingPartitionKey(
            f"query on {tables & SHARDED_TABLES} lacks a tenant_id filter"
        )

for engine in SHARDS.values():
    event.listen(engine, "before_execute", enforce_tenant_filter)

Operational note: Roll this out in warn-only mode first — log and increment a metric instead of raising — and run it for one full business cycle. Month-end jobs, admin exports, and data-science notebooks are where legitimate unfiltered queries hide; give each an allow-listed execution_options(tenant_guard="off") escape hatch you can grep for.

SRE tip: The guard raises inside the request path, so make the exception message actionable: include the offending table and the first 200 characters of compiled SQL. Five minutes saved per page during an incident pays for the verbosity.

Verification

First, confirm routing by echoing which engine executes each statement. With echo=True set temporarily on the shard engines, a tenant-filtered query logs on exactly one engine:

INFO sqlalchemy.engine.Engine [pg-shard-a] SELECT orders.tenant_id, orders.order_no,
     orders.total FROM orders WHERE orders.tenant_id = %(tenant_id_1)s
INFO sqlalchemy.engine.Engine [pg-shard-a] {'tenant_id_1': 'acme'}

No output appears from pg-shard-b or pg-shard-c — single-shard routing is working. Next, verify partition pruning inside the chosen shard by running the exact SQL SQLAlchemy emitted through EXPLAIN:

from sqlalchemy import text

with Session() as session:
    plan = session.execute(
        text("EXPLAIN (COSTS OFF) SELECT * FROM orders "
             "WHERE tenant_id = :t AND order_no = :o"),
        {"t": "acme", "o": 1001},
        bind_arguments={"shard_id": "shard_a"},
    ).scalars().all()
    print("\n".join(plan))

Expected output — one partition scanned, not all eight:

Index Scan using orders_p3_pkey on orders_p3 orders
  Index Cond: ((tenant_id = 'acme'::text) AND (order_no = 1001))

Finally, confirm the guard fires: session.scalars(select(Order)) (no filter, no set_shard_id) must raise MissingPartitionKey rather than fanning out.

Failure mode table

Failure mode Root cause SRE mitigation
Latency spike from queries silently fanning out to all shards execute_chooser finds no tenant_id equality — the filter sits on a joined table alias or inside an IN (subquery) the traversal does not match Alert on the sharded_query_fanout_total counter; rewrite the query with a direct tenant_id predicate or attach set_shard_id() where the tenant is known
Order loads return None for rows that exist Tenant remapped to a new shard while TENANT_SHARD_MAP in a long-lived worker still points at the old one Version the tenant map and export the version as a gauge per process; on remap, bump the version and recycle workers that lag; verify with a cross-shard count of the tenant’s rows
One shard exhausts its connection pool while others idle Hot tenant concentration on a single shard multiplies checkouts against one pool Track pool.checkedout() per shard; rebalance the heaviest tenants (see handling hot keys in list-partitioned tables) and cap per-tenant concurrency at the app layer

FAQ

Does ShardedSession work with SQLAlchemy's asyncio extension?

The horizontal sharding extension is built around the synchronous Session. For asyncio applications, resolve the shard from tenant_id in application code before creating the session, then build an AsyncSession bound to that shard’s AsyncEngine. You lose automatic fan-out, but single-tenant routing — the common case — works cleanly, and cross-shard reads can be issued as explicit asyncio.gather calls across per-shard sessions.

Do I still need table partitioning inside each shard if routing already isolates tenants?

Yes, when shards host many tenants each. Sharding bounds the working set per server; hash or list partitioning by tenant_id inside each shard keeps indexes small, enables per-tenant DETACH-and-archive operations, and lets PostgreSQL prune to one partition per query — which is exactly what the Verification section’s EXPLAIN confirms. The two layers solve different problems and compose well.

How do I add a new shard without breaking the shard_chooser mapping?

Never derive the shard from a modulo over the live shard count — adding a shard would remap most tenants. Use an explicit tenant-to-shard assignment table, or a consistent hash ring so only a small key range moves. Deploy the new engine into SHARDS first, then copy the moving tenants’ rows, then flip each tenant’s entry in TENANT_SHARD_MAP. The chooser functions themselves never change.

Back to Orm Integration Partition Routing