Tenant Isolation Strategies.

# tenant_isolation/context.py

"""
Tenant context management for multi-tenant applications.

The tenant context must flow through every layer:
- HTTP request → Service → Database → Cache → Queue → Background job
"""

from contextvars import ContextVar
from dataclasses import dataclass
from typing import Optional
import functools


# Thread-safe tenant context using contextvars
_tenant_context: ContextVar[Optional['TenantContext']] = ContextVar(
    'tenant_context', 
    default=None
)


@dataclass(frozen=True)
class TenantContext:
    """
    Immutable tenant context that flows through the application.
    
    Frozen to prevent accidental modification.
    """
    tenant_id: str
    tenant_name: str
    plan: str  # 'free', 'pro', 'enterprise'
    region: str  # For data residency
    features: frozenset  # Enabled features
    
    # Security context
    isolation_level: str  # 'shared', 'schema', 'database'
    encryption_key_id: Optional[str] = None  # For tenant-specific encryption
    is_active: bool = True  # False when the tenant is suspended


def get_current_tenant() -> TenantContext:
    """
    Get the current tenant context.
    
    Raises TenantContextError if not in tenant context.
    """
    ctx = _tenant_context.get()
    if ctx is None:
        raise TenantContextError("No tenant context set. This is a bug!")
    return ctx


def get_current_tenant_id() -> str:
    """Convenience method to get just the tenant ID."""
    return get_current_tenant().tenant_id


def set_tenant_context(ctx: TenantContext) -> None:
    """Set the tenant context for the current execution."""
    _tenant_context.set(ctx)


def clear_tenant_context() -> None:
    """Clear the tenant context."""
    _tenant_context.set(None)


class TenantContextError(Exception):
    """Raised when tenant context is missing or invalid."""
    pass


# Context manager for tenant scope
class tenant_scope:
    """
    Context manager for executing code in a tenant's context.
    
    Usage:
        with tenant_scope(tenant_ctx):
            # All code here runs in tenant context
            users = user_service.get_users()  # Automatically filtered
    """
    
    def __init__(self, tenant: TenantContext):
        self.tenant = tenant
        self.previous = None
    
    def __enter__(self):
        self.previous = _tenant_context.get()
        _tenant_context.set(self.tenant)
        return self.tenant
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        _tenant_context.set(self.previous)
        return False


# Decorator for tenant-required functions
def requires_tenant(func):
    """
    Decorator that ensures function runs within tenant context.
    
    Usage:
        @requires_tenant
        def get_user_data():
            tenant = get_current_tenant()
            # ...
    """
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        if _tenant_context.get() is None:
            raise TenantContextError(
                f"Function {func.__name__} requires tenant context"
            )
        return func(*args, **kwargs)
    return wrapper


# Async version
def requires_tenant_async(func):
    """Async version of requires_tenant decorator."""
    @functools.wraps(func)
    async def wrapper(*args, **kwargs):
        if _tenant_context.get() is None:
            raise TenantContextError(
                f"Function {func.__name__} requires tenant context"
            )
        return await func(*args, **kwargs)
    return wrapper

# tenant_isolation/middleware.py

"""
Middleware for extracting and setting tenant context from HTTP requests.
"""

from typing import Optional
from fastapi import Request
from fastapi.responses import JSONResponse
from starlette.middleware.base import BaseHTTPMiddleware
import logging

logger = logging.getLogger(__name__)


class TenantMiddleware(BaseHTTPMiddleware):
    """
    Middleware that extracts tenant context from requests.
    
    Tenant can be identified via:
    1. JWT claim: token.tenant_id = 'acme'  (authoritative)
    2. Subdomain: acme.myapp.com → 'acme'   (hint; must match the JWT)
    3. Header: X-Tenant-ID: acme            (hint; must match the JWT)
    4. Query param, for webhooks only        (after signature verification)

    Assumes an auth middleware has already verified the JWT and set
    request.state.user. Middleware exceptions bypass FastAPI's
    HTTPException handlers, so errors are returned as responses.
    """
    
    def __init__(self, app, tenant_service):
        super().__init__(app)
        self.tenant_service = tenant_service
    
    async def dispatch(self, request: Request, call_next):
        # Skip tenant resolution for public endpoints
        if self._is_public_endpoint(request.url.path):
            return await call_next(request)
        
        # Extract tenant identifier
        tenant_id = await self._extract_tenant_id(request)
        
        if not tenant_id:
            return JSONResponse(
                status_code=400,
                content={"detail": "Tenant identification required"}
            )
        
        # Load tenant context from database/cache
        tenant_ctx = await self.tenant_service.get_tenant_context(tenant_id)
        
        if not tenant_ctx:
            return JSONResponse(
                status_code=404,
                content={"detail": f"Tenant '{tenant_id}' not found"}
            )
        
        if not tenant_ctx.is_active:
            return JSONResponse(
                status_code=403,
                content={"detail": "Tenant account is suspended"}
            )
        
        # Set context for this request
        with tenant_scope(tenant_ctx):
            # Add tenant to request state for logging
            request.state.tenant_id = tenant_id
            
            # Process request
            response = await call_next(request)
            
            # Add tenant to response headers (for debugging)
            response.headers["X-Tenant-ID"] = tenant_id
            
            return response
    
    async def _extract_tenant_id(self, request: Request) -> Optional[str]:
        """
        Extract tenant ID from request using multiple strategies.
        Never trust a client-supplied value on its own (see Mistake 2).
        """
        
        # Strategy 1: JWT claim from the authenticated user (authoritative)
        user = getattr(request.state, "user", None)
        claimed = user.tenant_id if user else None
        
        # Strategies 2-3: subdomain or header are only hints
        host = request.headers.get("host", "").split(":")[0]
        subdomain = host.split(".")[0] if host.count(".") >= 2 else None
        if subdomain in ("www", "api", "app"):
            subdomain = None
        hint = subdomain or request.headers.get("X-Tenant-ID")
        
        if claimed:
            if hint and hint != claimed:
                return None  # Mismatch: reject, never pick the client's value
            return claimed
        
        # Strategy 4: unauthenticated webhooks. The caller must verify the
        # webhook signature and that it belongs to this tenant before use.
        if request.url.path.startswith("/webhooks/"):
            return request.query_params.get("tenant_id")
        
        return None
    
    def _is_public_endpoint(self, path: str) -> bool:
        """Check if endpoint doesn't require tenant context."""
        public_paths = [
            "/health",
            "/metrics",
            "/api/auth/login",
            "/api/auth/register",
            "/api/tenants/create",  # Creating new tenant
        ]
        return any(path.startswith(p) for p in public_paths)


