Foundation Platform Specification¶

from collections.abc import AsyncGenerator
from sqlalchemy.ext.asyncio import (
    AsyncSession,
    async_sessionmaker,
    create_async_engine,
)
from sqlalchemy.pool import NullPool
from .config import settings

# NullPool: SQLAlchemy must not maintain its own connection pool on top of PgBouncer.
# PgBouncer is the pool manager. SQLAlchemy acquires a connection, uses it, returns it.
engine = create_async_engine(
    settings.database_url,
    poolclass=NullPool,
    echo=settings.db_echo,
)

# Alembic uses direct connection (bypasses PgBouncer) for DDL operations
# that require session-level settings. Configured separately in alembic/env.py.

AsyncSessionLocal = async_sessionmaker(
    engine,
    class_=AsyncSession,
    expire_on_commit=False,
    autoflush=False,
)

async def get_db() -> AsyncGenerator[AsyncSession, None]:
    """
    Dependency that yields a tenant-scoped database session.

    Reads tenant and user from contextvars (set by middleware — never from
    request.state, which is mutable and not available in background jobs).

    SET LOCAL ROLE: assumes the per-tenant Postgres role (r_{slug}), which has
    GRANT on t_{slug} schema only. A missed SET LOCAL becomes a permission error
    rather than a silent cross-tenant fallback. Defence-in-depth for ADR-007.

    SET LOCAL search_path: scopes unqualified table names to the tenant schema.

    Both are SET LOCAL (transaction-scoped). PgBouncer must be in session mode
    for these to survive across multiple transactions on the same connection.
    See ADR-021 and Layer 1 docker-compose notes.
    """
    from .context import current_tenant, current_user, current_request_id

    tenant = current_tenant.get()
    if tenant is None:
        raise RuntimeError("No tenant context — TenantMiddleware did not set it")

    validate_schema_name(tenant.schema_name)   # guards against injection — ADR-007 M1
    role_name = f"r_{tenant.slug}"
    schema_name = tenant.schema_name

    async with AsyncSessionLocal() as session:
        # ROLE first, then search_path — both in one round trip
        await session.execute(text(
            f'SET LOCAL ROLE "{role_name}"; '
            f'SET LOCAL search_path TO "{schema_name}", public'
        ))
        # Bind session-level audit context (read by audit trigger) — ADR-010
        user = current_user.get()
        request_id = current_request_id.get()
        if user is not None:
            await session.execute(
                text("SET LOCAL app.user_id = :uid; SET LOCAL app.identity_id = :iid"),
                {"uid": str(user.id), "iid": str(user.identity_id)},
            )
        if request_id is not None:
            await session.execute(
                text("SET LOCAL app.request_id = :rid"),
                {"rid": str(request_id)},
            )
        await session.execute(
            text("SET LOCAL app.actor_type = :at"),
            {"at": "user" if user else "system"},
        )
        yield session

import uuid
from datetime import datetime
from sqlalchemy import DateTime, text
from sqlalchemy.dialects.postgresql import JSONB, UUID
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column

class PublicBase(DeclarativeBase):
    """Base for public schema tables: tenants, tenant_labels, api_keys, field_definitions."""
    pass

class TenantBase(DeclarativeBase):
    """Base for all tenant-schema entity tables."""
    pass

class TimestampMixin:
    """created_at and updated_at. updated_at maintained by DB trigger, not SQLAlchemy."""
    created_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True),
        server_default=text("NOW()"),
        nullable=False,
    )
    updated_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True),
        server_default=text("NOW()"),
        nullable=False,
    )

class EntityMixin(TimestampMixin):
    """Universal columns for every tenant entity table."""
    id: Mapped[uuid.UUID] = mapped_column(
        UUID(as_uuid=True),
        primary_key=True,
        default=uuid.uuid4,
    )
    created_by: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), nullable=False)
    updated_by: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), nullable=False)
    deleted_at: Mapped[datetime | None] = mapped_column(
        DateTime(timezone=True), nullable=True, default=None
    )
    custom_fields: Mapped[dict] = mapped_column(JSONB, nullable=False, default=dict)

