"""Author RAG Chatbot SaaS — FastAPI Dependency Injection. All shared dependencies (DB session, Redis, current user, subscription) are defined here and injected via FastAPI's Depends() system. """ from typing import AsyncGenerator import redis.asyncio as aioredis import structlog from fastapi import Depends, Header, HTTPException from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine from app.config import Settings, get_settings from app.exceptions.auth import InvalidTokenError, ExpiredTokenError logger = structlog.get_logger(__name__) # ─── Database ───────────────────────────────────────────────────────────────── def _create_engine(cfg: Settings): """Create SQLAlchemy async engine from settings.""" db_url = str(cfg.DATABASE_URL) is_sqlite = db_url.startswith("sqlite") engine_kwargs = { "echo": cfg.DEBUG, "pool_pre_ping": True, } if not is_sqlite: engine_kwargs["pool_size"] = cfg.DB_POOL_SIZE engine_kwargs["max_overflow"] = cfg.DB_MAX_OVERFLOW return create_async_engine(db_url, **engine_kwargs) def _create_session_factory(cfg: Settings) -> async_sessionmaker: """Create async session factory.""" engine = _create_engine(cfg) return async_sessionmaker(engine, expire_on_commit=False, class_=AsyncSession) _session_factory: async_sessionmaker | None = None def _get_session_factory() -> async_sessionmaker: global _session_factory if _session_factory is None: _session_factory = _create_session_factory(get_settings()) return _session_factory async def get_db() -> AsyncGenerator[AsyncSession, None]: """Yield a database session, ensuring proper cleanup.""" async with _get_session_factory()() as session: try: yield session await session.commit() except Exception: try: await session.rollback() except Exception: pass # Don't mask the original exception raise def get_engine(): """Return the underlying async engine (for DDL operations like create_all).""" return _create_engine(get_settings()) from contextlib import asynccontextmanager as _acm @_acm async def get_async_session(): """Async context manager that yields a standalone DB session (for startup tasks).""" async with _get_session_factory()() as session: try: yield session except Exception: try: await session.rollback() except Exception: pass raise # ─── Redis ──────────────────────────────────────────────────────────────────── _redis_pool: aioredis.Redis | None = None async def get_redis() -> aioredis.Redis: """Return the shared Redis connection pool.""" global _redis_pool if _redis_pool is None: cfg = get_settings() _redis_pool = aioredis.from_url( str(cfg.REDIS_URL), decode_responses=cfg.REDIS_DECODE_RESPONSES, encoding="utf-8", ) return _redis_pool # ─── Authentication ─────────────────────────────────────────────────────────── async def _get_authenticated_user( authorization: str, db: AsyncSession, ): """Validate JWT and return the authenticated user model.""" from app.core.access.token_crypto import decode_jwt if not authorization or not authorization.startswith("Bearer "): raise HTTPException(status_code=401, detail="Missing Bearer token") token = authorization.removeprefix("Bearer ").strip() if not token or token in ("null", "undefined", ""): raise HTTPException(status_code=401, detail="Missing Bearer token") try: payload = decode_jwt(token) except ExpiredTokenError: raise HTTPException(status_code=401, detail="Token has expired") except Exception: raise HTTPException(status_code=401, detail="Could not validate token") # R-004/R-005: Check if JWT has been blacklisted (logout/refresh rotation) jti = payload.get("jti") if jti: try: from app.core.access.jwt_blacklist import is_blacklisted redis = await get_redis() if await is_blacklisted(redis, jti): raise HTTPException(status_code=401, detail="Token has been revoked") except HTTPException: raise except Exception: pass # Redis down: fail-open for blacklist check user_id: str = payload.get("sub") if not user_id: raise HTTPException(status_code=401, detail="Token missing subject claim") try: from app.repositories.user_repo import UserRepository user = await UserRepository(db).get_by_id(user_id) except Exception: raise HTTPException(status_code=401, detail="Could not validate user") if user is None or not user.is_active: raise HTTPException(status_code=401, detail="User not found or inactive") return user async def get_current_user( authorization: str = Header(default=""), db: AsyncSession = Depends(get_db), ): """Validate JWT and return any authenticated user (author or superadmin).""" return await _get_authenticated_user(authorization, db) async def get_current_author( authorization: str = Header(default=""), db: AsyncSession = Depends(get_db), ): """Ensure the authenticated user is an Author (not SuperAdmin).""" user = await _get_authenticated_user(authorization, db) if user.role != "author": raise HTTPException( status_code=403, detail="Author access required. SuperAdmins must use /superadmin.", ) return user async def get_current_author_scoped( author_slug: str, current_user=Depends(get_current_author), ): """Ensure the author can only access their own tenant data.""" if author_slug != current_user.id: raise HTTPException(status_code=403, detail="Access denied to this author account") return current_user async def get_superadmin_jwt_only( current_user=Depends(get_current_user), ): """SuperAdmin role check only — for TOTP enrollment before 2FA is configured.""" if current_user.role != "superadmin": raise HTTPException(status_code=403, detail="SuperAdmin access required") return current_user async def get_current_superadmin( current_user=Depends(get_current_user), x_totp_code: str = Header(default="", alias="X-TOTP-Code"), x_totp_session: str = Header(default="", alias="X-TOTP-Session"), redis: aioredis.Redis = Depends(get_redis), ): """Ensure SuperAdmin JWT + TOTP session or one-time code on privileged routes.""" if current_user.role != "superadmin": raise HTTPException(status_code=403, detail="SuperAdmin access required") if not current_user.totp_secret: raise HTTPException( status_code=403, detail="TOTP setup required — complete two-factor enrollment first", ) from app.services.totp_session_service import TotpSessionService if x_totp_session and await TotpSessionService(redis).is_valid(current_user.id, x_totp_session): return current_user from app.core.access.totp import verify_totp if not x_totp_code: raise HTTPException(status_code=403, detail="TOTP verification required") if not verify_totp(current_user.totp_secret, x_totp_code): raise HTTPException(status_code=403, detail="Invalid TOTP code") return current_user async def get_subscription_author( x_subscription_token: str = Header(..., alias="X-Subscription-Token"), db: AsyncSession = Depends(get_db), redis: aioredis.Redis = Depends(get_redis), ): """Validate subscription token for chatbot endpoints.""" from app.core.access.subscription import validate_subscription_token from app.exceptions.access import ( AccessRevokedError, BudgetExhaustedError, InvalidSubscriptionTokenError, SubscriptionExpiredError, SubscriptionNotFoundError, ) try: return await validate_subscription_token(x_subscription_token, db, redis) except InvalidSubscriptionTokenError: raise HTTPException(status_code=401, detail="Invalid or expired subscription token") except SubscriptionExpiredError: raise HTTPException(status_code=403, detail="This chatbot service is currently unavailable") except AccessRevokedError: raise HTTPException(status_code=403, detail="This chatbot service is currently unavailable") except SubscriptionNotFoundError: raise HTTPException(status_code=403, detail="No active subscription. Contact support.") except BudgetExhaustedError: raise # R-016: Let global exception handler in main.py return HTTP 200 with friendly message # ─── Health Check ───────────────────────────────────────────────────── async def check_health() -> dict: """Run all dependency health checks. Returns status dict. B9 fix: ChromaDB heartbeat was a sync blocking call inside an async route, freezing the uvicorn event loop. Wrapped in asyncio.to_thread(). """ import asyncio from sqlalchemy import text health = {"status": "ok", "checks": {}} # DB check try: async with _get_session_factory()() as session: await session.execute(text("SELECT 1")) health["checks"]["database"] = "ok" except Exception as e: health["checks"]["database"] = f"error: {str(e)[:80]}" health["status"] = "degraded" # Redis check try: redis = await get_redis() await redis.ping() health["checks"]["redis"] = "ok" except Exception as e: health["checks"]["redis"] = f"error: {str(e)[:80]}" health["status"] = "degraded" # ChromaDB check # B9 fix: chroma.heartbeat() and list_collections() are SYNCHRONOUS. # Calling them directly in async context blocks the event loop. # asyncio.to_thread() offloads them to a thread pool — safe for HF Spaces. try: from app.core.chroma_client import get_chroma_client chroma = get_chroma_client() await asyncio.to_thread(chroma.heartbeat) collections = await asyncio.to_thread(chroma.list_collections) health["checks"]["chromadb"] = "ok" health["checks"]["chromadb_collections"] = len(collections) except Exception as e: health["checks"]["chromadb"] = f"degraded: {str(e)[:80]}" health["status"] = "degraded" return health