svc-infra 0.1.600__py3-none-any.whl → 0.1.640__py3-none-any.whl

svc_infra/api/fastapi/dependencies/ratelimit.py

@@ -1,12 +1,17 @@
 from __future__ import annotations
 
 import time
-from typing import Callable
+from typing import Callable, Optional
 
 from fastapi import HTTPException
 from starlette.requests import Request
 
 from svc_infra.api.fastapi.middleware.ratelimit_store import InMemoryRateLimitStore, RateLimitStore
+
+try:
+    from svc_infra.api.fastapi.tenancy.context import resolve_tenant_id as _resolve_tenant_id
+except Exception:  # pragma: no cover - minimal builds
+    _resolve_tenant_id = None  # type: ignore
 from svc_infra.obs.metrics import emit_rate_limited
 
 
@@ -17,20 +22,44 @@ class RateLimiter:
         limit: int,
         window: int = 60,
         key_fn: Callable = lambda r: "global",
+        limit_resolver: Optional[Callable[[Request, Optional[str]], Optional[int]]] = None,
+        scope_by_tenant: bool = False,
         store: RateLimitStore | None = None,
     ):
         self.limit = limit
         self.window = window
         self.key_fn = key_fn
+        self._limit_resolver = limit_resolver
+        self.scope_by_tenant = scope_by_tenant
         self.store = store or InMemoryRateLimitStore(limit=limit)
 
     async def __call__(self, request: Request):
+        # Try resolving tenant when asked
+        tenant_id = None
+        if self.scope_by_tenant or self._limit_resolver:
+            try:
+                if _resolve_tenant_id is not None:
+                    tenant_id = await _resolve_tenant_id(request)
+            except Exception:
+                tenant_id = None
+
         key = self.key_fn(request)
-        count, limit, reset = self.store.incr(str(key), self.window)
-        if count > limit:
+        if self.scope_by_tenant and tenant_id:
+            key = f"{key}:tenant:{tenant_id}"
+
+        eff_limit = self.limit
+        if self._limit_resolver:
+            try:
+                v = self._limit_resolver(request, tenant_id)
+                eff_limit = int(v) if v is not None else self.limit
+            except Exception:
+                eff_limit = self.limit
+
+        count, store_limit, reset = self.store.incr(str(key), self.window)
+        if count > eff_limit:
             retry = max(0, reset - int(time.time()))
             try:
