import logging from collections.abc import AsyncGenerator from dataclasses import dataclass from uuid import UUID from fastapi import Depends, Header, HTTPException, status from psycopg import AsyncConnection from sqlalchemy.exc import SQLAlchemyError from sqlalchemy.ext.asyncio import AsyncSession from app.auth.metadata_dependencies import get_current_metadata_user from app.core.config import settings from app.infra.db.dynamic_manager import project_connection_manager from app.infra.db.metadb.database import get_metadata_session from app.infra.db.metadb.repositories.metadata_repository import ( MetadataRepository, ProjectDbRouting, ) DB_ROLE_BIZ_DATA = "biz_data" DB_ROLE_IOT_DATA = "iot_data" DB_TYPE_POSTGRES = "postgresql" DB_TYPE_TIMESCALE = "timescaledb" logger = logging.getLogger(__name__) @dataclass(frozen=True) class ProjectContext: project_id: UUID project_code: str user_id: UUID project_role: str system_role: str = "user" is_superuser: bool = False async def get_metadata_repository( session: AsyncSession = Depends(get_metadata_session), ) -> MetadataRepository: return MetadataRepository(session) async def resolve_project_context( x_project_id: str, current_user, metadata_repo: MetadataRepository, ) -> ProjectContext: try: project_uuid = UUID(x_project_id) except ValueError as exc: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail="Invalid project id" ) from exc try: project = await metadata_repo.get_project_by_id(project_uuid) if not project: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, detail="Project not found" ) if project.status != "active": raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, detail="Project is not active" ) membership_role = await metadata_repo.get_membership_role( project_uuid, current_user.id ) if not membership_role: raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, detail="No access to project" ) except SQLAlchemyError as exc: logger.error( "Metadata DB error while resolving project context", exc_info=True, ) raise HTTPException( status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail="Metadata database is unavailable", ) from exc return ProjectContext( project_id=project.id, project_code=project.code, user_id=current_user.id, project_role=membership_role, system_role=current_user.role, is_superuser=current_user.is_superuser, ) async def get_project_context( x_project_id: str = Header(..., alias="X-Project-Id"), current_user=Depends(get_current_metadata_user), metadata_repo: MetadataRepository = Depends(get_metadata_repository), ) -> ProjectContext: return await resolve_project_context(x_project_id, current_user, metadata_repo) async def _get_project_routing( metadata_repo: MetadataRepository, project_id: UUID, db_role: str, expected_db_type: str, database_label: str, ) -> ProjectDbRouting: try: routing = await metadata_repo.get_project_db_routing(project_id, db_role) except ValueError as exc: logger.error( "Invalid project %s routing DSN configuration", database_label, exc_info=True, ) raise HTTPException( status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail=f"Project {database_label} routing DSN is invalid: {exc}", ) from exc if not routing: raise HTTPException( status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail=f"Project {database_label} not configured", ) if routing.db_type != expected_db_type: raise HTTPException( status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail=f"Project {database_label} type mismatch", ) return routing async def get_project_pg_session( ctx: ProjectContext = Depends(get_project_context), metadata_repo: MetadataRepository = Depends(get_metadata_repository), ) -> AsyncGenerator[AsyncSession, None]: routing = await _get_project_routing( metadata_repo, ctx.project_id, DB_ROLE_BIZ_DATA, DB_TYPE_POSTGRES, "PostgreSQL", ) pool_min_size = routing.pool_min_size or settings.PROJECT_PG_POOL_SIZE pool_max_size = routing.pool_max_size or settings.PROJECT_PG_POOL_SIZE sessionmaker = await project_connection_manager.get_pg_sessionmaker( ctx.project_id, DB_ROLE_BIZ_DATA, routing.dsn, pool_min_size, pool_max_size, ) async with sessionmaker() as session: yield session async def get_project_pg_connection( ctx: ProjectContext = Depends(get_project_context), metadata_repo: MetadataRepository = Depends(get_metadata_repository), ) -> AsyncGenerator[AsyncConnection, None]: routing = await _get_project_routing( metadata_repo, ctx.project_id, DB_ROLE_BIZ_DATA, DB_TYPE_POSTGRES, "PostgreSQL", ) pool_min_size = routing.pool_min_size or settings.PROJECT_PG_POOL_SIZE pool_max_size = routing.pool_max_size or settings.PROJECT_PG_POOL_SIZE pool = await project_connection_manager.get_pg_pool( ctx.project_id, DB_ROLE_BIZ_DATA, routing.dsn, pool_min_size, pool_max_size, ) async with pool.connection() as conn: yield conn async def get_project_timescale_connection( ctx: ProjectContext = Depends(get_project_context), metadata_repo: MetadataRepository = Depends(get_metadata_repository), ) -> AsyncGenerator[AsyncConnection, None]: routing = await _get_project_routing( metadata_repo, ctx.project_id, DB_ROLE_IOT_DATA, DB_TYPE_TIMESCALE, "TimescaleDB", ) pool_min_size = routing.pool_min_size or settings.PROJECT_TS_POOL_MIN_SIZE pool_max_size = routing.pool_max_size or settings.PROJECT_TS_POOL_MAX_SIZE pool = await project_connection_manager.get_timescale_pool( ctx.project_id, DB_ROLE_IOT_DATA, routing.dsn, pool_min_size, pool_max_size, ) async with pool.connection() as conn: yield conn