class TenantService:
    """
    Service for loading and caching tenant configuration.
    """
    
    def __init__(self, db, cache):
        self.db = db
        self.cache = cache
    
    async def get_tenant_context(self, tenant_id: str) -> Optional[TenantContext]:
        """
        Load tenant context, with caching.
        """
        
        # Check cache first
        cache_key = f"tenant:{tenant_id}"
        cached = await self.cache.get(cache_key)
        
        if cached:
            return TenantContext(**cached)
        
        # Load from database
        tenant = await self.db.fetchrow(
            """
            SELECT 
                id, name, plan, region, features,
                isolation_level, encryption_key_id, is_active
            FROM tenants
            WHERE id = $1
            """,
            tenant_id
        )
        
        if not tenant:
            return None
        
        ctx = TenantContext(
            tenant_id=tenant["id"],
            tenant_name=tenant["name"],
            plan=tenant["plan"],
            region=tenant["region"],
            features=frozenset(tenant["features"]),
            isolation_level=tenant["isolation_level"],
            encryption_key_id=tenant["encryption_key_id"],
            is_active=tenant["is_active"]
        )
        
        # Cache for 5 minutes
        await self.cache.set(cache_key, ctx.__dict__, ttl=300)
        
        return ctx

# tenant_isolation/repository.py

"""
Repository pattern with automatic tenant filtering.

All database access goes through repositories that automatically
apply tenant filtering.
"""

from typing import TypeVar, Generic, List, Optional
from dataclasses import dataclass
import logging

logger = logging.getLogger(__name__)

T = TypeVar('T')


class TenantAwareRepository(Generic[T]):
    """
    Base repository that automatically filters by tenant.
    
    All queries include tenant_id filter.
    All inserts include tenant_id.
    """
    
    def __init__(self, db, table_name: str, model_class: type):
        self.db = db
        self.table_name = table_name
        self.model_class = model_class
    
    async def find_by_id(self, id: str) -> Optional[T]:
        """
        Find record by ID within current tenant.
        """
        tenant_id = get_current_tenant_id()
        
        row = await self.db.fetchrow(
            f"""
            SELECT * FROM {self.table_name}
            WHERE id = $1 AND tenant_id = $2
            """,
            id, tenant_id
        )
        
        if not row:
            return None
        
        return self.model_class(**row)
    
    async def find_all(
        self,
        filters: dict = None,
        limit: int = 100,
        offset: int = 0
    ) -> List[T]:
        """
        Find all records within current tenant.
        """
        tenant_id = get_current_tenant_id()
        
        # Build query with tenant filter always first
        query = f"SELECT * FROM {self.table_name} WHERE tenant_id = $1"
        params = [tenant_id]
        param_idx = 2
        
        if filters:
            # Column names can't be bound as parameters: keys must come
            # from an allow-list of known columns, never from user input
            for key, value in filters.items():
                query += f" AND {key} = ${param_idx}"
                params.append(value)
                param_idx += 1
        
        query += f" LIMIT ${param_idx} OFFSET ${param_idx + 1}"
        params.extend([limit, offset])
        
        rows = await self.db.fetch(query, *params)
        return [self.model_class(**row) for row in rows]
    
    async def create(self, data: dict) -> T:
        """
        Create record within current tenant.
        """
        tenant_id = get_current_tenant_id()
        
        # Always set tenant_id
        data["tenant_id"] = tenant_id
        
        columns = ", ".join(data.keys())
        placeholders = ", ".join(f"${i+1}" for i in range(len(data)))
        
        row = await self.db.fetchrow(
            f"""
            INSERT INTO {self.table_name} ({columns})
            VALUES ({placeholders})
            RETURNING *
            """,
            *data.values()
        )
        
        return self.model_class(**row)
    
    async def update(self, id: str, data: dict) -> Optional[T]:
        """
        Update record within current tenant.
        """
        tenant_id = get_current_tenant_id()
        
        # Remove tenant_id from update data (can't change tenant)
        data.pop("tenant_id", None)
        
        set_clauses = ", ".join(
            f"{key} = ${i+1}" for i, key in enumerate(data.keys())
        )
        
        row = await self.db.fetchrow(
            f"""
            UPDATE {self.table_name}
            SET {set_clauses}
            WHERE id = ${len(data)+1} AND tenant_id = ${len(data)+2}
            RETURNING *
            """,
            *data.values(), id, tenant_id
        )
        
        if not row:
            return None
        
        return self.model_class(**row)
    
    async def delete(self, id: str) -> bool:
        """
        Delete record within current tenant.
        """
        tenant_id = get_current_tenant_id()
        
        result = await self.db.execute(
            f"""
            DELETE FROM {self.table_name}
            WHERE id = $1 AND tenant_id = $2
            """,
            id, tenant_id
        )
        
        return result.rowcount > 0
    
    async def count(self, filters: dict = None) -> int:
        """
        Count records within current tenant.
        """
        tenant_id = get_current_tenant_id()
        
        query = f"SELECT COUNT(*) FROM {self.table_name} WHERE tenant_id = $1"
        params = [tenant_id]
        
        if filters:
            for i, (key, value) in enumerate(filters.items()):
                query += f" AND {key} = ${i+2}"
                params.append(value)
        
        result = await self.db.fetchrow(query, *params)
        return result["count"]


