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,
)
from app.infra.db.project_routing import ActiveProjectRouting, activate_project_routing

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_business_routing(
    ctx: ProjectContext = Depends(get_project_context),
    metadata_repo: MetadataRepository = Depends(get_metadata_repository),
) -> ActiveProjectRouting:
    return await resolve_project_business_routing(ctx, metadata_repo)


async def use_project_business_routing(
    routing: ActiveProjectRouting = Depends(get_project_business_routing),
) -> AsyncGenerator[ActiveProjectRouting, None]:
    """Keep the routed BizDB active for an entire endpoint invocation."""
    with activate_project_routing(routing):
        yield routing


async def resolve_project_business_routing(
    ctx: ProjectContext,
    metadata_repo: MetadataRepository,
) -> ActiveProjectRouting:
    business = await _get_project_routing(
        metadata_repo,
        ctx.project_id,
        DB_ROLE_BIZ_DATA,
        DB_TYPE_POSTGRES,
        "PostgreSQL",
    )
    return ActiveProjectRouting(
        project_code=ctx.project_code,
        business_dsn=business.dsn,
    )


async def get_project_simulation_routing(
    ctx: ProjectContext = Depends(get_project_context),
    metadata_repo: MetadataRepository = Depends(get_metadata_repository),
) -> ActiveProjectRouting:
    return await resolve_project_simulation_routing(ctx, metadata_repo)


async def resolve_project_simulation_routing(
    ctx: ProjectContext,
    metadata_repo: MetadataRepository,
) -> ActiveProjectRouting:
    business = await _get_project_routing(
        metadata_repo,
        ctx.project_id,
        DB_ROLE_BIZ_DATA,
        DB_TYPE_POSTGRES,
        "PostgreSQL",
    )
    timescale = await _get_project_routing(
        metadata_repo,
        ctx.project_id,
        DB_ROLE_IOT_DATA,
        DB_TYPE_TIMESCALE,
        "TimescaleDB",
    )
    return ActiveProjectRouting(
        project_code=ctx.project_code,
        business_dsn=business.dsn,
        timescale_dsn=timescale.dsn,
    )


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_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
        if routing.pool_min_size is not None
        else settings.PROJECT_PG_POOL_MIN_SIZE
    )
    pool_max_size = (
        routing.pool_max_size
        if routing.pool_max_size is not None
        else settings.PROJECT_PG_POOL_SIZE
    )
    async with project_connection_manager.pg_connection(
        ctx.project_id,
        DB_ROLE_BIZ_DATA,
        routing.dsn,
        pool_min_size,
        pool_max_size,
    ) 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
        if routing.pool_min_size is not None
        else settings.PROJECT_TS_POOL_MIN_SIZE
    )
    pool_max_size = (
        routing.pool_max_size
        if routing.pool_max_size is not None
        else settings.PROJECT_TS_POOL_MAX_SIZE
    )
    async with project_connection_manager.timescale_connection(
        ctx.project_id,
        DB_ROLE_IOT_DATA,
        routing.dsn,
        pool_min_size,
        pool_max_size,
    ) as conn:
        yield conn