from sqlalchemy import text
from alembic import context

# Alembic uses DIRECT database URL (not PgBouncer) for DDL
DATABASE_URL_DIRECT = os.environ["DATABASE_URL_DIRECT"]

async def run_migrations_online() -> None:
    engine = create_async_engine(DATABASE_URL_DIRECT)

    # Public schema only — tenant schemas are handled by migrate_tenants.py
    async with engine.begin() as conn:
        await conn.execute(text("SET search_path TO public"))
        await conn.run_sync(lambda sync_conn: context.run_migrations())

    await engine.dispose()

from dataclasses import dataclass, field
from typing import Literal
import asyncio, argparse, sys
from sqlalchemy.ext.asyncio import create_async_engine
from sqlalchemy import text

@dataclass
class MigrationReport:
    succeeded: list = field(default_factory=list)
    failed: list = field(default_factory=list)
    skipped: list = field(default_factory=list)

async def migrate_all_tenants(
    target_revision: str | None = None,
    only: list[str] | None = None,
    from_state: Literal["all", "failed"] = "all",
    dry_run: bool = False,
) -> MigrationReport:
    """
    Migrate every active tenant schema to target_revision (default: head).
    Continues past failures, records outcomes in public.tenant_migrations.
    Exits 0 only if every tenant reaches target_revision.
    """
    target = target_revision or get_head_revision()
    tenants = await fetch_active_tenants()
    if only:
        tenants = [t for t in tenants if t.slug in only]
    if from_state == "failed":
        tenants = [t for t in tenants if await last_attempt_failed(t.id, target)]

    report = MigrationReport()

    for tenant in tenants:
        validate_schema_name(tenant.schema_name)
        current = await get_current_revision(tenant.schema_name)
        if current == target:
            report.skipped.append(tenant)
            continue
        if dry_run:
            print(f"[dry-run] Would migrate {tenant.slug}: {current} → {target}")
            continue

        attempt_id = await record_attempt(tenant.id, target)
        try:
            await run_alembic_for_schema(tenant.schema_name, target, timeout_seconds=60)
            await record_success(attempt_id)
            report.succeeded.append(tenant)
        except Exception as e:
            await record_failure(attempt_id, str(e))
            report.failed.append((tenant, e))
            # Continue to next tenant — do NOT abort the whole run

    return report

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--only", help="Comma-separated tenant slugs")
    parser.add_argument("--target", help="Target revision (default: head)")
    parser.add_argument("--from-state", choices=["all", "failed"], default="all")
    parser.add_argument("--dry-run", action="store_true")
    args = parser.parse_args()

    report = asyncio.run(migrate_all_tenants(
        target_revision=args.target,
        only=args.only.split(",") if args.only else None,
        from_state=args.from_state,
        dry_run=args.dry_run,
    ))

    print(f"Succeeded: {len(report.succeeded)}, Failed: {len(report.failed)}, Skipped: {len(report.skipped)}")

    if report.failed:
        for tenant, err in report.failed:
            print(f"  FAILED: {tenant.slug} — {err}", file=sys.stderr)
        sys.exit(1)   # 1 = one or more tenants not at target
    sys.exit(0)       # 0 = all active tenants at target
    # Exit 2 is reserved for runner-itself failures (config error, DB unreachable)

import boto3
from .tenant import validate_schema_name

