重构数据库连接管理，添加元数据支持

2026-02-11 18:57:47 +08:00
parent ff2011ae24
commit 780a48d927
21 changed files with 1195 additions and 305 deletions
--- a/app/auth/keycloak_dependencies.py
+++ b/app/auth/keycloak_dependencies.py
@@ -0,0 +1,56 @@
+from uuid import UUID
+
+from fastapi import Depends, HTTPException, status
+from fastapi.security import OAuth2PasswordBearer
+from jose import JWTError, jwt
+
+from app.core.config import settings
+
+oauth2_optional = OAuth2PasswordBearer(
+    tokenUrl=f"{settings.API_V1_STR}/auth/login", auto_error=False
+)
+
+
+async def get_current_keycloak_sub(
+    token: str | None = Depends(oauth2_optional),
+) -> UUID:
+    if settings.AUTH_DISABLED:
+        return UUID(int=0)
+    if not token:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Not authenticated",
+            headers={"WWW-Authenticate": "Bearer"},
+        )
+    if settings.KEYCLOAK_PUBLIC_KEY:
+        key = settings.KEYCLOAK_PUBLIC_KEY.replace("\\n", "\n")
+        algorithms = [settings.KEYCLOAK_ALGORITHM]
+    else:
+        key = settings.SECRET_KEY
+        algorithms = [settings.ALGORITHM]
+
+    try:
+        payload = jwt.decode(token, key, algorithms=algorithms)
+    except JWTError as exc:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Invalid token",
+            headers={"WWW-Authenticate": "Bearer"},
+        ) from exc
+
+    sub = payload.get("sub")
+    if not sub:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Missing subject claim",
+            headers={"WWW-Authenticate": "Bearer"},
+        )
+
+    try:
+        return UUID(sub)
+    except ValueError as exc:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Invalid subject claim",
+            headers={"WWW-Authenticate": "Bearer"},
+        ) from exc
--- a/app/auth/metadata_dependencies.py
+++ b/app/auth/metadata_dependencies.py
@@ -0,0 +1,50 @@
+from dataclasses import dataclass
+from uuid import UUID
+
+from fastapi import Depends, HTTPException, status
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from app.auth.keycloak_dependencies import get_current_keycloak_sub
+from app.core.config import settings
+from app.infra.db.metadata.database import get_metadata_session
+from app.infra.repositories.metadata_repository import MetadataRepository
+
+
+async def get_metadata_repository(
+    session: AsyncSession = Depends(get_metadata_session),
+) -> MetadataRepository:
+    return MetadataRepository(session)
+
+
+async def get_current_metadata_user(
+    keycloak_sub: UUID = Depends(get_current_keycloak_sub),
+    metadata_repo: MetadataRepository = Depends(get_metadata_repository),
+):
+    if settings.AUTH_DISABLED:
+        return _AuthBypassUser()
+    user = await metadata_repo.get_user_by_keycloak_id(keycloak_sub)
+    if not user or not user.is_active:
+        raise HTTPException(
+            status_code=status.HTTP_403_FORBIDDEN, detail="Inactive user"
+        )
+    return user
+
+
+async def get_current_metadata_admin(
+    user=Depends(get_current_metadata_user),
+):
+    if settings.AUTH_DISABLED:
+        return user
+    if user.is_superuser or user.role == "admin":
+        return user
+    raise HTTPException(
+        status_code=status.HTTP_403_FORBIDDEN, detail="Admin access required"
+    )
+
+
+@dataclass(frozen=True)
+class _AuthBypassUser:
+    id: UUID = UUID(int=0)
+    role: str = "admin"
+    is_superuser: bool = True
+    is_active: bool = True
--- a/app/auth/project_dependencies.py
+++ b/app/auth/project_dependencies.py
@@ -0,0 +1,176 @@
+from dataclasses import dataclass
+from typing import AsyncGenerator
+from uuid import UUID
+
+from fastapi import Depends, Header, HTTPException, status
+from psycopg import AsyncConnection
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from app.auth.keycloak_dependencies import get_current_keycloak_sub
+from app.core.config import settings
+from app.infra.db.dynamic_manager import project_connection_manager
+from app.infra.db.metadata.database import get_metadata_session
+from app.infra.repositories.metadata_repository import MetadataRepository
+
+DB_ROLE_BIZ_PG = "biz_pg"
+DB_ROLE_IOT_TS = "iot_ts"
+DB_TYPE_POSTGRES = "postgresql"
+DB_TYPE_TIMESCALE = "timescaledb"
+
+
+@dataclass(frozen=True)
+class ProjectContext:
+    project_id: UUID
+    user_id: UUID
+    project_role: str
+
+
+async def get_metadata_repository(
+    session: AsyncSession = Depends(get_metadata_session),
+) -> MetadataRepository:
+    return MetadataRepository(session)
+
+
+async def get_project_context(
+    x_project_id: str = Header(..., alias="X-Project-Id"),
+    keycloak_sub: UUID = Depends(get_current_keycloak_sub),
+    metadata_repo: MetadataRepository = Depends(get_metadata_repository),
+) -> 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
+
+    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"
+        )
+
+    if settings.AUTH_DISABLED:
+        return ProjectContext(
+            project_id=project.id,
+            user_id=UUID(int=0),
+            project_role="owner",
+        )
+
+    user = await metadata_repo.get_user_by_keycloak_id(keycloak_sub)
+    if not user:
+        raise HTTPException(
+            status_code=status.HTTP_403_FORBIDDEN, detail="User not registered"
+        )
+    if not user.is_active:
+        raise HTTPException(
+            status_code=status.HTTP_403_FORBIDDEN, detail="Inactive user"
+        )
+
+    membership_role = await metadata_repo.get_membership_role(project_uuid, user.id)
+    if not membership_role:
+        raise HTTPException(
+            status_code=status.HTTP_403_FORBIDDEN, detail="No access to project"
+        )
+
+    return ProjectContext(
+        project_id=project.id,
+        user_id=user.id,
+        project_role=membership_role,
+    )
+
+
+async def get_project_pg_session(
+    ctx: ProjectContext = Depends(get_project_context),
+    metadata_repo: MetadataRepository = Depends(get_metadata_repository),
+) -> AsyncGenerator[AsyncSession, None]:
+    routing = await metadata_repo.get_project_db_routing(
+        ctx.project_id, DB_ROLE_BIZ_PG
+    )
+    if not routing:
+        raise HTTPException(
+            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
+            detail="Project PostgreSQL not configured",
+        )
+    if routing.db_type != DB_TYPE_POSTGRES:
+        raise HTTPException(
+            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
+            detail="Project PostgreSQL type mismatch",
+        )
+
+    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_PG,
+        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 metadata_repo.get_project_db_routing(
+        ctx.project_id, DB_ROLE_BIZ_PG
+    )
+    if not routing:
+        raise HTTPException(
+            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
+            detail="Project PostgreSQL not configured",
+        )
+    if routing.db_type != DB_TYPE_POSTGRES:
+        raise HTTPException(
+            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
+            detail="Project PostgreSQL type mismatch",
+        )
+
+    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_PG,
+        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 metadata_repo.get_project_db_routing(
+        ctx.project_id, DB_ROLE_IOT_TS
+    )
+    if not routing:
+        raise HTTPException(
+            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
+            detail="Project TimescaleDB not configured",
+        )
+    if routing.db_type != DB_TYPE_TIMESCALE:
+        raise HTTPException(
+            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
+            detail="Project TimescaleDB type mismatch",
+        )
+
+    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_TS,
+        routing.dsn,
+        pool_min_size,
+        pool_max_size,
+    )
+    async with pool.connection() as conn:
+        yield conn