Database Design

from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker, scoped_session
from contextlib import contextmanager

class DatabasePool:
    """Production-ready database connection pool"""
    
    def __init__(self, database_url: str):
        self.engine = create_engine(
            database_url,
            pool_size=20,           # Persistent connections
            max_overflow=30,        # Extra connections under load
            pool_timeout=30,        # Wait time for connection
            pool_recycle=3600,      # Recycle connections hourly
            pool_pre_ping=True,     # Verify connection before use
            echo=False              # Set True for SQL logging
        )
        
        self._session_factory = sessionmaker(bind=self.engine)
        self.Session = scoped_session(self._session_factory)
    
    @contextmanager
    def get_session(self):
        """Context manager for database sessions"""
        session = self.Session()
        try:
            yield session
            session.commit()
        except Exception as e:
            session.rollback()
            raise e
        finally:
            session.close()
    
    def execute_with_retry(self, query, params=None, max_retries=3):
        """Execute query with automatic retry on transient failures"""
        import time
        
        for attempt in range(max_retries):
            try:
                with self.get_session() as session:
                    result = session.execute(query, params or {})
                    return result.fetchall()
            except Exception as e:
                if "connection" in str(e).lower() and attempt < max_retries - 1:
                    time.sleep(2 ** attempt)  # Exponential backoff
                    continue
                raise


# Usage
db = DatabasePool("postgresql://user:pass@localhost/mydb")

with db.get_session() as session:
    users = session.query(User).filter(User.active == True).all()

import random
from datetime import datetime, timedelta
from typing import Optional, List, Any
from enum import Enum

class QueryType(Enum):
    READ = "read"
    WRITE = "write"

class ReadWriteRouter:
    """
    Routes queries to primary or replicas based on query type
    and session consistency requirements.
    """
    
    def __init__(self, primary_pool, replica_pools: list):
        self.primary = primary_pool
        self.replicas = replica_pools
        self.user_writes = {}  # Track recent writes per user
        self.replication_lag_threshold = 5  # seconds
    
    def get_connection(self, query_type: QueryType, 
                       user_id: Optional[str] = None):
        """
        Get appropriate connection based on query type and user session.
        Ensures read-your-writes consistency.
        """
        if query_type == QueryType.WRITE:
            # All writes go to primary
            self._record_write(user_id)
            return self.primary
        
        # For reads, check if user recently wrote
        if user_id and self._should_read_from_primary(user_id):
            return self.primary
        
        # Otherwise, load balance across replicas
        return self._get_healthy_replica()
    
    def _record_write(self, user_id: Optional[str]):
        if user_id:
            self.user_writes[user_id] = datetime.utcnow()
    
    def _should_read_from_primary(self, user_id: str) -> bool:
        """Read from primary if user wrote recently"""
        last_write = self.user_writes.get(user_id)
        if not last_write:
            return False
        
        time_since_write = (datetime.utcnow() - last_write).total_seconds()
        return time_since_write < self.replication_lag_threshold
    
    def _get_healthy_replica(self):
        """Get a healthy replica with lowest lag"""
        healthy = [r for r in self.replicas if self._is_healthy(r)]
        
        if not healthy:
            # Fallback to primary if no healthy replicas
            return self.primary
        
        # Random selection (could be weighted by lag)
        return random.choice(healthy)
    
    def _is_healthy(self, replica) -> bool:
        try:
            # Check replication lag
            result = replica.execute("SELECT pg_last_xact_replay_timestamp()")
            lag = (datetime.utcnow() - result).total_seconds()
            return lag < 30  # Healthy if lag < 30 seconds
        except:
            return False