# Example usage
@dataclass
class User:
    id: str
    tenant_id: str
    email: str
    name: str
    role: str


class UserRepository(TenantAwareRepository[User]):
    def __init__(self, db):
        super().__init__(db, "users", User)
    
    async def find_by_email(self, email: str) -> Optional[User]:
        """Custom query - still tenant-scoped."""
        tenant_id = get_current_tenant_id()
        
        row = await self.db.fetchrow(
            """
            SELECT * FROM users
            WHERE email = $1 AND tenant_id = $2
            """,
            email, tenant_id
        )
        
        if not row:
            return None
        
        return User(**row)

# tenant_isolation/database_rls.py

"""
PostgreSQL Row-Level Security setup for multi-tenant isolation.

RLS provides database-enforced tenant isolation - even if
application code has a bug, the database won't return wrong data.
"""

RLS_SETUP_SQL = """
-- =============================================================================
-- ROW-LEVEL SECURITY SETUP
-- =============================================================================

-- Enable RLS on tenant tables
ALTER TABLE users ENABLE ROW LEVEL SECURITY;
ALTER TABLE orders ENABLE ROW LEVEL SECURITY;
ALTER TABLE documents ENABLE ROW LEVEL SECURITY;

-- Force RLS even for table owners (important!)
ALTER TABLE users FORCE ROW LEVEL SECURITY;
ALTER TABLE orders FORCE ROW LEVEL SECURITY;
ALTER TABLE documents FORCE ROW LEVEL SECURITY;


-- =============================================================================
-- TENANT CONTEXT FUNCTIONS
-- =============================================================================

-- Function to get current tenant from session variable
CREATE OR REPLACE FUNCTION current_tenant_id() 
RETURNS TEXT AS $$
BEGIN
    RETURN current_setting('app.current_tenant', true);
END;
$$ LANGUAGE plpgsql STABLE;


-- =============================================================================
-- ROW-LEVEL SECURITY POLICIES
-- =============================================================================

-- Users table policies
CREATE POLICY tenant_isolation_users ON users
    FOR ALL
    TO application_role
    USING (tenant_id = current_tenant_id())
    WITH CHECK (tenant_id = current_tenant_id());

-- Orders table policies  
CREATE POLICY tenant_isolation_orders ON orders
    FOR ALL
    TO application_role
    USING (tenant_id = current_tenant_id())
    WITH CHECK (tenant_id = current_tenant_id());

-- Documents table policies
CREATE POLICY tenant_isolation_documents ON documents
    FOR ALL
    TO application_role
    USING (tenant_id = current_tenant_id())
    WITH CHECK (tenant_id = current_tenant_id());


-- =============================================================================
-- ADMIN BYPASS POLICY (for support/admin operations)
-- =============================================================================

-- Admin role can see all data (use with extreme caution!)
CREATE POLICY admin_bypass_users ON users
    FOR ALL
    TO admin_role
    USING (true)
    WITH CHECK (true);

-- Note: Admin access should be heavily audited
"""


class RLSConnectionManager:
    """
    Manages database connections with RLS tenant context.
    
    Sets the tenant context on each connection before use.
    """
    
    def __init__(self, pool):
        self.pool = pool
    
    async def get_connection(self):
        """
        Get a connection with tenant context set.
        """
        conn = await self.pool.acquire()
        
        # Set tenant context for RLS. SET can't take bind parameters,
        # so use set_config() with a parameter (never an f-string: injection)
        tenant_id = get_current_tenant_id()
        await conn.execute(
            "SELECT set_config('app.current_tenant', $1, false)",
            tenant_id
        )
        
        return conn
    
    async def release_connection(self, conn):
        """
        Release connection and clear tenant context.
        """
        # Clear tenant context before returning to pool
        await conn.execute("RESET app.current_tenant")
        await self.pool.release(conn)


class TenantAwareConnectionPool:
    """
    Connection pool that automatically sets tenant context.
    
    Usage:
        async with await pool.connection() as conn:
            # tenant_id is automatically set
            result = await conn.fetch("SELECT * FROM users")
            # RLS filters to current tenant
    """
    
    def __init__(self, dsn: str):
        self.dsn = dsn
        self._pool = None
    
    async def initialize(self):
        """Create the connection pool."""
        import asyncpg
        self._pool = await asyncpg.create_pool(
            self.dsn,
            min_size=5,
            max_size=20,
            setup=self._setup_connection
        )
    
    async def _setup_connection(self, conn):
        """Called for each new connection."""
        # Connection starts with no tenant context
        pass
    
    async def connection(self):
        """Get connection context manager with tenant set."""
        return TenantConnection(self._pool)


class TenantConnection:
    """Async context manager for tenant-scoped connections."""
    
    def __init__(self, pool):
        self.pool = pool
        self.conn = None
    
    async def __aenter__(self):
        self.conn = await self.pool.acquire()
        
        # Set tenant context. is_local=false makes it session-level:
        # with true it would reset at the end of the implicit transaction,
        # i.e. right after this statement. __aexit__ clears it.
        try:
            tenant_id = get_current_tenant_id()
            await self.conn.execute(
                "SELECT set_config('app.current_tenant', $1, false)",
                tenant_id
            )
        except TenantContextError:
            # No tenant context - queries will fail RLS
            # This is intentional - forces proper context
            pass
        
        return self.conn
    
    async def __aexit__(self, exc_type, exc_val, exc_tb):
        if self.conn:
            # Clear tenant context
            await self.conn.execute(
                "SELECT set_config('app.current_tenant', '', false)"
            )
            await self.pool.release(self.conn)
        return False

# tenant_isolation/schema_manager.py

"""
Schema-per-tenant implementation.

Each tenant gets their own database schema with identical structure.
Provides stronger isolation than shared tables.
"""

import logging
import re
from typing import List

logger = logging.getLogger(__name__)


