TJWaterServerBinary/app/auth/project_dependencies.py

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