async def provision_tenant(
    db: AsyncSession,
    slug: str,
    name: str,
    plan: str = "starter",
    region: str = "eu-west-2",
) -> Tenant:
    schema_name = f"t_{slug}"
    role_name = f"r_{slug}"
    validate_schema_name(schema_name)  # ADR-007 M1

    # Generate per-tenant DEK (ADR-020)
    kms = boto3.client("kms", region_name=region)
    dek_response = kms.generate_data_key(
        KeyId=settings.kms_tenant_key_arn,
        KeySpec="AES_256",
    )
    encrypted_dek = dek_response["CiphertextBlob"]
    dek_key_id = settings.kms_tenant_key_arn

    async with db.begin():
        tenant = Tenant(
            slug=slug, name=name, schema_name=schema_name,
            plan=plan, region=region,
            encrypted_dek=encrypted_dek, dek_key_id=dek_key_id,
        )
        db.add(tenant)
        await db.flush()

        # Create tenant schema
        await db.execute(text(f'CREATE SCHEMA "{schema_name}"'))

        # Create per-tenant Postgres role (ADR-007 M2)
        # NOLOGIN — assumed via SET LOCAL ROLE by the application, never used directly
        await db.execute(text(f'CREATE ROLE "{role_name}" NOLOGIN'))
        await db.execute(text(
            f'GRANT USAGE ON SCHEMA "{schema_name}" TO "{role_name}"'
        ))
        await db.execute(text(
            f'ALTER DEFAULT PRIVILEGES IN SCHEMA "{schema_name}" '
            f'GRANT SELECT, INSERT, UPDATE, DELETE ON TABLES TO "{role_name}"'
        ))
        # Allow the app's main DB user to assume this role via SET LOCAL ROLE
        await db.execute(text(
            f'GRANT "{role_name}" TO "{settings.app_db_user}"'
        ))

        # Per-tenant resource limits (ADR-008)
        # These apply when connected as the tenant role
        timeouts = PLAN_TIMEOUTS[plan]  # e.g. {"statement": "10s", "idle": "30s", "lock": "5s"}
        await db.execute(text(
            f'ALTER ROLE "{role_name}" '
            f'SET statement_timeout = \'{timeouts["statement"]}\' '
            f'SET idle_in_transaction_session_timeout = \'{timeouts["idle"]}\' '
            f'SET lock_timeout = \'{timeouts["lock"]}\''
        ))

        # Run tenant migrations via standalone runner (ADR-001)
        await _run_tenant_migrations(schema_name, role_name)
        # Seed defaults into tenant schema
        await _seed_default_labels(db, tenant.id)
        await _seed_field_definitions(db, tenant.id)

    return tenant

PLAN_TIMEOUTS: dict[str, dict[str, str]] = {
    "trial":        {"statement": "10s",  "idle": "30s", "lock": "5s"},
    "starter":      {"statement": "10s",  "idle": "30s", "lock": "5s"},
    "professional": {"statement": "30s",  "idle": "60s", "lock": "10s"},
    "enterprise":   {"statement": "60s",  "idle": "120s","lock": "15s"},
}

from contextvars import ContextVar
from uuid import UUID

# Type annotations use forward references to avoid circular imports
current_tenant: ContextVar["Tenant | None"] = ContextVar("current_tenant", default=None)
current_user: ContextVar["CurrentUser | None"] = ContextVar("current_user", default=None)
current_request_id: ContextVar["UUID | None"] = ContextVar("current_request_id", default=None)

import re

_SCHEMA_NAME_RE = re.compile(r"^t_[a-z0-9_]{1,40}$")

def validate_schema_name(name: str) -> str:
    """
    Validates a tenant schema name against the allowed pattern.
    Called at every boundary where a schema name is used in dynamic SQL.
    Guards against injection if a slug ever reaches provisioning with bad characters.
    """
    if not _SCHEMA_NAME_RE.match(name):
        raise ValueError(f"Invalid schema name: {name!r}")
    return name

from starlette.middleware.base import BaseHTTPMiddleware
from .context import current_tenant, current_request_id
from .tenant import validate_schema_name
import uuid

class RequestIDMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next):
        request_id = uuid.uuid4()
        token = current_request_id.set(request_id)
        structlog.contextvars.bind_contextvars(
            request_id=str(request_id),
            path=request.url.path,
            method=request.method,
        )
        try:
            response = await call_next(request)
            response.headers["X-Request-ID"] = str(request_id)
            return response
        finally:
            current_request_id.reset(token)
            structlog.contextvars.clear_contextvars()

class TenantMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next):
        if request.url.path == "/health":
            return await call_next(request)

        host = request.headers.get("host", "").split(":")[0]
        slug = _extract_slug(host)
        if slug is None:
            return JSONResponse(
                {"error": "not_found", "message": "Unknown host"}, status_code=404
            )

        tenant = await _load_tenant(slug, request.app.state.redis)
        if tenant is None:
            return JSONResponse(
                {"error": "not_found", "message": "Tenant not found"}, status_code=404
            )

        # Validate schema name at the boundary (ADR-007 M1)
        try:
            validate_schema_name(tenant.schema_name)
        except ValueError:
            return JSONResponse(
                {"error": "internal_error", "message": "Invalid tenant configuration"},
                status_code=500,
            )

        token = current_tenant.set(tenant)
        structlog.contextvars.bind_contextvars(tenant_slug=tenant.slug)
        try:
            return await call_next(request)
        finally:
            current_tenant.reset(token)

def _extract_slug(host: str) -> str | None:
    """
    'acme.construo.build' → 'acme'
    'acme.local.test'     → 'acme'   (local dev)
    'localhost'           → None
    'construo.build'      → None     (no subdomain)
    """
    parts = host.split(".")
    if len(parts) < 2:
        return None
    slug = parts[0]
    return slug if slug else None

async def _load_tenant(slug: str, redis: Redis) -> Tenant | None:
    cached = await redis.get(f"tenant:{slug}")
    if cached:
        return Tenant.model_validate_json(cached)

    # Use lookup_tenant() SECURITY DEFINER function — not direct table read (ADR-009)
    async with AsyncSessionLocal() as db:
        result = await db.execute(
            text("SELECT * FROM public.lookup_tenant(:slug)"), {"slug": slug}
        )
        row = result.mappings().one_or_none()

    if row:
        tenant = Tenant(**row)
        await redis.setex(
            f"tenant:{tenant.id}:{slug}", 300, tenant.model_dump_json()
        )
        return tenant
    return None

class TenantAssertionMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next):
        response = await call_next(request)
        if request.url.path != "/health":
            # If we set a tenant at the start, it should still be set at the end.
            # If it's None here, middleware failed to reset — context leaked.
            tenant = current_tenant.get()
            if tenant is None:
                logger.error("tenant_context_missing_after_request",
                             path=request.url.path)
        return response

async def run_with_tenant(tenant_id: UUID, fn: "Callable") -> None:
    """Sets tenant context for a background job, then calls fn."""
    tenant = await load_tenant_by_id(tenant_id)
    token_t = current_tenant.set(tenant)
    token_r = current_request_id.set(uuid.uuid4())
    try:
        await fn()
    finally:
        current_tenant.reset(token_t)
        current_request_id.reset(token_r)

async def run_as_system(tenant_id: UUID, reason: str, fn: "Callable") -> None:
    """
    Like run_with_tenant but sets actor_type='system' for audit.
    Use for background jobs, scanners, and auto-actions (ADR-010).
    """
    from .context import current_user
    # current_user left as None — get_db will set actor_type='system'
    await run_with_tenant(tenant_id, fn)

async def get_platform_admin_db(
    user: "CurrentUser",
    action: str,
    target_tenant_id: UUID | None,
    details: dict,
) -> "AsyncGenerator[AsyncSession, None]":
    """
    Yields a session connected as the platform superuser (bypasses per-tenant role).
    Requires caller to have platform_admin role. Writes audit row before yielding.
    """
    if "platform_admin" not in user.roles:
        raise ForbiddenError("platform_admin role required")

    async with PlatformAdminSessionLocal() as session:
        # Pre-write audit before the operation (ADR-007 M6)
        await session.execute(
            text("""
                INSERT INTO public.platform_admin_actions
                    (actor_id, action, target_tenant_id, details)
                VALUES (:aid, :action, :tid, :details::jsonb)
            """),
            {"aid": str(user.id), "action": action,
             "tid": str(target_tenant_id) if target_tenant_id else None,
             "details": json.dumps(details)},
        )
        yield session

from redis.asyncio import Redis, from_url

async def init_redis() -> Redis:
    return await from_url(settings.redis_url, encoding="utf-8", decode_responses=False)

def get_redis(request: Request) -> Redis:
    return request.app.state.redis