class SchemaPerTenantManager:
    """
    Manages schema-per-tenant database architecture.
    """
    
    def __init__(self, admin_pool, schema_prefix: str = "tenant_"):
        self.admin_pool = admin_pool
        self.schema_prefix = schema_prefix
    
    async def create_tenant_schema(self, tenant_id: str):
        """
        Create a new schema for a tenant.
        
        Called during tenant onboarding.
        """
        # Identifiers can't be bound as parameters: validate before interpolating
        if not re.fullmatch(r"[a-z0-9_]{1,50}", tenant_id):
            raise ValueError(f"Invalid tenant_id: {tenant_id!r}")
        schema_name = f"{self.schema_prefix}{tenant_id}"
        
        async with self.admin_pool.acquire() as conn:
            # Create schema
            await conn.execute(f"CREATE SCHEMA IF NOT EXISTS {schema_name}")
            
            # Create tables in schema
            await self._create_schema_tables(conn, schema_name)
            
            # Grant permissions
            await conn.execute(f"""
                GRANT USAGE ON SCHEMA {schema_name} TO app_user
            """)
            await conn.execute(f"""
                GRANT ALL ON ALL TABLES IN SCHEMA {schema_name} TO app_user
            """)
            
            logger.info(f"Created schema for tenant: {tenant_id}")
    
    async def _create_schema_tables(self, conn, schema_name: str):
        """Create standard tables in tenant schema."""
        
        await conn.execute(f"""
            CREATE TABLE {schema_name}.users (
                id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
                email VARCHAR(255) UNIQUE NOT NULL,
                name VARCHAR(255) NOT NULL,
                role VARCHAR(50) DEFAULT 'member',
                created_at TIMESTAMP DEFAULT NOW(),
                updated_at TIMESTAMP DEFAULT NOW()
            )
        """)
        
        await conn.execute(f"""
            CREATE TABLE {schema_name}.documents (
                id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
                title VARCHAR(500) NOT NULL,
                content TEXT,
                created_by UUID REFERENCES {schema_name}.users(id),
                created_at TIMESTAMP DEFAULT NOW(),
                updated_at TIMESTAMP DEFAULT NOW()
            )
        """)
        
        await conn.execute(f"""
            CREATE TABLE {schema_name}.settings (
                key VARCHAR(255) PRIMARY KEY,
                value JSONB,
                updated_at TIMESTAMP DEFAULT NOW()
            )
        """)
    
    async def delete_tenant_schema(self, tenant_id: str):
        """
        Delete a tenant's schema.
        
        Called during tenant offboarding. Requires confirmation!
        """
        schema_name = f"{self.schema_prefix}{tenant_id}"
        
        async with self.admin_pool.acquire() as conn:
            # CASCADE drops all objects in schema
            await conn.execute(f"DROP SCHEMA IF EXISTS {schema_name} CASCADE")
            
            logger.info(f"Deleted schema for tenant: {tenant_id}")
    
    async def get_all_tenant_schemas(self) -> List[str]:
        """Get list of all tenant schemas."""
        
        async with self.admin_pool.acquire() as conn:
            rows = await conn.fetch(f"""
                SELECT schema_name 
                FROM information_schema.schemata
                WHERE schema_name LIKE '{self.schema_prefix}%'
            """)
            
            return [row["schema_name"] for row in rows]
    
    async def run_migration(self, migration_sql: str):
        """
        Run migration across all tenant schemas.
        """
        schemas = await self.get_all_tenant_schemas()
        
        async with self.admin_pool.acquire() as conn:
            for schema in schemas:
                try:
                    # Set search path to tenant schema
                    await conn.execute(f"SET search_path TO {schema}")
                    
                    # Run migration
                    await conn.execute(migration_sql)
                    
                    logger.info(f"Migration complete for schema: {schema}")
                    
                except Exception as e:
                    logger.error(f"Migration failed for {schema}: {e}")
                    raise
        
        logger.info(f"Migration complete for {len(schemas)} schemas")


class SchemaRoutingConnectionPool:
    """
    Connection pool that routes to tenant-specific schema.
    """
    
    def __init__(self, pool, schema_prefix: str = "tenant_"):
        self.pool = pool
        self.schema_prefix = schema_prefix
    
    async def connection(self):
        """Get connection with schema routing."""
        return SchemaRoutedConnection(self.pool, self.schema_prefix)


class SchemaRoutedConnection:
    """Connection that automatically routes to tenant schema."""
    
    def __init__(self, pool, schema_prefix: str):
        self.pool = pool
        self.schema_prefix = schema_prefix
        self.conn = None
    
    async def __aenter__(self):
        self.conn = await self.pool.acquire()
        
        # Set search path to tenant schema
        tenant_id = get_current_tenant_id()
        schema_name = f"{self.schema_prefix}{tenant_id}"
        
        await self.conn.execute(f"SET search_path TO {schema_name}, public")
        
        return self.conn
    
    async def __aexit__(self, exc_type, exc_val, exc_tb):
        if self.conn:
            # Reset search path
            await self.conn.execute("SET search_path TO public")
            await self.pool.release(self.conn)
        return False

# tenant_isolation/database_per_tenant.py

"""
Database-per-tenant implementation.

Each tenant gets their own database instance.
Provides strongest isolation but highest operational complexity.
"""

from dataclasses import dataclass
from typing import Dict, Optional
import logging

logger = logging.getLogger(__name__)


@dataclass
class TenantDatabase:
    """Configuration for a tenant's database."""
    tenant_id: str
    host: str
    port: int
    database: str
    username: str
    password: str  # Should be encrypted/from secrets manager
    ssl_mode: str = "require"
    
    @property
    def connection_string(self) -> str:
        return (
            f"postgresql://{self.username}:{self.password}@"
            f"{self.host}:{self.port}/{self.database}"
            f"?sslmode={self.ssl_mode}"
        )