-                emit_rate_limited(str(key), limit, retry)
+                emit_rate_limited(str(key), eff_limit, retry)
             except Exception:
                 pass
             raise HTTPException(
@@ -46,17 +75,38 @@ def rate_limiter(
     limit: int,
     window: int = 60,
     key_fn: Callable = lambda r: "global",
+    limit_resolver: Optional[Callable[[Request, Optional[str]], Optional[int]]] = None,
+    scope_by_tenant: bool = False,
     store: RateLimitStore | None = None,
 ):
     store_ = store or InMemoryRateLimitStore(limit=limit)
 
     async def dep(request: Request):
+        tenant_id = None
+        if scope_by_tenant or limit_resolver:
+            try:
+                if _resolve_tenant_id is not None:
+                    tenant_id = await _resolve_tenant_id(request)
+            except Exception:
+                tenant_id = None
+
         key = key_fn(request)
-        count, lim, reset = store_.incr(str(key), window)
-        if count > lim:
+        if scope_by_tenant and tenant_id:
+            key = f"{key}:tenant:{tenant_id}"
+
+        eff_limit = limit
+        if limit_resolver:
+            try:
+                v = limit_resolver(request, tenant_id)
+                eff_limit = int(v) if v is not None else limit
+            except Exception:
+                eff_limit = limit
+
+        count, _store_limit, reset = store_.incr(str(key), window)
+        if count > eff_limit:
             retry = max(0, reset - int(time.time()))
             try:
-                emit_rate_limited(str(key), lim, retry)
+                emit_rate_limited(str(key), eff_limit, retry)
             except Exception:
                 pass
             raise HTTPException(

svc_infra/api/fastapi/docs/add.py

@@ -0,0 +1,160 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Optional
+
+from fastapi import FastAPI, Request
+from fastapi.openapi.docs import get_redoc_html, get_swagger_ui_html
+from fastapi.responses import HTMLResponse, JSONResponse
+
+from .landing import CardSpec, DocTargets, render_index_html
+from .scoped import DOC_SCOPES
+
+
+def add_docs(
+    app: FastAPI,
+    *,
+    redoc_url: str = "/redoc",
+    swagger_url: str = "/docs",
+    openapi_url: str = "/openapi.json",
+    export_openapi_to: Optional[str] = None,
+    # Landing page options
+    landing_url: str = "/",
+    include_landing: bool = True,
+) -> None:
+    """Enable docs endpoints and optionally export OpenAPI schema to disk on startup.
+
+    We mount docs and OpenAPI routes explicitly so this works even when configured post-init.
+    """
+
+    # OpenAPI JSON route
+    async def openapi_handler() -> JSONResponse:  # noqa: ANN201
+        return JSONResponse(app.openapi())
+
+    app.add_api_route(openapi_url, openapi_handler, methods=["GET"], include_in_schema=False)
+
+    # Swagger UI route
+    async def swagger_ui(request: Request) -> HTMLResponse:  # noqa: ANN201
+        resp = get_swagger_ui_html(openapi_url=openapi_url, title="API Docs")
+        theme = request.query_params.get("theme")
+        if theme == "dark":
+            return _with_dark_mode(resp)
+        return resp
+
+    app.add_api_route(swagger_url, swagger_ui, methods=["GET"], include_in_schema=False)
+
+    # Redoc route
+    async def redoc_ui(request: Request) -> HTMLResponse:  # noqa: ANN201
+        resp = get_redoc_html(openapi_url=openapi_url, title="API ReDoc")
+        theme = request.query_params.get("theme")
+        if theme == "dark":
+            return _with_dark_mode(resp)
+        return resp
+
+    app.add_api_route(redoc_url, redoc_ui, methods=["GET"], include_in_schema=False)
+
+    # Optional export to disk on startup
+    if export_openapi_to:
+        export_path = Path(export_openapi_to)
+
+        async def _export_docs() -> None:
+            # Startup export
+            spec = app.openapi()
+            export_path.parent.mkdir(parents=True, exist_ok=True)
+            export_path.write_text(json.dumps(spec, indent=2))
+
+        app.add_event_handler("startup", _export_docs)
+
+    # Optional landing page with the same look/feel as setup_service_api
+    if include_landing:
+        # Avoid path collision; if landing_url is already taken for GET, fallback to "/_docs"
+        existing_paths = {
+            (getattr(r, "path", None) or getattr(r, "path_format", None))
+            for r in getattr(app, "routes", [])
+            if getattr(r, "methods", None) and "GET" in r.methods
+        }
+        landing_path = landing_url or "/"
+        if landing_path in existing_paths:
+            landing_path = "/_docs"
+
+        async def _landing() -> HTMLResponse:  # noqa: ANN201
+            cards: list[CardSpec] = []
+            # Root docs card using the provided paths
+            cards.append(
+                CardSpec(
+                    tag="",
+                    docs=DocTargets(swagger=swagger_url, redoc=redoc_url, openapi_json=openapi_url),
+                )
+            )
+            # Scoped docs (if any were registered via add_prefixed_docs)
+            for scope, swagger, redoc, openapi_json, _title in DOC_SCOPES:
+                cards.append(
+                    CardSpec(
+                        tag=scope.strip("/"),
+                        docs=DocTargets(swagger=swagger, redoc=redoc, openapi_json=openapi_json),
+                    )
+                )
+            html = render_index_html(
+                service_name=app.title or "API", release=app.version or "", cards=cards
+            )
+            return HTMLResponse(html)
+
+        app.add_api_route(landing_path, _landing, methods=["GET"], include_in_schema=False)
+
+
+def _with_dark_mode(resp: HTMLResponse) -> HTMLResponse:
+    """Return a copy of the HTMLResponse with a minimal dark-theme CSS injected.
+
+    We avoid depending on custom Swagger/ReDoc builds; this works by inlining a small CSS
+    block and toggling a `.dark` class on the body element.
+    """
+    try:
+        body = resp.body.decode("utf-8", errors="ignore")
+    except Exception:  # pragma: no cover - very unlikely
+        return resp
+
+    css = _DARK_CSS
+    if "</head>" in body:
+        body = body.replace("</head>", f"<style>\n{css}\n</style></head>", 1)
+    # add class to body to allow stronger selectors
+    body = body.replace("<body>", '<body class="dark">', 1)
+    return HTMLResponse(content=body, status_code=resp.status_code, headers=dict(resp.headers))
+
+
+_DARK_CSS = """
+/* Minimal dark mode override for Swagger/ReDoc */
+@media (prefers-color-scheme: dark) { :root { color-scheme: dark; } }
+html.dark, body.dark { background: #0b0e14; color: #e0e6f1; }
+#swagger, .redoc-wrap { background: transparent; }
+a { color: #62aef7; }
+"""
+
+
+def add_sdk_generation_stub(
+    app: FastAPI,
+    *,
+    on_generate: Optional[callable] = None,
+    openapi_path: str = "/openapi.json",
+) -> None:
+    """Hook to add an SDK generation stub.
+
+    Provide `on_generate()` to run generation (e.g., openapi-generator). This is a stub only; we
+    don't ship a hard dependency. If `on_generate` is provided, we expose `/_docs/generate-sdk`.
+    """
+    from svc_infra.api.fastapi.dual.public import public_router
+
+    if not on_generate:
+        return
+
+    router = public_router(prefix="/_docs", include_in_schema=False)
+
+    @router.post("/generate-sdk")
+    async def _generate() -> dict:  # noqa: ANN201
+        on_generate()
+        return {"status": "ok"}
+
+    app.include_router(router)
+
+
+__all__ = ["add_docs", "add_sdk_generation_stub"]

svc_infra/api/fastapi/docs/landing.py

@@ -115,7 +115,7 @@ def render_index_html(*, service_name: str, release: str, cards: Iterable[CardSp
     <section class="grid">
       {grid}
     </section>
-    <footer>Tip: each card exposes Swagger, ReDoc, and a pretty JSON view.</footer>
+    <footer>Tip: each card exposes Swagger, ReDoc, and a JSON view.</footer>
   </div>
 </body>
 </html>

svc_infra/api/fastapi/middleware/errors/handlers.py

@@ -2,6 +2,7 @@ import logging
 import traceback
 from typing import Any, Dict, Optional
 
+import httpx
 from fastapi import Request
 from fastapi.exceptions import HTTPException, RequestValidationError
 from fastapi.responses import JSONResponse, Response
@@ -46,6 +47,7 @@ def problem_response(
     code: str | None = None,
     errors: list[dict] | None = None,
     trace_id: str | None = None,
+    headers: dict[str, str] | None = None,
 ) -> Response:
     body: Dict[str, Any] = {
         "type": type_uri,
@@ -62,10 +64,24 @@ def problem_response(
         body["errors"] = errors
     if trace_id:
         body["trace_id"] = trace_id
-    return JSONResponse(status_code=status, content=body, media_type=PROBLEM_MT)
+    return JSONResponse(status_code=status, content=body, media_type=PROBLEM_MT, headers=headers)
 
 
 def register_error_handlers(app):
+    @app.exception_handler(httpx.TimeoutException)
+    async def handle_httpx_timeout(request: Request, exc: httpx.TimeoutException):
+        trace_id = _trace_id_from_request(request)
+        # Map outbound HTTP client timeouts to 504 Gateway Timeout
+        # Keep details generic in prod
+        return problem_response(
+            status=504,
+            title="Gateway Timeout",
+            detail=("Upstream request timed out." if IS_PROD else (str(exc) or "httpx timeout")),
+            code="GATEWAY_TIMEOUT",
+            instance=str(request.url),
+            trace_id=trace_id,
+        )
+
     @app.exception_handler(FastApiException)
     async def handle_app_exception(request: Request, exc: FastApiException):
         trace_id = _trace_id_from_request(request)
@@ -104,14 +120,25 @@ def register_error_handlers(app):
     @app.exception_handler(HTTPException)
     async def handle_http_exception(request: Request, exc: HTTPException):
         trace_id = _trace_id_from_request(request)
-        title = {401: "Unauthorized", 403: "Forbidden", 404: "Not Found"}.get(
-            exc.status_code, "Error"
-        )
+        title = {
+            401: "Unauthorized",
+            403: "Forbidden",
+            404: "Not Found",
+            429: "Too Many Requests",
+        }.get(exc.status_code, "Error")
         detail = (
             exc.detail
             if not IS_PROD or exc.status_code < 500
             else "Something went wrong. Please contact support."
         )
+        # Preserve headers set on the exception (e.g., Retry-After for rate limits)
+        hdrs: dict[str, str] | None = None
+        try:
+            if getattr(exc, "headers", None):
+                # FastAPI/Starlette exceptions store headers as a dict[str, str]
+                hdrs = dict(getattr(exc, "headers"))  # type: ignore[arg-type]
+        except Exception:
+            hdrs = None
         return problem_response(
             status=exc.status_code,
             title=title,
@@ -119,19 +146,29 @@ def register_error_handlers(app):
             code=title.replace(" ", "_").upper(),
             instance=str(request.url),
             trace_id=trace_id,
+            headers=hdrs,
         )
 
     @app.exception_handler(StarletteHTTPException)
     async def handle_starlette_http_exception(request: Request, exc: StarletteHTTPException):
         trace_id = _trace_id_from_request(request)
-        title = {401: "Unauthorized", 403: "Forbidden", 404: "Not Found"}.get(
-            exc.status_code, "Error"
-        )
+        title = {
+            401: "Unauthorized",
+            403: "Forbidden",
+            404: "Not Found",
+            429: "Too Many Requests",
+        }.get(exc.status_code, "Error")
         detail = (
             exc.detail
             if not IS_PROD or exc.status_code < 500
             else "Something went wrong. Please contact support."
         )
+        hdrs: dict[str, str] | None = None
+        try:
+            if getattr(exc, "headers", None):
+                hdrs = dict(getattr(exc, "headers"))  # type: ignore[arg-type]
+        except Exception:
+            hdrs = None
         return problem_response(
             status=exc.status_code,
             title=title,
@@ -139,6 +176,7 @@ def register_error_handlers(app):
             code=title.replace(" ", "_").upper(),
             instance=str(request.url),
             trace_id=trace_id,
+            headers=hdrs,
         )
 
     @app.exception_handler(IntegrityError)

svc_infra/api/fastapi/middleware/graceful_shutdown.py

@@ -0,0 +1,87 @@
+from __future__ import annotations
+
+import asyncio
+import logging
+import os
+from contextlib import asynccontextmanager
+from typing import Optional
+
+from fastapi import FastAPI
+from starlette.types import ASGIApp, Receive, Scope, Send
+
+from svc_infra.app.env import pick
+
+logger = logging.getLogger(__name__)
+
+
+def _get_grace_period_seconds() -> float:
+    default = pick(prod=20.0, nonprod=5.0)
+    raw = os.getenv("SHUTDOWN_GRACE_PERIOD_SECONDS")
+    if raw is None or raw == "":
+        return float(default)
+    try:
+        return float(raw)
+    except ValueError:
+        return float(default)
+
+
+class InflightTrackerMiddleware:
+    """Tracks number of in-flight requests to support graceful shutdown drains."""
+
+    def __init__(self, app: ASGIApp):
+        self.app = app
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send):
+        if scope.get("type") != "http":
+            await self.app(scope, receive, send)
+            return
+        state = scope.get("app").state  # type: ignore[attr-defined]
+        state._inflight_requests = getattr(state, "_inflight_requests", 0) + 1
+        try:
+            await self.app(scope, receive, send)
+        finally:
+            state._inflight_requests = max(0, getattr(state, "_inflight_requests", 1) - 1)
+
+
+async def _wait_for_drain(app: FastAPI, grace: float) -> None:
+    interval = 0.1
+    waited = 0.0
+    while waited < grace:
+        inflight = int(getattr(app.state, "_inflight_requests", 0))
+        if inflight <= 0:
+            return
+        await asyncio.sleep(interval)
+        waited += interval
+    inflight = int(getattr(app.state, "_inflight_requests", 0))
+    if inflight > 0:
+        logger.warning(
+            "Graceful shutdown timeout: %s in-flight request(s) after %.2fs", inflight, waited
+        )
+
+
+def install_graceful_shutdown(app: FastAPI, *, grace_seconds: Optional[float] = None) -> None:
+    """Install inflight tracking and lifespan hooks to wait for requests to drain.
+
+    - Adds InflightTrackerMiddleware
+    - Registers a lifespan handler that initializes state and waits up to grace_seconds on shutdown
+    """
+    app.add_middleware(InflightTrackerMiddleware)
+
+    g = float(grace_seconds) if grace_seconds is not None else _get_grace_period_seconds()
+
+    # Preserve any existing lifespan and wrap it so our drain runs on shutdown.
+    previous_lifespan = getattr(app.router, "lifespan_context", None)
+
+    @asynccontextmanager
+    async def _lifespan(a: FastAPI):  # noqa: ANN202
+        # Startup: initialize inflight counter
+        a.state._inflight_requests = 0
+        if previous_lifespan is not None:
+            async with previous_lifespan(a):
+                yield
+        else:
+            yield
+        # Shutdown: wait for in-flight requests to drain (up to grace period)
+        await _wait_for_drain(a, g)
+
+    app.router.lifespan_context = _lifespan

svc_infra/api/fastapi/middleware/ratelimit.py

@@ -7,6 +7,12 @@ from svc_infra.obs.metrics import emit_rate_limited
 
 from .ratelimit_store import InMemoryRateLimitStore, RateLimitStore
 
+try:
+    # Optional import: tenancy may not be enabled in all apps
+    from svc_infra.api.fastapi.tenancy.context import resolve_tenant_id as _resolve_tenant_id
+except Exception:  # pragma: no cover - fallback for minimal builds
+    _resolve_tenant_id = None  # type: ignore
+
 
 class SimpleRateLimitMiddleware(BaseHTTPMiddleware):
     def __init__(
@@ -15,18 +21,70 @@ class SimpleRateLimitMiddleware(BaseHTTPMiddleware):
         limit: int = 120,
         window: int = 60,
         key_fn=None,
+        *,
+        # When provided, dynamically computes a limit for the current request (e.g. per-tenant quotas)
+        # Signature: (request: Request, tenant_id: Optional[str]) -> int | None
+        limit_resolver=None,
+        # If True, automatically scopes the bucket key by tenant id when available
+        scope_by_tenant: bool = False,
+        # When True, allows unresolved tenant IDs to fall back to an "X-Tenant-Id" header value.
+        # Disabled by default to avoid trusting arbitrary client-provided headers which could
+        # otherwise be used to evade per-tenant limits when authentication fails.
+        allow_untrusted_tenant_header: bool = False,
         store: RateLimitStore | None = None,
     ):
         super().__init__(app)
         self.limit, self.window = limit, window
         self.key_fn = key_fn or (lambda r: r.headers.get("X-API-Key") or r.client.host)
+        self._limit_resolver = limit_resolver
+        self.scope_by_tenant = scope_by_tenant
+        self._allow_untrusted_tenant_header = allow_untrusted_tenant_header
         self.store = store or InMemoryRateLimitStore(limit=limit)
 
     async def dispatch(self, request, call_next):
+        # Resolve tenant when possible
+        tenant_id = None
+        if self.scope_by_tenant or self._limit_resolver:
+            try:
+                if _resolve_tenant_id is not None:
+                    tenant_id = await _resolve_tenant_id(request)
+            except Exception:
+                tenant_id = None
+            # Fallback header behavior:
+            # - If tenancy context is unavailable (minimal builds), accept header by default so
+            #   unit/integration tests can exercise per-tenant scoping without full auth state.
+            # - If tenancy is available, only trust the header when explicitly allowed.
+            if not tenant_id:
+                if _resolve_tenant_id is None:
+                    tenant_id = request.headers.get("X-Tenant-Id") or request.headers.get(
+                        "X-Tenant-ID"
+                    )
+                elif self._allow_untrusted_tenant_header:
+                    tenant_id = request.headers.get("X-Tenant-Id") or request.headers.get(
+                        "X-Tenant-ID"
+                    )
+
         key = self.key_fn(request)
+        if self.scope_by_tenant and tenant_id:
+            key = f"{key}:tenant:{tenant_id}"
+
+        # Allow dynamic limit overrides
+        eff_limit = self.limit
+        if self._limit_resolver:
+            try:
+                v = self._limit_resolver(request, tenant_id)
+                eff_limit = int(v) if v is not None else self.limit
+            except Exception:
+                eff_limit = self.limit
+
         now = int(time.time())
         # Increment counter in store
-        count, limit, reset = self.store.incr(str(key), self.window)
+        # Update store limit if it differs; stores capture configured limit internally
+        # For in-memory store, we can temporarily adjust per-request by swapping a new store instance
+        # but to keep API simple, we reuse store and clamp by eff_limit below.
+        count, store_limit, reset = self.store.incr(str(key), self.window)
+        # Enforce the effective limit selected for this request
+        limit = eff_limit
         remaining = max(0, limit - count)
 
         if remaining < 0:  # defensive clamp

svc_infra/api/fastapi/middleware/ratelimit_store.py

@@ -16,14 +16,20 @@ class RateLimitStore(Protocol):
 class InMemoryRateLimitStore:
     def __init__(self, limit: int = 120):
         self.limit = limit
-        self._buckets: dict[tuple[str, int], int] = {}
+        # Track per-key rolling windows: key -> (count, window_start_epoch)
+        self._state: dict[str, tuple[int, float]] = {}
 
     def incr(self, key: str, window: int) -> Tuple[int, int, int]:
-        now = int(time.time())
-        win = now - (now % window)
-        count = self._buckets.get((key, win), 0) + 1
-        self._buckets[(key, win)] = count
-        reset = win + window
+        now = time.time()
+        count, window_start = self._state.get(key, (0, now))
+        # If outside the rolling window, reset
+        if now >= window_start + window:
+            count = 1
+            window_start = now
+        else:
+            count += 1
+        self._state[key] = (count, window_start)
+        reset = int(window_start + window)
         return count, self.limit, reset