async def get_label(key: str, redis: Redis, db: AsyncSession) -> str:
    """Returns tenant label override or falls back to DEFAULT_LABELS."""
    tenant = current_tenant.get()
    cache_key = f"labels:{tenant.id}:{key}"   # UUID key, not slug
    cached = await redis.get(cache_key)
    if cached:
        return cached.decode()

    # Labels table is in the tenant schema — get_db already set search_path
    result = await db.execute(
        select(Label.label_value).where(Label.label_key == key)
    )
    value = result.scalar_one_or_none()
    if value:
        await redis.setex(cache_key, 300, value)
        return value

    return DEFAULT_LABELS.get(key, key)

import boto3
from construo_shared import (
    get_tenant_for_client_id,
    get_membership,
    get_user_roles,
)
# construo_shared is a Lambda layer containing DB query utilities.
# Lambda connects to RDS via VPC ENI using IAM role — no hardcoded credentials.
# Results are cached in Lambda memory with a short TTL to reduce DB load (ADR-013).

_tenant_cache: dict = {}

def handler(event: dict, context) -> dict:
    """
    Cognito Pre-Token-Generation V2 trigger.
    Adds custom claims to the JWT before it is issued.
    """
    client_id = event["callerContext"]["clientId"]
    user_sub = event["userName"]

    # Resolve tenant from client ID (cached) — ADR-004, ADR-013
    tenant = _tenant_cache.get(client_id) or get_tenant_for_client_id(client_id)
    _tenant_cache[client_id] = tenant

    # Verify active membership
    membership = get_membership(user_sub, tenant.id)
    if not membership or membership.status != "active":
        raise Exception("No active membership for this tenant")

    # platform_roles as a native Cognito list claim (not JSON string) — ADR-003 P3
    roles = get_user_roles(user_sub, tenant.id)

    event["response"]["claimsOverrideDetails"] = {
        "claimsToAddOrOverride": {
            "custom:tenant_id":    str(tenant.id),
            "custom:tenant_slug":  tenant.slug,
            "custom:identity_id":  str(membership.identity_id),
        },
        # claimsToAddOrOverrideList is Cognito's native list-claim shape (2021+)
        "claimsToAddOrOverrideList": {
            "custom:platform_roles": roles,
        },
    }
    return event

import httpx
from jose import jwt, JWTError

class CognitoJWTValidator:
    def __init__(self, user_pool_id: str, region: str, client_id: str):
        self.jwks_url = (
            f"https://cognito-idp.{region}.amazonaws.com"
            f"/{user_pool_id}/.well-known/jwks.json"
        )
        self.client_id = client_id
        self._jwks: dict | None = None

    async def fetch_jwks(self) -> None:
        async with httpx.AsyncClient() as client:
            r = await client.get(self.jwks_url)
            r.raise_for_status()
            self._jwks = r.json()

    async def validate(self, token: str) -> dict:
        if self._jwks is None:
            await self.fetch_jwks()
        try:
            return jwt.decode(
                token, self._jwks, algorithms=["RS256"], audience=self.client_id
            )
        except JWTError as e:
            raise UnauthorisedError(str(e))

from pydantic import BaseModel
from .context import current_user

class CurrentUser(BaseModel):
    id: UUID           # user_id within tenant schema
    identity_id: UUID  # global identity (ADR-004)
    tenant_id: UUID
    tenant_slug: str
    roles: list[str]
    # project_access is NOT here — fetched live via core/access.py (ADR-003)