# ORM Integration Example
class UserRepository:
    def __init__(self, router: ReadWriteRouter):
        self.router = router
    
    def find_by_id(self, user_id: str, session_user: str = None) -> dict:
        """Read operation - goes to replica"""
        conn = self.router.get_connection(QueryType.READ, session_user)
        return conn.execute(
            "SELECT * FROM users WHERE id = %s", 
            (user_id,)
        ).fetchone()
    
    def update(self, user_id: str, data: dict, session_user: str) -> bool:
        """Write operation - goes to primary"""
        conn = self.router.get_connection(QueryType.WRITE, session_user)
        conn.execute(
            "UPDATE users SET name = %s WHERE id = %s",
            (data['name'], user_id)
        )
        return True


# Usage
router = ReadWriteRouter(primary_pool, [replica1, replica2])
repo = UserRepository(router)

# This write goes to primary
repo.update("user_123", {"name": "Alice"}, session_user="user_123")

# This read goes to primary (recent write by same user)
user = repo.find_by_id("user_123", session_user="user_123")

# This read goes to replica (different user)
user = repo.find_by_id("user_123", session_user="user_456")

import hashlib
from typing import Dict, Any, List

class ShardRouter:
    """
    Routes queries to appropriate shard based on shard key.
    Uses consistent hashing for even distribution.
    """
    
    def __init__(self, shards: Dict[int, Any]):
        self.shards = shards  # shard_id -> connection pool
        self.num_shards = len(shards)
    
    def get_shard(self, shard_key: str) -> Any:
        """Get shard connection for a given key"""
        shard_id = self._compute_shard_id(shard_key)
        return self.shards[shard_id]
    
    def _compute_shard_id(self, key: str) -> int:
        """Consistent hash to determine shard"""
        hash_value = int(hashlib.md5(key.encode()).hexdigest(), 16)
        return hash_value % self.num_shards
    
    def execute_on_shard(self, shard_key: str, query: str, params: tuple):
        """Execute query on appropriate shard"""
        shard = self.get_shard(shard_key)
        return shard.execute(query, params)
    
    def scatter_gather(self, query: str, params: tuple = None) -> List[Any]:
        """Execute query on ALL shards and combine results"""
        results = []
        for shard_id, shard in self.shards.items():
            try:
                result = shard.execute(query, params or ())
                results.extend(result.fetchall())
            except Exception as e:
                print(f"Shard {shard_id} failed: {e}")
        return results


class ShardedUserRepository:
    """User repository with sharding by user_id"""
    
    def __init__(self, shard_router: ShardRouter):
        self.router = shard_router
    
    def create_user(self, user_id: str, data: dict) -> bool:
        """Create user on appropriate shard"""
        query = """
            INSERT INTO users (id, name, email, created_at)
            VALUES (%s, %s, %s, NOW())
        """
        self.router.execute_on_shard(
            shard_key=user_id,
            query=query,
            params=(user_id, data['name'], data['email'])
        )
        return True
    
    def find_by_id(self, user_id: str) -> dict:
        """Find user - knows exactly which shard"""
        query = "SELECT * FROM users WHERE id = %s"
        result = self.router.execute_on_shard(user_id, query, (user_id,))
        return result.fetchone()
    
    def find_by_email(self, email: str) -> dict:
        """Find by email - must scatter-gather all shards"""
        query = "SELECT * FROM users WHERE email = %s"
        results = self.router.scatter_gather(query, (email,))
        return results[0] if results else None
    
    def get_total_count(self) -> int:
        """Count across all shards"""
        query = "SELECT COUNT(*) as count FROM users"
        results = self.router.scatter_gather(query)
        return sum(r['count'] for r in results)


# Setup
shards = {
    0: create_pool("postgresql://localhost:5432/users_shard_0"),
    1: create_pool("postgresql://localhost:5433/users_shard_1"),
    2: create_pool("postgresql://localhost:5434/users_shard_2"),
    3: create_pool("postgresql://localhost:5435/users_shard_3"),
}

router = ShardRouter(shards)
repo = ShardedUserRepository(router)

# Usage
repo.create_user("user_123", {"name": "Alice", "email": "alice@example.com"})
user = repo.find_by_id("user_123")  # Goes to one shard
user = repo.find_by_email("alice@example.com")  # Queries ALL shards