svc-infra 0.1.595__py3-none-any.whl → 0.1.706__py3-none-any.whl

svc_infra/api/fastapi/auth/ws_security.py

@@ -0,0 +1,275 @@
+"""WebSocket authentication primitives.
+
+This module provides lightweight JWT-based authentication for WebSocket endpoints.
+Unlike HTTP auth which requires DB access, WS auth uses JWT claims only, making it
+suitable for high-frequency real-time connections.
+
+Usage:
+    from svc_infra.api.fastapi.auth.ws_security import WSIdentity
+
+    @router.websocket("/ws")
+    async def ws_handler(websocket: WebSocket, user: WSIdentity):
+        # user.id, user.email, user.scopes available from JWT claims
+        await websocket.accept()
+        ...
+
+For router-level dependencies (protects all endpoints):
+    from svc_infra.api.fastapi.auth.ws_security import RequireWSIdentity
+
+    router = DualAPIRouter(dependencies=[RequireWSIdentity])
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Annotated, Any, cast
+
+import jwt
+from fastapi import Depends, WebSocket, WebSocketException, status
+
+from svc_infra.api.fastapi.auth.settings import get_auth_settings
+
+
+# ---------- WSPrincipal ----------
+@dataclass
+class WSPrincipal:
+    """Lightweight principal for WebSocket connections.
+
+    Unlike the HTTP `Principal` which loads the full user from DB,
+    `WSPrincipal` contains only JWT claims. This makes it suitable
+    for high-frequency real-time connections without DB overhead.
+
+    Attributes:
+        id: User ID from JWT 'sub' claim (typically UUID string)
+        email: User email from JWT 'email' claim (if present)
+        scopes: List of scopes/permissions from JWT 'scopes' claim
+        claims: Full JWT payload for custom claim access
+        via: Authentication method ('query', 'header', 'subprotocol')
+    """
+
+    id: str
+    email: str | None = None
+    scopes: list[str] = field(default_factory=list)
+    claims: dict = field(default_factory=dict)
+    via: str = "query"  # 'query' | 'header' | 'subprotocol'
+
+
+# ---------- Token extraction ----------
+def _extract_token(websocket: WebSocket) -> tuple[str | None, str]:
+    """Extract JWT token from WebSocket connection.
+
+    Tries extraction in order:
+    1. Query parameter: ?token=xxx
+    2. Authorization header: Bearer xxx
+    3. Sec-WebSocket-Protocol header (for browser clients that can't set headers)
+
+    Returns:
+        Tuple of (token, source) where source is 'query', 'header', or 'subprotocol'
+    """
+    # 1. Query parameter (most common for WebSocket)
+    token = websocket.query_params.get("token")
+    if token:
+        return token.strip(), "query"
+
+    # 2. Authorization header
+    auth_header = websocket.headers.get("authorization", "")
+    if auth_header.lower().startswith("bearer "):
+        token = auth_header.split(" ", 1)[1].strip()
+        if token:
+            return token, "header"
+
+    # 3. Sec-WebSocket-Protocol (browser workaround)
+    # Some clients send token as: Sec-WebSocket-Protocol: bearer, <token>
+    protocol = websocket.headers.get("sec-websocket-protocol", "")
+    if protocol:
+        parts = [p.strip() for p in protocol.split(",")]
+        # Look for token after 'bearer' protocol
+        for i, part in enumerate(parts):
+            if part.lower() == "bearer" and i + 1 < len(parts):
+                return parts[i + 1], "subprotocol"
+
+    return None, ""
+
+
+def _decode_jwt(token: str) -> dict:
+    """Decode and validate JWT token.
+
+    Uses the same JWT settings as HTTP auth (AUTH_JWT__SECRET).
+    Supports key rotation via old_secrets.
+
+    Returns:
+        JWT payload dict
+
+    Raises:
+        WebSocketException: If token is invalid or expired
+    """
+    settings = get_auth_settings()
+
+    if not settings.jwt:
+        raise WebSocketException(
+            code=status.WS_1008_POLICY_VIOLATION,
+            reason="JWT not configured",
+        )
+
+    secret = settings.jwt.secret.get_secret_value()
+    old_secrets = [s.get_secret_value() for s in (settings.jwt.old_secrets or [])]
+    all_secrets = [secret] + old_secrets
+
+    last_error: Exception | None = None
+
+    for s in all_secrets:
+        try:
+            payload = jwt.decode(
+                token,
+                s,
+                algorithms=["HS256"],
+                options={"require": ["sub", "exp"]},
+            )
+            return cast(dict[Any, Any], payload)
+        except jwt.ExpiredSignatureError:
+            raise WebSocketException(
+                code=status.WS_1008_POLICY_VIOLATION,
+                reason="Token expired",
+            )
+        except jwt.InvalidTokenError as e:
+            last_error = e
+            continue
+
+    # None of the secrets worked
+    raise WebSocketException(
+        code=status.WS_1008_POLICY_VIOLATION,
+        reason=f"Invalid token: {last_error}",
+    )
+
+
+# ---------- Resolvers ----------
+async def resolve_ws_bearer_principal(websocket: WebSocket) -> WSPrincipal | None:
+    """Extract and validate JWT from WebSocket, returning WSPrincipal or None.
+
+    This is the optional resolver - returns None if no token present.
+    Use `_ws_current_principal` for required authentication.
+
+    Token sources (in order):
+        1. Query parameter: ?token=xxx
+        2. Authorization header: Bearer xxx
+        3. Sec-WebSocket-Protocol: bearer, xxx
+    """
+    token, source = _extract_token(websocket)
+    if not token:
+        return None
+
+    payload = _decode_jwt(token)
+
+    return WSPrincipal(
+        id=str(payload.get("sub", "")),
+        email=payload.get("email"),
+        scopes=payload.get("scopes", []) or payload.get("scope", "").split(),
+        claims=payload,
+        via=source,
+    )
+
+
+async def _ws_current_principal(
+    websocket: WebSocket,
+    principal: WSPrincipal | None = Depends(resolve_ws_bearer_principal),
+) -> WSPrincipal:
+    """Require authenticated WebSocket connection.
+
+    Use this as a dependency to require authentication.
+    Closes connection with 1008 (Policy Violation) if no valid token.
+    """
+    if not principal:
+        raise WebSocketException(
+            code=status.WS_1008_POLICY_VIOLATION,
+            reason="Missing or invalid authentication",
+        )
+    return principal
+
+
+async def _ws_optional_principal(
+    websocket: WebSocket,
+    principal: WSPrincipal | None = Depends(resolve_ws_bearer_principal),
+) -> WSPrincipal | None:
+    """Optional WebSocket authentication.
+
+    Returns None if no token present, WSPrincipal if valid token.
+    """
+    return principal
+
+
+# ---------- DX: types for endpoint params ----------
+WSIdentity = Annotated[WSPrincipal, Depends(_ws_current_principal)]
+"""Annotated type for required WebSocket authentication.
+
+Usage:
+    @router.websocket("/ws")
+    async def handler(websocket: WebSocket, user: WSIdentity):
+        # user.id, user.email, user.scopes available
+        ...
+"""
+
+OptionalWSIdentity = Annotated[WSPrincipal | None, Depends(_ws_optional_principal)]
+"""Annotated type for optional WebSocket authentication.
+
+Usage:
+    @router.websocket("/ws")
+    async def handler(websocket: WebSocket, user: OptionalWSIdentity):
+        if user:
+            # authenticated
+        else:
+            # anonymous
+        ...
+"""
+
+
+# ---------- DX: constants for router-level dependencies ----------
+RequireWSIdentity = Depends(_ws_current_principal)
+"""Router-level dependency for required WebSocket authentication.
+
+Usage:
+    router = DualAPIRouter(dependencies=[RequireWSIdentity])
+"""
+
+AllowWSIdentity = Depends(_ws_optional_principal)
+"""Router-level dependency for optional WebSocket authentication.
+
+Usage:
+    router = DualAPIRouter(dependencies=[AllowWSIdentity])
+"""
+
+
+# ---------- DX: guard factories ----------
+def RequireWSScopes(*needed: str):
+    """Require specific scopes for WebSocket connection.
+
+    Usage:
+        router = DualAPIRouter(dependencies=[RequireWSScopes("chat:read", "chat:write")])
+    """
+
+    async def _guard(principal: WSIdentity) -> WSPrincipal:
+        if not set(needed).issubset(set(principal.scopes or [])):
+            raise WebSocketException(
+                code=status.WS_1008_POLICY_VIOLATION,
+                reason="Insufficient scope",
+            )
+        return principal
+
+    return Depends(_guard)
+
+
+def RequireWSAnyScope(*candidates: str):
+    """Require at least one of the specified scopes.
+
+    Usage:
+        router = DualAPIRouter(dependencies=[RequireWSAnyScope("admin", "moderator")])
+    """
+
+    async def _guard(principal: WSIdentity) -> WSPrincipal:
+        if not set(principal.scopes or []) & set(candidates):
+            raise WebSocketException(
+                code=status.WS_1008_POLICY_VIOLATION,
+                reason="Insufficient scope",
+            )
+        return principal
+
+    return Depends(_guard)

svc_infra/api/fastapi/billing/router.py

@@ -0,0 +1,73 @@
+from __future__ import annotations
+
+from datetime import datetime, timezone
+from typing import Annotated, Optional
+
+from fastapi import APIRouter, Depends, Response, status
+
+from svc_infra.api.fastapi.db.sql.session import SqlSessionDep
+from svc_infra.api.fastapi.middleware.idempotency import require_idempotency_key
+from svc_infra.api.fastapi.tenancy.context import TenantId
+from svc_infra.billing.async_service import AsyncBillingService
+from svc_infra.billing.schemas import (
+    UsageAckOut,
+    UsageAggregateRow,
+    UsageAggregatesOut,
+    UsageIn,
+)
+
+router = APIRouter(prefix="/_billing", tags=["Billing"])
+
+
+def get_service(tenant_id: TenantId, session: SqlSessionDep) -> AsyncBillingService:
+    return AsyncBillingService(session=session, tenant_id=tenant_id)
+
+
+@router.post(
+    "/usage",
+    name="billing_record_usage",
+    status_code=status.HTTP_202_ACCEPTED,
+    response_model=UsageAckOut,
+    dependencies=[Depends(require_idempotency_key)],
+)
+async def record_usage(
+    data: UsageIn,
+    svc: Annotated[AsyncBillingService, Depends(get_service)],
+    response: Response,
+):
+    at = data.at or datetime.now(tz=timezone.utc)
+    evt_id = await svc.record_usage(
+        metric=data.metric,
+        amount=int(data.amount),
+        at=at,
+        idempotency_key=data.idempotency_key,
+        metadata=data.metadata,
+    )
+    # For 202, no Location header is required, but we can surface the id in the body
+    return UsageAckOut(id=evt_id, accepted=True)
+
+
+@router.get(
+    "/usage",
+    name="billing_list_aggregates",
+    response_model=UsageAggregatesOut,
+)
+async def list_aggregates(
+    metric: str,
+    date_from: Optional[datetime] = None,
+    date_to: Optional[datetime] = None,
+    svc: Annotated[AsyncBillingService, Depends(get_service)] = None,  # type: ignore[assignment]
+):
+    rows = await svc.list_daily_aggregates(
+        metric=metric, date_from=date_from, date_to=date_to
+    )
+    items = [
+        UsageAggregateRow(
+            period_start=r.period_start,
+            granularity=r.granularity,
+            metric=r.metric,
+            total=int(r.total),
+        )
+        for r in rows
+    ]
+    return UsageAggregatesOut(items=items, next_cursor=None)

svc_infra/api/fastapi/billing/setup.py

@@ -0,0 +1,19 @@
+from __future__ import annotations
+
+from fastapi import FastAPI
+
+from .router import router as billing_router
+
+
+def add_billing(app: FastAPI, *, prefix: str = "/_billing") -> None:
+    # Mount under the chosen prefix; default is /_billing
+    if prefix and prefix != "/_billing":
+        # If a custom prefix is desired, clone router with new prefix
+        from fastapi import APIRouter
+
+        custom = APIRouter(prefix=prefix, tags=["Billing"])
+        for route in billing_router.routes:
+            custom.routes.append(route)
+        app.include_router(custom)
+    else:
+        app.include_router(billing_router)

svc_infra/api/fastapi/cache/add.py

@@ -1,3 +1,5 @@
+from contextlib import asynccontextmanager
+
 from fastapi import FastAPI
 
 from svc_infra.cache.backend import shutdown_cache
@@ -5,10 +7,12 @@ from svc_infra.cache.decorators import init_cache
 
 
 def setup_caching(app: FastAPI) -> None:
-    @app.on_event("startup")
-    async def _startup():
+    @asynccontextmanager
+    async def lifespan(_app: FastAPI):
         init_cache()
+        try:
+            yield
+        finally:
+            await shutdown_cache()
 
-    @app.on_event("shutdown")
-    async def _shutdown():
-        await shutdown_cache()
+    app.router.lifespan_context = lifespan

svc_infra/api/fastapi/db/__init__.py

@@ -1,4 +1,8 @@
-from svc_infra.api.fastapi.db.nosql import add_mongo_db, add_mongo_health, add_mongo_resources
+from svc_infra.api.fastapi.db.nosql import (
+    add_mongo_db,
+    add_mongo_health,
+    add_mongo_resources,
+)
 from svc_infra.api.fastapi.db.sql import add_sql_db, add_sql_health, add_sql_resources
 
 __all__ = [

svc_infra/api/fastapi/db/http.py

@@ -25,7 +25,9 @@ class OrderParams(BaseModel):
 
 
 def dep_order(
-    order_by: Optional[str] = Query(None, description="Comma-separated fields; '-' for DESC"),
+    order_by: Optional[str] = Query(
+        None, description="Comma-separated fields; '-' for DESC"
+    ),
 ) -> OrderParams:
     return OrderParams(order_by=order_by)
 

svc_infra/api/fastapi/db/nosql/__init__.py

@@ -1,4 +1,42 @@
-from .mongo.add import add_mongo_db, add_mongo_health, add_mongo_resources
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from fastapi import FastAPI
+
+    from svc_infra.db.nosql.resource import NoSqlResource
+
+
+def _missing_mongo_dependency() -> ModuleNotFoundError:
+    return ModuleNotFoundError(
+        "MongoDB support is an optional dependency. Install pymongo (and motor) to use "
+        "Mongo helpers like add_mongo_db/add_mongo_health/add_mongo_resources."
+    )
+
+
+try:
+    from .mongo.add import add_mongo_db, add_mongo_health, add_mongo_resources
+except ModuleNotFoundError as exc:
+    mongo_import_error = exc
+
+    # NOTE: pymongo provides `bson`, which can be absent in minimal installs/CI.
+    # We keep imports working for non-mongo users/tests by providing stubs.
+    def add_mongo_db(app: FastAPI, *, dsn_env: str = "MONGO_URL") -> None:
+        raise _missing_mongo_dependency() from mongo_import_error
+
+    def add_mongo_health(
+        app: FastAPI,
+        *,
+        prefix: str = "/_mongo/health",
+        include_in_schema: bool = False,
+    ) -> None:
+        raise _missing_mongo_dependency() from mongo_import_error
+
+    def add_mongo_resources(app: FastAPI, resources: Sequence[NoSqlResource]) -> None:
+        raise _missing_mongo_dependency() from mongo_import_error
+
 
 __all__ = [
     # MongoDB

svc_infra/api/fastapi/db/nosql/mongo/add.py

@@ -24,12 +24,15 @@ from .health import make_mongo_health_router
 def add_mongo_db_with_url(app: FastAPI, url: str, db_name: str) -> None:
     @asynccontextmanager
     async def lifespan(_app: FastAPI):
-        await init_mongo(MongoSettings(url=url, db_name=db_name))
+        # MongoSettings expects url as AnyUrl, which can be constructed from str via Pydantic
+        await init_mongo(MongoSettings(url=url, db_name=db_name))  # type: ignore[arg-type]  # Pydantic coerces str to AnyUrl
         try:
             expected = get_mongo_dbname_from_env(required=False)
             db = await acquire_db()
             if expected and db.name != expected:
-                raise RuntimeError(f"Connected to Mongo DB '{db.name}', expected '{expected}'.")
+                raise RuntimeError(
+                    f"Connected to Mongo DB '{db.name}', expected '{expected}'."
+                )
             yield
         finally:
             await close_mongo()
@@ -38,19 +41,23 @@ def add_mongo_db_with_url(app: FastAPI, url: str, db_name: str) -> None:
 
 
 def add_mongo_db(app: FastAPI, *, dsn_env: str = "MONGO_URL") -> None:
-    @app.on_event("startup")
-    async def _startup() -> None:
+    @asynccontextmanager
+    async def lifespan(_app: FastAPI):
         if not os.getenv(dsn_env):
             raise RuntimeError(f"Missing environment variable {dsn_env} for Mongo URL")
         await init_mongo()
         expected = get_mongo_dbname_from_env(required=False)
         db = await acquire_db()
         if expected and db.name != expected:
-            raise RuntimeError(f"Connected to Mongo DB '{db.name}', expected '{expected}'.")
+            raise RuntimeError(
+                f"Connected to Mongo DB '{db.name}', expected '{expected}'."
+            )
+        try:
+            yield
+        finally:
+            await close_mongo()
 
-    @app.on_event("shutdown")
-    async def _shutdown() -> None:
-        await close_mongo()
+    app.router.lifespan_context = lifespan
 
 
 def add_mongo_health(
@@ -58,50 +65,58 @@ def add_mongo_health(
 ) -> None:
     if include_in_schema is None:
         include_in_schema = CURRENT_ENVIRONMENT == LOCAL_ENV
-    app.include_router(make_mongo_health_router(prefix=prefix, include_in_schema=include_in_schema))
+    app.include_router(
+        make_mongo_health_router(prefix=prefix, include_in_schema=include_in_schema)
+    )
 
 
 def add_mongo_resources(app: FastAPI, resources: Sequence[NoSqlResource]) -> None:
-    for r in resources:
+    for resource in resources:
         repo = NoSqlRepository(
-            collection_name=r.resolved_collection(),
-            id_field=r.id_field,
-            soft_delete=r.soft_delete,
-            soft_delete_field=r.soft_delete_field,
-            soft_delete_flag_field=r.soft_delete_flag_field,
+            collection_name=resource.resolved_collection(),
+            id_field=resource.id_field,
+            soft_delete=resource.soft_delete,
+            soft_delete_field=resource.soft_delete_field,
+            soft_delete_flag_field=resource.soft_delete_flag_field,
+        )
+        svc = (
+            resource.service_factory(repo)
+            if resource.service_factory
+            else NoSqlService(repo)
         )
-        svc = r.service_factory(repo) if r.service_factory else NoSqlService(repo)
 
-        if r.read_schema and r.create_schema and r.update_schema:
-            Read, Create, Update = r.read_schema, r.create_schema, r.update_schema
-        elif r.document_model is not None:
+        if resource.read_schema and resource.create_schema and resource.update_schema:
+            Read, Create, Update = (
+                resource.read_schema,
+                resource.create_schema,
+                resource.update_schema,
+            )
+        elif resource.document_model is not None:
             # CRITICAL: teach Pydantic to dump ObjectId/PyObjectId
             Read, Create, Update = make_document_crud_schemas(
-                r.document_model,
-                create_exclude=r.create_exclude,
-                read_name=r.read_name,
-                create_name=r.create_name,
-                update_name=r.update_name,
-                read_exclude=r.read_exclude,
-                update_exclude=r.update_exclude,
+                resource.document_model,
+                create_exclude=resource.create_exclude,
+                read_name=resource.read_name,
+                create_name=resource.create_name,
+                update_name=resource.update_name,
+                read_exclude=resource.read_exclude,
+                update_exclude=resource.update_exclude,
                 json_encoders={ObjectId: str, PyObjectId: str},
             )
         else:
             raise RuntimeError(
-                f"Resource for collection '{r.collection}' requires either explicit schemas "
+                f"Resource for collection '{resource.collection}' requires either explicit schemas "
                 f"(read/create/update) or a 'document_model' to derive them."
             )
 
         router = make_crud_router_plus_mongo(
-            collection=r.resolved_collection(),
-            repo=repo,
             service=svc,
             read_schema=Read,
             create_schema=Create,
             update_schema=Update,
-            prefix=r.prefix,
-            tags=r.tags,
-            search_fields=r.search_fields,
+            prefix=resource.prefix,
+            tags=resource.tags,
+            search_fields=resource.search_fields,
             default_ordering=None,
             allowed_order_fields=None,
         )

svc_infra/api/fastapi/db/nosql/mongo/crud_router.py

@@ -1,7 +1,14 @@
 from typing import Annotated, Any, Optional, Sequence, Type, cast
 
 from fastapi import APIRouter, Body, Depends, HTTPException
-from motor.motor_asyncio import AsyncIOMotorDatabase
+
+try:
+    from motor.motor_asyncio import AsyncIOMotorDatabase
+
+    HAS_MOTOR = True
+except ImportError:  # pragma: no cover
+    HAS_MOTOR = False
+    AsyncIOMotorDatabase = Any  # type: ignore[assignment, misc]
 
 from svc_infra.api.fastapi.db.http import (
     LimitOffsetParams,
@@ -46,6 +53,9 @@ def make_crud_router_plus_mongo(
     allowed_order_fields: Optional[list[str]] = None,
     mount_under_db_prefix: bool = True,
 ) -> APIRouter:
+    read_model = cast(Any, read_schema)
+    page_model = cast(Any, Page[read_schema])  # type: ignore[valid-type]
+
     router_prefix = ("/_mongo" + prefix) if mount_under_db_prefix else prefix
     router = public_router(
         prefix=router_prefix,
@@ -56,7 +66,7 @@ def make_crud_router_plus_mongo(
     # LIST
     @router.get(
         "",
-        response_model=cast(Any, Page[read_schema]),
+        response_model=page_model,
         description=f"List items in {prefix} collection",
     )
     async def list_items(
@@ -68,20 +78,25 @@ def make_crud_router_plus_mongo(
         sort = _parse_sort(op.order_by or default_ordering, allowed_order_fields)
         if sp.q and search_fields:
             items = await service.search(
-                db, q=sp.q, fields=search_fields, limit=lp.limit, offset=lp.offset, sort=sort
+                db,
+                q=sp.q,
+                fields=search_fields,
+                limit=lp.limit,
+                offset=lp.offset,
+                sort=sort,
             )
             total = await service.count_filtered(db, q=sp.q, fields=search_fields)
         else:
             items = await service.list(db, limit=lp.limit, offset=lp.offset, sort=sort)
             total = await service.count(db)
-        return Page[read_schema].from_items(
+        return Page[Any].from_items(
             total=total, items=items, limit=lp.limit, offset=lp.offset
         )
 
     # GET by id
     @router.get(
         "/{item_id}",
-        response_model=cast(Any, read_schema),
+        response_model=read_model,
         description=f"Get item from {prefix} collection",
     )
     async def get_item(db: DBDep, item_id: Any):
@@ -93,22 +108,26 @@ def make_crud_router_plus_mongo(
     # CREATE
     @router.post(
         "",
-        response_model=cast(Any, read_schema),
+        response_model=read_model,
         status_code=201,
         description=f"Create item in {prefix} collection",
     )
-    async def create_item(db: DBDep, payload: create_schema = Body(...)):
-        data = payload.model_dump(exclude_unset=True)
+    async def create_item(db: DBDep, payload: create_schema = Body(...)):  # type: ignore[valid-type]
+        data = cast(Any, payload).model_dump(exclude_unset=True)
         return await service.create(db, data)
 
     # UPDATE
     @router.patch(
         "/{item_id}",
-        response_model=cast(Any, read_schema),
+        response_model=read_model,
         description=f"Update item in {prefix} collection",
     )
-    async def update_item(db: DBDep, item_id: Any, payload: update_schema = Body(...)):
-        data = payload.model_dump(exclude_unset=True)
+    async def update_item(
+        db: DBDep,
+        item_id: Any,
+        payload: update_schema = Body(...),  # type: ignore[valid-type]
+    ):
+        data = cast(Any, payload).model_dump(exclude_unset=True)
         row = await service.update(db, item_id, data)
         if not row:
             raise HTTPException(404, "Not found")