Clean Code Python: Multi-Tenant Foundation — Context, Isolation, and the Data Boundary

from dataclasses import dataclass
from uuid import UUID


@dataclass(frozen=True, slots=True)
class Tenant:
    """Immutable tenant context for the current request."""
    id: UUID
    slug: str
    plan: str  # "starter" | "growth" | "enterprise"

from contextvars import ContextVar
from uuid import UUID

from src.core.tenant import Tenant

_current_tenant: ContextVar[Tenant] = ContextVar("current_tenant")


def get_current_tenant() -> Tenant:
    """Return the tenant for the current request.

    Raises LookupError if called outside a tenant context.
    This is intentional — code that accesses tenant data
    without a tenant context is a bug.
    """
    try:
        return _current_tenant.get()
    except LookupError:
        raise RuntimeError(
            "No tenant context. This code path requires a tenant-scoped request."
        )


def set_current_tenant(tenant: Tenant) -> None:
    """Set the tenant for the current async context."""
    _current_tenant.set(tenant)


def get_current_tenant_id() -> UUID:
    """Shortcut — most callers only need the ID."""
    return get_current_tenant().id

from uuid import UUID

from fastapi import Request, HTTPException
from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
from starlette.responses import Response

from src.core.context import set_current_tenant
from src.core.tenant import Tenant


class TenantMiddleware(BaseHTTPMiddleware):
    """Extract tenant from JWT claims and propagate via contextvars."""

    TENANT_FREE_PATHS = {"/health", "/docs", "/openapi.json"}

    async def dispatch(
        self, request: Request, call_next: RequestResponseEndpoint
    ) -> Response:
        if request.url.path in self.TENANT_FREE_PATHS:
            return await call_next(request)

        tenant = self._extract_tenant(request)
        set_current_tenant(tenant)
        return await call_next(request)

    def _extract_tenant(self, request: Request) -> Tenant:
        # In production, decode and verify the JWT.
        # Simplified here to show the extraction pattern.
        claims = request.state.jwt_claims  # set by auth middleware
        tenant_id = claims.get("tenant_id")
        tenant_slug = claims.get("tenant_slug")
        tenant_plan = claims.get("tenant_plan", "starter")

        if not tenant_id or not tenant_slug:
            raise HTTPException(
                status_code=403,
                detail="Missing tenant context in access token",
            )

        return Tenant(
            id=UUID(tenant_id),
            slug=tenant_slug,
            plan=tenant_plan,
        )

from uuid import UUID

from sqlalchemy import ForeignKey, Index
from sqlalchemy.orm import Mapped, mapped_column, declared_attr


class TenantMixin:
    """Mixin that adds tenant_id to any SQLAlchemy model.

    Every query on a tenant-scoped model MUST filter by tenant_id.
    The session event listener enforces this automatically.
    """

    tenant_id: Mapped[UUID] = mapped_column(
        nullable=False,
        index=True,
    )

    @declared_attr
    @classmethod
    def __table_args__(cls):
        # Composite index: tenant_id + id for fast tenant-scoped lookups
        return (
            Index(f"ix_{cls.__tablename__}_tenant_id", "tenant_id"),
            {"extend_existing": True},
        )

from sqlalchemy.orm import Mapped, mapped_column

from src.db.base import Base
from src.db.mixins import TenantMixin


class Book(TenantMixin, Base):
    __tablename__ = "books"

    id: Mapped[int] = mapped_column(primary_key=True)
    title: Mapped[str] = mapped_column(nullable=False)
    isbn: Mapped[str] = mapped_column(unique=False)  # unique per tenant, not global
    price: Mapped[float] = mapped_column(nullable=False)

from sqlalchemy import event, inspect
from sqlalchemy.orm import Session, Query

from src.core.context import get_current_tenant_id
from src.db.mixins import TenantMixin


def _is_tenant_model(mapper) -> bool:
    """Check if a model uses TenantMixin."""
    return issubclass(mapper.class_, TenantMixin)