async def get_current_user(
    request: Request,
    authorization: str = Header(..., alias="Authorization"),
) -> CurrentUser:
    tenant = current_tenant.get()
    if tenant is None:
        raise UnauthorisedError("No tenant context")

    if authorization.startswith("Bearer "):
        token = authorization[7:]
        claims = await request.app.state.jwt_validator.validate(token)
        token_tenant_id = claims.get("custom:tenant_id")
        if str(tenant.id) != token_tenant_id:
            raise UnauthorisedError("Token/tenant mismatch")
        user = CurrentUser(
            id=UUID(claims["sub"]),
            identity_id=UUID(claims["custom:identity_id"]),
            tenant_id=UUID(token_tenant_id),
            tenant_slug=claims["custom:tenant_slug"],
            # platform_roles is now a native Cognito list claim (ADR-003 P3)
            roles=claims.get("custom:platform_roles", []),
        )
        # Bind to contextvar so get_db can set app.user_id in session (ADR-007, ADR-010)
        token_cv = current_user.set(user)
        structlog.contextvars.bind_contextvars(user_id=str(user.id))
        request.state._current_user_token = token_cv  # reset in middleware
        return user

    if authorization.startswith("ApiKey "):
        key = authorization[7:]
        return await _validate_api_key(key, request)

    raise UnauthorisedError("Bearer token or ApiKey required")

def require_permission(permission: str):
    async def _check(
        user: CurrentUser = Depends(get_current_user)
    ) -> CurrentUser:
        if not _has_permission(user.roles, permission):
            raise ForbiddenError(f"Requires {permission}")
        return user
    return _check

from typing import Literal
from uuid import UUID
from redis.asyncio import Redis
from sqlalchemy.ext.asyncio import AsyncSession

async def get_project_access(
    user_id: UUID,
    tenant_id: UUID,
    redis: Redis,
    db: AsyncSession,
) -> list[UUID] | Literal["all"]:
    """
    Returns the project UUIDs the user can access, or sentinel 'all'
    if the user is a tenant admin (implicit access to all projects).

    Redis-cached with 60-second TTL. Explicitly invalidated on write to
    user_project_access. ADR-003.
    """
    cache_key = f"upa:{user_id}:{tenant_id}"
    cached = await redis.get(cache_key)
    if cached is not None:
        return _decode_access(cached)

    # Tenant admins implicitly access all projects
    user = await fetch_user(user_id, tenant_id, db)
    if "tenant_admin" in user.roles or "platform_admin" in user.roles:
        await redis.setex(cache_key, 60, "ALL")
        return "all"

    result = await db.execute(
        select(UserProjectAccess.project_id)
        .where(UserProjectAccess.user_id == user_id)
        .where(UserProjectAccess.tenant_id == tenant_id)
        .where(UserProjectAccess.revoked_at.is_(None))
    )
    projects = [row.project_id for row in result.fetchall()]
    await redis.setex(cache_key, 60, _encode_access(projects))
    return projects

async def grant_project_access(
    user_id: UUID, tenant_id: UUID, project_id: UUID,
    role: str, granted_by: UUID, redis: Redis, db: AsyncSession,
) -> None:
    """Insert into user_project_access and invalidate cache. Never write directly."""
    await db.execute(
        insert(UserProjectAccess).values(
            user_id=user_id, tenant_id=tenant_id, project_id=project_id,
            role=role, granted_by=granted_by,
        )
    )
    await redis.delete(f"upa:{user_id}:{tenant_id}")

async def revoke_project_access(
    user_id: UUID, tenant_id: UUID, project_id: UUID,
    redis: Redis, db: AsyncSession,
) -> None:
    """Mark revoked_at and invalidate cache. Never write directly."""
    await db.execute(
        update(UserProjectAccess)
        .where(UserProjectAccess.user_id == user_id)
        .where(UserProjectAccess.tenant_id == tenant_id)
        .where(UserProjectAccess.project_id == project_id)
        .values(revoked_at=func.now())
    )
    await redis.delete(f"upa:{user_id}:{tenant_id}")

def require_project_access(project_id_param: str = "projectId"):
    """
    FastAPI dependency for project-scoped routes.
    Use in place of (or in addition to) require_permission() on any route
    whose URL contains a project or site ID.
    """
    async def _check(
        request: Request,
        user: CurrentUser = Depends(get_current_user),
        redis: Redis = Depends(get_redis),
        db: AsyncSession = Depends(get_db),
    ) -> CurrentUser:
        project_id = UUID(request.path_params[project_id_param])
        access = await get_project_access(user.id, user.tenant_id, redis, db)
        if access != "all" and project_id not in access:
            raise ForbiddenError("No access to this project")
        return user
    return _check