"""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, to_regclass('app.admin_health_event') IS NOT NULL AS has_admin_health_event, to_regclass('app.support_ticket') IS NOT NULL AS has_support_ticket """ ) return bool( row and row["has_user"] and row["has_auth_session"] and row["has_preferences"] and row["has_engine_config"] and row["has_admin_health_event"] and row["has_support_ticket"] ) 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)