@event.listens_for(Session, "do_orm_execute")
def _apply_tenant_filter(execute_state):
    """Intercept every ORM query and inject tenant_id filtering.

    This fires on SELECT, UPDATE, and DELETE operations.
    INSERT is handled by _set_tenant_on_insert below.
    """
    if execute_state.is_orm_statement:
        # Only filter if the query targets a tenant-scoped model
        stmt = execute_state.statement
        for column_desc in stmt.column_descriptions:
            entity = column_desc.get("entity")
            if entity is not None and issubclass(entity, TenantMixin):
                tenant_id = get_current_tenant_id()
                execute_state.statement = stmt.filter(
                    entity.tenant_id == tenant_id
                )


@event.listens_for(Session, "before_flush")
def _set_tenant_on_insert(session, flush_context, instances):
    """Auto-set tenant_id on new objects before they are flushed to the database.

    This means services never need to manually set tenant_id — the system
    handles it. A developer literally cannot forget.
    """
    tenant_id = get_current_tenant_id()
    for obj in session.new:
        if isinstance(obj, TenantMixin):
            if obj.tenant_id is None:
                obj.tenant_id = tenant_id
            elif obj.tenant_id != tenant_id:
                raise ValueError(
                    f"Attempted to insert {type(obj).__name__} with tenant_id "
                    f"{obj.tenant_id} in context of tenant {tenant_id}. "
                    "This is a cross-tenant write violation."
                )

from src.core.context import system_context


async def aggregate_monthly_billing(operator_email: str):
    async with system_context(
        reason="Monthly billing aggregation",
        operator=operator_email,
    ):
        # This query runs WITHOUT tenant filtering — returns all tenants
        result = await session.execute(
            select(
                Order.tenant_id,
                func.sum(Order.total).label("revenue"),
            ).group_by(Order.tenant_id)
        )
        return result.all()

import pytest
from uuid import uuid4

from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker

from src.core.context import set_current_tenant
from src.core.tenant import Tenant
from src.db.base import Base
from src.db.tenant_filter import _apply_tenant_filter, _set_tenant_on_insert
from src.models.book import Book


TENANT_A = Tenant(id=uuid4(), slug="acme-books", plan="growth")
TENANT_B = Tenant(id=uuid4(), slug="beta-reads", plan="starter")


@pytest.fixture
async def session():
    engine = create_async_engine("sqlite+aiosqlite:///:memory:")
    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)
    factory = async_sessionmaker(engine)
    async with factory() as s:
        yield s
    await engine.dispose()


async def test_tenant_a_cannot_see_tenant_b_books(session):
    """The foundational isolation test. If this fails, nothing else matters."""

    # Insert books as Tenant A
    set_current_tenant(TENANT_A)
    session.add(Book(title="Tenant A Book", isbn="9780000000001", price=10.0))
    await session.flush()

    # Insert books as Tenant B
    set_current_tenant(TENANT_B)
    session.add(Book(title="Tenant B Book", isbn="9780000000002", price=20.0))
    await session.flush()

    # Query as Tenant A — must see ONLY Tenant A's book
    set_current_tenant(TENANT_A)
    from sqlalchemy import select
    result = await session.execute(select(Book))
    books = result.scalars().all()

    assert len(books) == 1
    assert books[0].title == "Tenant A Book"
    assert books[0].tenant_id == TENANT_A.id


async def test_tenant_cannot_insert_with_wrong_tenant_id(session):
    """Prevent a service from inserting data into another tenant's space."""
    set_current_tenant(TENANT_A)

    book = Book(title="Sneaky Book", isbn="9780000000003", price=15.0)
    book.tenant_id = TENANT_B.id  # Explicitly set wrong tenant

    session.add(book)
    with pytest.raises(ValueError, match="cross-tenant write violation"):
        await session.flush()


async def test_system_context_sees_all_tenants(session):
    """System context bypasses tenant filtering for admin operations."""
    from src.core.context import system_context

    # Insert as both tenants
    set_current_tenant(TENANT_A)
    session.add(Book(title="A Book", isbn="9780000000004", price=10.0))
    await session.flush()

    set_current_tenant(TENANT_B)
    session.add(Book(title="B Book", isbn="9780000000005", price=20.0))
    await session.flush()

    # System context sees everything
    async with system_context(reason="test", operator="test@shelfwise.io"):
        from sqlalchemy import select
        result = await session.execute(select(Book))
        books = result.scalars().all()
        assert len(books) == 2