class DatabasePerTenantManager:
    """
    Manages database-per-tenant architecture.
    
    In practice, this often integrates with:
    - AWS RDS for database provisioning
    - Terraform for infrastructure
    - Secrets Manager for credentials
    """
    
    def __init__(self, db_cluster_endpoint: str, admin_credentials: dict):
        self.cluster_endpoint = db_cluster_endpoint
        self.admin_creds = admin_credentials
        self.tenant_pools: Dict[str, any] = {}
    
    async def provision_tenant_database(
        self,
        tenant_id: str,
        tier: str = "standard"
    ) -> TenantDatabase:
        """
        Provision a new database for a tenant.
        
        In production, this might:
        - Call AWS RDS API to create database
        - Use Terraform to provision infrastructure
        - Set up replication, backups, etc.
        """
        
        import asyncpg
        import secrets
        
        database_name = f"tenant_{tenant_id}"
        username = f"user_{tenant_id}"
        password = secrets.token_urlsafe(32)
        
        # Connect as admin to create database
        admin_conn = await asyncpg.connect(
            host=self.cluster_endpoint,
            user=self.admin_creds["username"],
            password=self.admin_creds["password"],
            database="postgres"
        )
        
        try:
            # Create user first: it will own the database. On PostgreSQL 15+
            # only the owner can create tables in the new database's public
            # schema (the admin role must be able to SET ROLE to the owner)
            await admin_conn.execute(f"""
                CREATE USER {username} WITH PASSWORD '{password}'
            """)
            
            # Create database
            await admin_conn.execute(f"""
                CREATE DATABASE {database_name}
                WITH OWNER {username} ENCODING 'UTF8'
            """)
            
            # Grant permissions
            await admin_conn.execute(f"""
                GRANT ALL PRIVILEGES ON DATABASE {database_name} TO {username}
            """)
            
            logger.info(f"Provisioned database for tenant: {tenant_id}")
            
        finally:
            await admin_conn.close()
        
        # Connect to new database to set up schema
        tenant_conn = await asyncpg.connect(
            host=self.cluster_endpoint,
            user=username,
            password=password,
            database=database_name
        )
        
        try:
            await self._create_tenant_schema(tenant_conn)
        finally:
            await tenant_conn.close()
        
        tenant_db = TenantDatabase(
            tenant_id=tenant_id,
            host=self.cluster_endpoint,
            port=5432,
            database=database_name,
            username=username,
            password=password
        )
        
        # Store credentials in secrets manager (not shown)
        await self._store_credentials(tenant_id, tenant_db)
        
        return tenant_db
    
    async def _create_tenant_schema(self, conn):
        """Create tables in tenant database."""
        
        await conn.execute("""
            CREATE TABLE users (
                id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
                email VARCHAR(255) UNIQUE NOT NULL,
                name VARCHAR(255) NOT NULL,
                role VARCHAR(50) DEFAULT 'member',
                created_at TIMESTAMP DEFAULT NOW()
            )
        """)
        
        await conn.execute("""
            CREATE TABLE documents (
                id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
                title VARCHAR(500) NOT NULL,
                content TEXT,
                created_by UUID REFERENCES users(id),
                created_at TIMESTAMP DEFAULT NOW()
            )
        """)
    
    async def _store_credentials(self, tenant_id: str, db: TenantDatabase):
        """Store credentials in secrets manager."""
        # In production: AWS Secrets Manager, HashiCorp Vault, etc.
        pass
    
    async def get_tenant_pool(self, tenant_id: str):
        """
        Get connection pool for a specific tenant.
        
        Pools are created lazily and cached.
        """
        import asyncpg
        
        if tenant_id in self.tenant_pools:
            return self.tenant_pools[tenant_id]
        
        # Load tenant database config
        tenant_db = await self._load_tenant_config(tenant_id)
        
        if not tenant_db:
            raise ValueError(f"No database configured for tenant: {tenant_id}")
        
        # Create pool
        pool = await asyncpg.create_pool(
            tenant_db.connection_string,
            min_size=2,
            max_size=10
        )
        
        self.tenant_pools[tenant_id] = pool
        return pool
    
    async def _load_tenant_config(self, tenant_id: str) -> Optional[TenantDatabase]:
        """Load tenant database configuration."""
        # In production: Load from secrets manager
        pass


class TenantDatabaseRouter:
    """
    Routes database connections to the correct tenant database.
    """
    
    def __init__(self, db_manager: DatabasePerTenantManager):
        self.db_manager = db_manager
    
    async def connection(self):
        """Get connection to current tenant's database."""
        tenant_id = get_current_tenant_id()
        pool = await self.db_manager.get_tenant_pool(tenant_id)
        return pool.acquire()

# tenant_isolation/service.py

"""
Complete example of a tenant-aware service.

This shows how all the pieces fit together:
- Middleware extracts tenant
- Service uses tenant-aware repository
- Caching is tenant-scoped
- Logging includes tenant context
"""

from fastapi import FastAPI, Depends, HTTPException
from typing import List, Optional
from dataclasses import dataclass
from datetime import datetime, timezone
import logging

logger = logging.getLogger(__name__)


@dataclass
class CreateDocumentRequest:
    title: str
    content: str


@dataclass
class Document:
    id: str
    tenant_id: str
    title: str
    content: str
    created_by: str
    created_at: datetime


