saas_backend/app/middleware/idempotency_middleware.py

"""Honour an `Idempotency-Key` on requests that change something.

Middleware rather than a dependency, for the same reason the tenant scope is:
the work spans the endpoint. A dependency can claim the key before the handler
runs but has no way to see the response afterwards, and the response is the half
that matters — remembering it is the entire point.

Only `POST`, `PUT` and `PATCH`. `GET` and `DELETE` are already idempotent by
definition; adding a key to them would be a cache, which is a different feature
with different rules about staleness.

The key is scoped to the caller — workspace and user — read from the same token
the scope middleware reads, or from an API key. An unauthenticated request with a
key is passed straight through: it will be refused a moment later anyway, and
claiming a key on the caller's behalf before knowing who they are would let
anybody occupy anybody's key.
"""

from __future__ import annotations

import logging
import uuid
from typing import Optional

from starlette.middleware.base import BaseHTTPMiddleware
from starlette.requests import Request
from starlette.responses import JSONResponse, Response

from app.services.system import idempotency_service

logger = logging.getLogger(__name__)

HEADER = "Idempotency-Key"
METHODS = {"POST", "PUT", "PATCH"}


def _caller(request: Request) -> Optional[tuple[Optional[uuid.UUID], Optional[uuid.UUID]]]:
    """(tenant_id, user_id) for the caller, or None if they are not identified.

    Read from the token rather than from the database, exactly as the scope
    middleware does, so that this costs nothing on the overwhelming majority of
    requests that carry no key at all.
    """
    from app.middleware.tenant_scope_middleware import _claims, _presented
    from app.services.auth import api_key_service

    claims = _claims(request)
    if claims is not None:
        subject = claims.get("sub")
        raw_tenant = claims.get("tenant_id")
        try:
            user_id = uuid.UUID(str(subject)) if subject else None
            tenant_id = uuid.UUID(str(raw_tenant)) if raw_tenant else None
        except ValueError:
            return None
        return (tenant_id, user_id)

    presented = _presented(request)
    if api_key_service.looks_like_a_key(presented):
        from app.config.database import SessionLocal

        db = SessionLocal()
        try:
            key = api_key_service.resolve(db, presented)
            return (key.tenant_id, key.user_id) if key else None
        except Exception:
            logger.exception("could not resolve an API key for idempotency")
            return None
        finally:
            db.close()

    return None


class IdempotencyMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next):
        key = request.headers.get(HEADER)
        if not key or request.method not in METHODS:
            return await call_next(request)

        if len(key) > idempotency_service.MAX_KEY_LENGTH:
            return JSONResponse(
                {"detail": f"{HEADER} must be at most "
                           f"{idempotency_service.MAX_KEY_LENGTH} characters"},
                status_code=400,
            )

        caller = _caller(request)
        if caller is None:
            return await call_next(request)

        tenant_id, user_id = caller
        endpoint = f"{request.method} {request.url.path}"

        body = await request.body()

        async def _receive():
            return {"type": "http.request", "body": body, "more_body": False}

        request = Request(request.scope, _receive)

        from app.config.database import SessionLocal

        db = SessionLocal()
        try:
            replay = idempotency_service.claim(
                db,
                key=key,
                endpoint=endpoint,
                request_hash=idempotency_service.hash_request(body),
                tenant_id=tenant_id,
                user_id=user_id,
            )
        except idempotency_service.Conflict as conflict:
            db.close()
            return JSONResponse({"detail": conflict.detail},
                                status_code=conflict.status_code)
        except Exception:
            logger.exception("idempotency claim failed; proceeding without it")
            db.close()
            return await call_next(request)

        if replay is not None:
            db.close()
            return JSONResponse(
                replay.response_body,
                status_code=replay.response_status or 200,
                headers={"Idempotent-Replay": "true"},
            )

        db.commit()

        try:
            response = await call_next(request)
        except Exception:
            idempotency_service.release(
                db, key=key, endpoint=endpoint,
                tenant_id=tenant_id, user_id=user_id,
            )
            db.close()
            raise

        payload = b""
        async for chunk in response.body_iterator:
            payload += chunk

        try:
            idempotency_service.complete(
                db, key=key, endpoint=endpoint,
                tenant_id=tenant_id, user_id=user_id,
                status_code=response.status_code, body=payload,
            )
        except Exception:
            logger.exception("could not record idempotent response")
        finally:
            db.close()

        return Response(
            content=payload,
            status_code=response.status_code,
            headers=dict(response.headers),
            media_type=response.media_type,
        )