Database sharding: the strategies that work and the ones that cause regret

import hashlib
from typing import Any

class HashShardRouter:
    def __init__(self, shard_connections: list):
        self.shards = shard_connections
        self.shard_count = len(shard_connections)
    
    def get_shard(self, shard_key: str) -> Any:
        """Deterministically route to a shard based on the key."""
        hash_value = int(hashlib.md5(shard_key.encode()).hexdigest(), 16)
        shard_index = hash_value % self.shard_count
        return self.shards[shard_index]
    
    async def execute(self, shard_key: str, query: str, params: list):
        conn = self.get_shard(shard_key)
        return await conn.fetch(query, *params)
    
    async def execute_all_shards(self, query: str, params: list):
        """For queries without a shard key — fan out to all shards."""
        results = []
        for shard in self.shards:
            rows = await shard.fetch(query, *params)
            results.extend(rows)
        return results

router = HashShardRouter([shard0_pool, shard1_pool, shard2_pool, shard3_pool])

# All user's data goes to the same shard
await router.execute(user_id, "SELECT * FROM orders WHERE user_id = $1", [user_id])

# Cross-user queries hit all shards
total_revenue = await router.execute_all_shards(
    "SELECT SUM(amount) FROM orders WHERE date > $1", [start_date]
)

from datetime import datetime
from typing import Any

class RangeShardRouter:
    """
    Route based on a range key (e.g., created_at date).
    Good for time-series data where you query by date range.
    Risk: hotspots — the latest shard gets all writes.
    """
    def __init__(self, shard_ranges: list[tuple[datetime, datetime, Any]]):
        # [(start_date, end_date, connection)]
        self.ranges = shard_ranges
    
    def get_shard(self, date: datetime) -> Any:
        for start, end, conn in self.ranges:
            if start <= date < end:
                return conn
        raise ValueError(f"No shard for date: {date}")
    
    def get_shards_for_range(
        self,
        start: datetime,
        end: datetime
    ) -> list[Any]:
        """Return all shards that overlap with the given date range."""
        return [
            conn for range_start, range_end, conn in self.ranges
            if range_start < end and range_end > start
        ]

class TenantShardRouter:
    """
    Route by tenant ID — each tenant's data in one shard.
    Enables tenant-level isolation and easy data export/migration.
    Directory stored in a central 'routing' database.
    """
    def __init__(self, routing_db, shard_pools: dict[str, Any]):
        self.routing_db = routing_db
        self.shards = shard_pools
        self._cache: dict[str, str] = {}  # tenant_id -> shard_id
    
    async def get_shard(self, tenant_id: str) -> Any:
        if tenant_id not in self._cache:
            row = await self.routing_db.fetchrow(
                "SELECT shard_id FROM tenant_routing WHERE tenant_id = $1",
                tenant_id
            )
            if not row:
                # Assign to least loaded shard
                shard_id = await self.assign_tenant(tenant_id)
            else:
                shard_id = row["shard_id"]
            
            self._cache[tenant_id] = shard_id
        
        return self.shards[self._cache[tenant_id]]
    
    async def assign_tenant(self, tenant_id: str) -> str:
        """Assign new tenant to the shard with fewest tenants."""
        row = await self.routing_db.fetchrow("""
            SELECT shard_id, COUNT(*) as tenant_count
            FROM tenant_routing
            GROUP BY shard_id
            ORDER BY tenant_count ASC
            LIMIT 1
        """)
        shard_id = row["shard_id"] if row else "shard-0"
        
        await self.routing_db.execute(
            "INSERT INTO tenant_routing (tenant_id, shard_id) VALUES ($1, $2)",
            tenant_id, shard_id
        )
        return shard_id

async def count_all_users(router: HashShardRouter) -> int:
    """Fan-out: must query all shards and sum results."""
    counts = await asyncio.gather(*[
        shard.fetchval("SELECT COUNT(*) FROM users")
        for shard in router.shards
    ])
    return sum(counts)

async def find_user_by_email(email: str, router: HashShardRouter) -> dict | None:
    """Must check all shards — email is not the shard key."""
    results = await asyncio.gather(*[
        shard.fetchrow("SELECT * FROM users WHERE email = $1", email)
        for shard in router.shards
    ], return_exceptions=True)
    
    for result in results:
        if result and not isinstance(result, Exception):
            return dict(result)
    return None