class DocumentService:
    """
    Tenant-aware document service.
    
    All operations are automatically scoped to current tenant.
    """
    
    def __init__(self, document_repo, user_repo, cache, event_publisher,
                 search_client=None):
        self.documents = document_repo
        self.users = user_repo
        self.cache = cache
        self.events = event_publisher
        self.search_client = search_client
    
    @requires_tenant_async
    async def create_document(
        self,
        request: CreateDocumentRequest,
        created_by_user_id: str
    ) -> Document:
        """
        Create a new document in current tenant.
        """
        tenant = get_current_tenant()
        
        # Validate user exists in this tenant
        user = await self.users.find_by_id(created_by_user_id)
        if not user:
            raise ValueError("User not found")
        
        # Create document
        document = await self.documents.create({
            "title": request.title,
            "content": request.content,
            "created_by": created_by_user_id,
            "created_at": datetime.now(timezone.utc)
        })
        
        # Invalidate cache
        await self._invalidate_document_cache()
        
        # Publish event
        await self.events.publish(
            topic="documents",
            event={
                "type": "document.created",
                "tenant_id": tenant.tenant_id,
                "document_id": document.id,
                "created_by": created_by_user_id
            }
        )
        
        logger.info(
            "Document created",
            extra={
                "tenant_id": tenant.tenant_id,
                "document_id": document.id,
                "user_id": created_by_user_id
            }
        )
        
        return document
    
    @requires_tenant_async
    async def get_document(self, document_id: str) -> Optional[Document]:
        """
        Get document by ID within current tenant.
        """
        tenant = get_current_tenant()
        cache_key = f"tenant:{tenant.tenant_id}:doc:{document_id}"
        
        # Check cache
        cached = await self.cache.get(cache_key)
        if cached:
            return Document(**cached)
        
        # Query database
        document = await self.documents.find_by_id(document_id)
        
        if document:
            # Cache for 5 minutes
            await self.cache.set(cache_key, document.__dict__, ttl=300)
        
        return document
    
    @requires_tenant_async
    async def list_documents(
        self,
        limit: int = 50,
        offset: int = 0
    ) -> List[Document]:
        """
        List documents in current tenant.
        """
        return await self.documents.find_all(
            limit=limit,
            offset=offset
        )
    
    @requires_tenant_async
    async def search_documents(self, query: str) -> List[Document]:
        """
        Search documents within current tenant.
        
        Search index is tenant-scoped.
        """
        tenant = get_current_tenant()
        
        # Search in tenant-scoped index
        results = await self.search_client.search(
            index=f"documents_{tenant.tenant_id}",
            query=query
        )
        
        return [Document(**r) for r in results]
    
    @requires_tenant_async
    async def delete_document(self, document_id: str) -> bool:
        """
        Delete document within current tenant.
        """
        tenant = get_current_tenant()
        
        # Delete from database
        deleted = await self.documents.delete(document_id)
        
        if deleted:
            # Invalidate cache
            cache_key = f"tenant:{tenant.tenant_id}:doc:{document_id}"
            await self.cache.delete(cache_key)
            
            # Publish event
            await self.events.publish(
                topic="documents",
                event={
                    "type": "document.deleted",
                    "tenant_id": tenant.tenant_id,
                    "document_id": document_id
                }
            )
        
        return deleted
    
    async def _invalidate_document_cache(self):
        """Invalidate document list caches for current tenant."""
        tenant = get_current_tenant()
        pattern = f"tenant:{tenant.tenant_id}:docs:*"
        await self.cache.delete_pattern(pattern)


# FastAPI router with tenant middleware

app = FastAPI()

# Dependency to get document service
async def get_document_service() -> DocumentService:
    # In production, use dependency injection
    pass


@app.post("/api/documents")
async def create_document(
    request: CreateDocumentRequest,
    service: DocumentService = Depends(get_document_service)
):
    """
    Create a new document.
    
    Tenant context is set by middleware before this runs.
    """
    # Get current user from auth context
    user_id = get_current_user_id()
    
    document = await service.create_document(request, user_id)
    
    return {"id": document.id, "title": document.title}


@app.get("/api/documents/{document_id}")
async def get_document(
    document_id: str,
    service: DocumentService = Depends(get_document_service)
):
    """Get a document by ID."""
    document = await service.get_document(document_id)
    
    if not document:
        raise HTTPException(status_code=404, detail="Document not found")
    
    return document


@app.get("/api/documents")
async def list_documents(
    limit: int = 50,
    offset: int = 0,
    service: DocumentService = Depends(get_document_service)
):
    """List documents in current tenant."""
    documents = await service.list_documents(limit=limit, offset=offset)
    return {"documents": documents, "count": len(documents)}

# tenant_isolation/events.py

"""
Tenant-aware event processing.

Events in multi-tenant systems must carry tenant context
so consumers can process them correctly.
"""

from dataclasses import dataclass
from typing import Dict, Any, Callable
from datetime import datetime, timezone
import json


@dataclass
class TenantEvent:
    """
    Event with tenant context.
    
    All events must include tenant_id for proper routing and processing.
    """
    event_id: str
    event_type: str
    tenant_id: str
    timestamp: datetime
    payload: Dict[str, Any]
    
    def to_json(self) -> str:
        return json.dumps({
            "event_id": self.event_id,
            "event_type": self.event_type,
            "tenant_id": self.tenant_id,
            "timestamp": self.timestamp.isoformat(),
            "payload": self.payload
        })
    
    @classmethod
    def from_json(cls, data: str) -> 'TenantEvent':
        d = json.loads(data)
        return cls(
            event_id=d["event_id"],
            event_type=d["event_type"],
            tenant_id=d["tenant_id"],
            timestamp=datetime.fromisoformat(d["timestamp"]),
            payload=d["payload"]
        )


class TenantEventPublisher:
    """
    Publishes events with tenant context.
    """
    
    def __init__(self, kafka_producer):
        self.producer = kafka_producer
    
    async def publish(self, topic: str, event: Dict[str, Any]):
        """
        Publish event with current tenant context.
        """
        import uuid
        
        tenant = get_current_tenant()
        
        tenant_event = TenantEvent(
            event_id=str(uuid.uuid4()),
            event_type=event.get("type", "unknown"),
            tenant_id=tenant.tenant_id,
            timestamp=datetime.now(timezone.utc),
            payload=event
        )
        
        # Use tenant_id as partition key for ordering
        await self.producer.send(
            topic=topic,
            key=tenant.tenant_id,
            value=tenant_event.to_json()
        )


