vignette/apps/api/app/db.py

"""asyncpg 연결 풀 + pgvector 등록.

DB = NAS PostgreSQL 16 단일 SoR (마스터플랜 §0). 스키마 4분할: app / kb / audit / ds.
RLS 이중강제: 커넥션 획득 시 SET LOCAL app.current_role / app.current_uid /
app.current_cohort / app.ai_context / app.current_ai_view / app.current_sens_max 주입
(deps.py 의 RBAC 의존성과 짝). 여기선 풀 + 헬퍼만 제공한다.
"""

from __future__ import annotations

import json
from contextlib import asynccontextmanager
from typing import Any, AsyncIterator, Optional, Sequence

import asyncpg

from .config import settings

# 전역 풀 핸들. main.py lifespan 에서 init/close.
_pool: Optional[asyncpg.Pool] = None


async def _init_connection(conn: asyncpg.Connection) -> None:
    """커넥션 단위 코덱 등록.

    - jsonb: dict 직렬화 자동 (asyncpg 기본은 str 반환)
    - vector(1024): pgvector. 런타임 인코딩은 RAG 경로에서 처리(여기선 텍스트 캐스트 허용).
    TODO: pgvector 바이너리 코덱 등록(register_vector) — RAG 라우터 구현 시 BGE-M3 1024d 연동.
    """
    await conn.set_type_codec(
        "jsonb",
        encoder=lambda v: json.dumps(v, ensure_ascii=False),
        decoder=json.loads,
        schema="pg_catalog",
    )
    await conn.set_type_codec(
        "json",
        encoder=lambda v: json.dumps(v, ensure_ascii=False),
        decoder=json.loads,
        schema="pg_catalog",
    )


async def init_pool() -> asyncpg.Pool:
    """풀 생성 (main lifespan startup)."""
    global _pool
    if _pool is not None:
        return _pool
    _pool = await asyncpg.create_pool(
        dsn=settings.database_url,
        min_size=settings.db_pool_min_size,
        max_size=settings.db_pool_max_size,
        command_timeout=settings.db_command_timeout,
        init=_init_connection,
    )
    return _pool


async def close_pool() -> None:
    """풀 종료 (main lifespan shutdown)."""
    global _pool
    if _pool is not None:
        await _pool.close()
        _pool = None


def get_pool() -> asyncpg.Pool:
    """초기화된 풀 반환. lifespan 밖 호출 시 RuntimeError."""
    if _pool is None:
        raise RuntimeError("DB pool not initialized — init_pool() must run in lifespan startup")
    return _pool


def _db_role(role: str) -> str:
    return "instructor" if role == "teacher" else role


def _default_sensitivity_max(ai_view: str | None) -> int | None:
    return {
        "client": 1,
        "counselor": 0,
        "evaluator": 2,
    }.get(ai_view or "")


def _cohort_value(cohort_ids: Sequence[str] | None, cohort: str | None) -> str:
    if cohort:
        return cohort
    if not cohort_ids:
        return ""
    return cohort_ids[0] or ""


@asynccontextmanager
async def acquire(
    *,
    role: Optional[str] = None,
    user_id: Optional[str] = None,
    cohort_ids: Optional[Sequence[str]] = None,
    cohort: Optional[str] = None,
    ai_view: Optional[str] = None,
    ai_context: Optional[bool] = None,
    sensitivity_max: Optional[int] = None,
) -> AsyncIterator[asyncpg.Connection]:
    """커넥션 획득 + RLS 컨텍스트 주입.

    RLS 이중강제 (설계서 §4.1, F-30):
      레이어1 = AI 정보비대칭: app.ai_context + app.current_ai_view +
                app.current_sens_max (visible_to[]/sensitivity WHERE 강제)
      레이어2 = 인간 RBAC×cohort: app.current_role + app.current_uid +
                app.current_cohort (RLS 정책)
    트랜잭션 내 SET LOCAL 로 주입해 커넥션 풀 재사용 시 누수 방지.

    NOTE: RLS 정책/세션변수는 Phase 0 마이그레이션에서 정의(설계서 §3.3 / §4).
          여기선 변수 주입 계약만 확정.
    """
    pool = get_pool()
    async with pool.acquire() as conn:
        async with conn.transaction():
            is_ai = ai_context if ai_context is not None else ai_view is not None
            await conn.execute("SELECT set_config('app.ai_context', $1, true)", "1" if is_ai else "")
            if role is not None:
                await conn.execute("SELECT set_config('app.current_role', $1, true)", _db_role(role))
            if user_id is not None:
                await conn.execute("SELECT set_config('app.current_uid', $1, true)", user_id)
            cohort_name = _cohort_value(cohort_ids, cohort)
            if cohort_name:
                await conn.execute("SELECT set_config('app.current_cohort', $1, true)", cohort_name)
            if ai_view is not None:
                await conn.execute("SELECT set_config('app.current_ai_view', $1, true)", ai_view)
                sens = sensitivity_max if sensitivity_max is not None else _default_sensitivity_max(ai_view)
                if sens is not None:
                    await conn.execute(
                        "SELECT set_config('app.current_sens_max', $1, true)",
                        str(sens),
                    )
            yield conn


async def healthcheck() -> bool:
    """Return true only when the DB is reachable and required app tables exist."""
    try:
        pool = get_pool()
        async with pool.acquire() as conn:
            row = await conn.fetchrow(
                """
                SELECT
                  to_regclass('app.app_user') IS NOT NULL AS has_user,
                  to_regclass('app.auth_session') IS NOT NULL AS has_auth_session,
                  to_regclass('app.user_preferences') IS NOT NULL AS has_preferences,
                  to_regclass('app.admin_engine_config') IS NOT NULL AS has_engine_config
                """
            )
            return bool(
                row
                and row["has_user"]
                and row["has_auth_session"]
                and row["has_preferences"]
                and row["has_engine_config"]
            )
    except Exception:
        return False


async def fetch(query: str, *args: Any, role: Optional[str] = None, ai_view: Optional[str] = None):
    async with acquire(role=role, ai_view=ai_view) as conn:
        return await conn.fetch(query, *args)


async def fetchrow(
    query: str, *args: Any, role: Optional[str] = None, ai_view: Optional[str] = None
):
    async with acquire(role=role, ai_view=ai_view) as conn:
        return await conn.fetchrow(query, *args)


async def execute(
    query: str, *args: Any, role: Optional[str] = None, ai_view: Optional[str] = None
) -> str:
    async with acquire(role=role, ai_view=ai_view) as conn:
        return await conn.execute(query, *args)