class TenantEventConsumer:
    """
    Consumes events and sets tenant context before processing.
    """
    
    def __init__(self, kafka_consumer, tenant_service):
        self.consumer = kafka_consumer
        self.tenant_service = tenant_service
        self.handlers: Dict[str, Callable] = {}
    
    def register_handler(self, event_type: str, handler: Callable):
        """Register handler for event type."""
        self.handlers[event_type] = handler
    
    async def start(self):
        """Start consuming events."""
        async for message in self.consumer:
            await self._process_message(message)
    
    async def _process_message(self, message):
        """Process a single message with tenant context."""
        
        event = TenantEvent.from_json(message.value)
        
        # Load tenant context
        tenant_ctx = await self.tenant_service.get_tenant_context(
            event.tenant_id
        )
        
        if not tenant_ctx:
            logger.error(f"Unknown tenant: {event.tenant_id}")
            return
        
        # Set tenant context and process
        with tenant_scope(tenant_ctx):
            handler = self.handlers.get(event.event_type)
            
            if handler:
                try:
                    await handler(event)
                except Exception as e:
                    logger.error(
                        f"Event processing failed: {e}",
                        extra={
                            "tenant_id": event.tenant_id,
                            "event_id": event.event_id,
                            "event_type": event.event_type
                        }
                    )
            else:
                logger.warning(f"No handler for event type: {event.event_type}")

# tenant_isolation/jobs.py

"""
Tenant-aware background job processing.

All background jobs must run within tenant context.
"""

from celery import Celery
from functools import wraps


celery_app = Celery('tasks')


def tenant_task(func):
    """
    Decorator for tenant-aware Celery tasks.
    
    Usage:
        @celery_app.task
        @tenant_task
        def process_document(document_id):
            # Runs in tenant context
            doc = document_repo.find_by_id(document_id)
    """
    @wraps(func)
    def wrapper(tenant_id: str, *args, **kwargs):
        # Load tenant context
        tenant_ctx = load_tenant_sync(tenant_id)
        
        if not tenant_ctx:
            raise ValueError(f"Unknown tenant: {tenant_id}")
        
        # Set context and run
        with tenant_scope(tenant_ctx):
            return func(*args, **kwargs)
    
    return wrapper


# Example tasks

@celery_app.task
@tenant_task
def send_document_notification(document_id: str, user_ids: list):
    """
    Send notifications about a document.
    
    Note: callers pass tenant_id as the first argument; @tenant_task consumes it
    """
    document = document_repo.find_by_id(document_id)
    
    if not document:
        return
    
    for user_id in user_ids:
        user = user_repo.find_by_id(user_id)
        if user:
            notification_service.send(
                user=user,
                template="document_shared",
                context={"document": document}
            )


@celery_app.task
@tenant_task
def generate_tenant_report():
    """
    Generate usage report for current tenant.
    """
    tenant = get_current_tenant()
    
    stats = {
        "tenant_id": tenant.tenant_id,
        "user_count": user_repo.count(),
        "document_count": document_repo.count(),
        "storage_used": storage_service.get_usage(),
    }
    
    report_service.generate(stats)


# Scheduling tenant jobs
class TenantJobScheduler:
    """
    Schedule jobs for all tenants or specific tenant.
    """
    
    def __init__(self, tenant_service):
        self.tenant_service = tenant_service
    
    async def schedule_for_all_tenants(self, task, *args, **kwargs):
        """
        Schedule a task to run for each tenant.
        
        Useful for nightly batch jobs.
        """
        tenants = await self.tenant_service.get_all_active_tenants()
        
        for tenant in tenants:
            # Include tenant_id as first argument
            task.delay(tenant.tenant_id, *args, **kwargs)
    
    async def schedule_for_tenant(self, tenant_id: str, task, *args, **kwargs):
        """Schedule a task for specific tenant."""
        task.delay(tenant_id, *args, **kwargs)

# tests/test_tenant_isolation.py

"""
Tests to verify tenant isolation.

These tests are CRITICAL - they verify that tenants can't see each other's data.
Run these in CI/CD and before every deployment.
"""

import pytest
from uuid import uuid4


class TestTenantIsolation:
    """
    Test suite for verifying tenant isolation.
    """
    
    @pytest.fixture
    async def tenant_a(self, tenant_service):
        """Create test tenant A."""
        return await tenant_service.create_tenant(
            name="Test Tenant A",
            plan="standard"
        )
    
    @pytest.fixture
    async def tenant_b(self, tenant_service):
        """Create test tenant B."""
        return await tenant_service.create_tenant(
            name="Test Tenant B",
            plan="standard"
        )
    
    async def test_tenant_cannot_see_other_tenant_data(
        self,
        tenant_a,
        tenant_b,
        document_service
    ):
        """
        Verify that Tenant A cannot see Tenant B's documents.
        """
        # Create document as Tenant A
        with tenant_scope(tenant_a):
            doc_a = await document_service.create_document(
                CreateDocumentRequest(
                    title="Tenant A Secret",
                    content="This is confidential"
                ),
                created_by_user_id="user_a"
            )
        
        # Try to access from Tenant B
        with tenant_scope(tenant_b):
            # Should not find Tenant A's document
            doc = await document_service.get_document(doc_a.id)
            assert doc is None, "Tenant B should not see Tenant A's document!"
            
            # List should not include Tenant A's documents
            docs = await document_service.list_documents()
            doc_ids = [d.id for d in docs]
            assert doc_a.id not in doc_ids, "Tenant A's document leaked to Tenant B!"
    
    async def test_tenant_cannot_modify_other_tenant_data(
        self,
        tenant_a,
        tenant_b,
        document_repo
    ):
        """
        Verify that Tenant A cannot modify Tenant B's documents.
        """
        # Create document as Tenant A
        with tenant_scope(tenant_a):
            doc_a = await document_repo.create({
                "title": "Original Title",
                "content": "Original Content"
            })
        
        # Try to update from Tenant B
        with tenant_scope(tenant_b):
            result = await document_repo.update(
                doc_a.id,
                {"title": "Hacked Title"}
            )
            assert result is None, "Tenant B should not modify Tenant A's document!"
        
        # Verify document unchanged
        with tenant_scope(tenant_a):
            doc = await document_repo.find_by_id(doc_a.id)
            assert doc.title == "Original Title"
    
    async def test_tenant_cannot_delete_other_tenant_data(
        self,
        tenant_a,
        tenant_b,
        document_repo
    ):
        """
        Verify that Tenant A cannot delete Tenant B's documents.
        """
        # Create document as Tenant A
        with tenant_scope(tenant_a):
            doc_a = await document_repo.create({
                "title": "Important Document",
                "content": "Do not delete"
            })
        
        # Try to delete from Tenant B
        with tenant_scope(tenant_b):
            deleted = await document_repo.delete(doc_a.id)
            assert not deleted, "Tenant B should not delete Tenant A's document!"
        
        # Verify document still exists
        with tenant_scope(tenant_a):
            doc = await document_repo.find_by_id(doc_a.id)
            assert doc is not None
    
    async def test_rls_filters_unscoped_queries(
        self,
        tenant_a,
        tenant_b,
        tenant_pool
    ):
        """
        Test the RLS policy directly: a raw query with NO tenant filter
        must still return only the current tenant's rows.
        
        (RLS alone does not stop SQL injection: injected SQL could call
        set_config() to switch tenant. Always parameterise queries.)
        """
        with tenant_scope(tenant_a):
            async with await tenant_pool.connection() as conn:
                rows = await conn.fetch("SELECT tenant_id FROM documents")
        
        for row in rows:
            assert row["tenant_id"] == tenant_a.tenant_id
    
    async def test_cache_isolation(
        self,
        tenant_a,
        tenant_b,
        document_service,
        cache
    ):
        """
        Verify that cached data is tenant-isolated.
        """
        # Create and cache document as Tenant A
        with tenant_scope(tenant_a):
            doc_a = await document_service.create_document(
                CreateDocumentRequest(title="Cached Doc", content="Secret"),
                created_by_user_id="user_a"
            )
            # Access to populate cache
            await document_service.get_document(doc_a.id)
        
        # Verify cache key is tenant-scoped
        cache_key_a = f"tenant:{tenant_a.tenant_id}:doc:{doc_a.id}"
        assert await cache.exists(cache_key_a)
        
        # Tenant B should not hit this cache
        with tenant_scope(tenant_b):
            cache_key_b = f"tenant:{tenant_b.tenant_id}:doc:{doc_a.id}"
            assert not await cache.exists(cache_key_b)


class TestTenantContextPropagation:
    """
    Test that tenant context flows correctly through the system.
    """
    
    async def test_context_available_in_service(self, tenant_a, document_service):
        """Verify tenant context is available in service layer."""
        with tenant_scope(tenant_a):
            # Service should be able to get current tenant
            tenant = get_current_tenant()
            assert tenant.tenant_id == tenant_a.tenant_id
    
    async def test_context_propagates_to_repository(self, tenant_a, document_repo):
        """Verify tenant context propagates to repository."""
        with tenant_scope(tenant_a):
            # Repository should automatically use current tenant
            doc = await document_repo.create({"title": "Test", "content": "Test"})
            assert doc.tenant_id == tenant_a.tenant_id
    
    async def test_no_context_raises_error(self, document_service):
        """Verify operations fail without tenant context."""
        # Outside tenant scope
        with pytest.raises(TenantContextError):
            await document_service.list_documents()
    
    async def test_context_cleared_after_scope(self, tenant_a):
        """Verify context is cleared after scope exits."""
        with tenant_scope(tenant_a):
            assert get_current_tenant().tenant_id == tenant_a.tenant_id
        
        # After scope, context should be cleared
        with pytest.raises(TenantContextError):
            get_current_tenant()

# tests/load_test_tenants.py

"""
Load testing for multi-tenant systems.

Verify that:
1. Multiple tenants can operate concurrently
2. One tenant's load doesn't affect others
3. System handles tenant context correctly under load
"""

import asyncio
import httpx
from locust import HttpUser, task, between


class TenantLoadTest(HttpUser):
    """
    Locust load test simulating multiple tenants.
    """
    wait_time = between(1, 3)
    
    def on_start(self):
        """Setup - assign this user to a random tenant."""
        import random
        self.tenant_id = f"tenant_{random.randint(1, 100)}"
        # The tenant comes from the token; the header is only a hint
        self.headers = {
            "Authorization": f"Bearer {test_token_for(self.tenant_id)}",
            "X-Tenant-ID": self.tenant_id,
        }
    
    @task(3)
    def list_documents(self):
        """List documents for this tenant."""
        self.client.get("/api/documents", headers=self.headers)
    
    @task(2)
    def create_document(self):
        """Create a document for this tenant."""
        self.client.post(
            "/api/documents",
            json={"title": "Load Test Doc", "content": "Test content"},
            headers=self.headers
        )
    
    @task(1)
    def get_document(self):
        """Get a specific document."""
        self.client.get(f"/api/documents/doc_1", headers=self.headers)


async def verify_isolation_under_load():
    """
    Verify tenant isolation while system is under load.
    
    Run this while load test is active.
    """
    tenants = ["tenant_1", "tenant_2", "tenant_3"]
    
    async def check_tenant(tenant_id):
        """Verify tenant only sees their own data."""
        async with httpx.AsyncClient() as client:
            response = await client.get(
                "http://localhost:8000/api/documents",
                headers={"Authorization": f"Bearer {test_token_for(tenant_id)}"}
            )
            
            docs = response.json()["documents"]
            for doc in docs:
                assert doc["tenant_id"] == tenant_id, \
                    f"Tenant {tenant_id} saw document from {doc['tenant_id']}!"
    
    # Check all tenants concurrently
    await asyncio.gather(*[check_tenant(t) for t in tenants])
    print("Isolation verified under load!")