"""사용자 실행 컨텍스트 레지스트리 (설계 문서 §2, §5.6). `UserContext` 조립(자격증명 복호화 → KIS 클라이언트 풀 → 리스크 매니저 팩토리 → 상태 복원)을 한곳에 모아 캐싱한다. FastAPI 의존성 `get_user_context()`도 여기서 제공한다(트레이딩 실행/엔드포인트 양쪽에서 재사용). - DB 세션은 항상 호출 측이 넘겨준다 — 이 모듈은 `AsyncSessionLocal()`을 직접 열지 않는다(§CLAUDE.md 아키텍처 원칙 1). `iter_active()`만 예외적으로 스케줄러가 세션을 만들어 넘겨주는 것을 전제로 `db` 인자를 필수로 받는다. - 복호화 실패 등 개별 사용자 오류는 그 사용자만 None/스킵 처리하고 다른 사용자에 전파하지 않는다. """ import logging import time import asyncio from fastapi import Depends, HTTPException from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession from app.core.config import settings from app.core.crypto import CredentialDecryptError, decrypt from app.core.database import get_db from app.core.tenant_auth import CurrentUser, get_current_user, get_legacy_owner from app.models.models import UserKisCredential from app.services.tenant.context import UserContext from app.services.tenant.kis_pool import kis_client_pool from app.services.tenant.risk_factory import ( create_loss_tracker_for_user, create_mdd_manager_for_user, load_user_risk_state, ) logger = logging.getLogger(__name__) _LEGACY_OWNER_TTL_SEC = 60 _ACTIVE_LIST_TTL_SEC = 60 class UserContextRegistry: """사용자별 `UserContext` 캐시 + 조립 로직.""" def __init__(self) -> None: self._contexts: dict[int, UserContext] = {} self._lock = asyncio.Lock() self._legacy_ctx: UserContext | None = None self._legacy_owner_id: int | None = None self._legacy_checked_at: float = 0.0 self._legacy_lock = asyncio.Lock() self._legacy_owner_confirmed: bool = False self._active_cache: list[UserContext] | None = None self._active_cached_at: float = 0.0 # ── 사용자별 컨텍스트 조립 ──────────────────────────────────────────── async def get(self, user_id: int, db: AsyncSession) -> UserContext | None: """캐시 히트면 즉시 반환, 미스면 DB에서 자격증명을 읽어 새로 조립한다. M4: `asyncio.Lock` 은 캐시 dict 갱신 구간(더블체크 패턴)만 감싼다 — DB I/O, PBKDF2 복호화(스레드로 위임), KIS 클라이언트 생성, Redis/DB 리스크 상태 복원은 모두 락 밖에서 수행해 이벤트 루프 직렬화를 피한다. 동시에 같은 user_id 로 호출되면 중복 조립이 일어날 수 있으나(락 밖 조립) 마지막에 dict 갱신만 락으로 보호하고 먼저 캐시된 인스턴스를 우선한다. """ cached = self._contexts.get(user_id) if cached is not None: return cached result = await db.execute( select(UserKisCredential).where(UserKisCredential.user_id == user_id) ) cred = result.scalar_one_or_none() if cred is None or not cred.is_enabled: logger.debug(f"[UserContextRegistry] user={user_id} 자격증명 없음/비활성 — None 반환") return None try: app_key = await asyncio.to_thread(decrypt, cred.app_key_enc, key_version=cred.key_version) app_secret = await asyncio.to_thread(decrypt, cred.app_secret_enc, key_version=cred.key_version) account_no = await asyncio.to_thread(decrypt, cred.account_no_enc, key_version=cred.key_version) except CredentialDecryptError as e: logger.error(f"[UserContextRegistry] user={user_id} 자격증명 복호화 실패: {e}") raise HTTPException( status_code=400, detail="자격증명 복호화 실패 — 서버 키 변경 여부 확인", ) from e try: kis = await kis_client_pool.get_or_create( user_id=user_id, app_key=app_key, app_secret=app_secret, account_no=account_no, mode=cred.mode, ) except Exception as e: logger.error(f"[UserContextRegistry] user={user_id} KIS 클라이언트 생성 실패: {e}") return None mdd = create_mdd_manager_for_user(user_id) loss_tracker = create_loss_tracker_for_user(user_id) await load_user_risk_state(user_id, mdd, loss_tracker, db) ctx = UserContext( user_id=user_id, username="", mode=cred.mode, kis=kis, mdd=mdd, loss_tracker=loss_tracker, auto_trade_enabled=bool(cred.auto_trade_enabled), is_legacy=False, ) async with self._lock: existing = self._contexts.get(user_id) if existing is not None: return existing self._contexts[user_id] = ctx return ctx # ── legacy(단일 계좌) 컨텍스트 ──────────────────────────────────────── async def legacy(self, db: AsyncSession | None = None) -> UserContext: """기존 전역 싱글톤(`korea_investment_api`/`mdd_manager`/`loss_tracker`)을 감싼 컨텍스트. owner id는 60초 TTL로 재확인한다(관리자 계정이 나중에 생기는 초기 부팅 구간 대응). owner 가 아직 없으면(관리자 미존재) `user_id=None`으로 반환하고 TTL 캐시를 걸지 않는다(M1) — `users.id=0`은 존재하지 않는 FK 위반 값이라 0 대신 None 으로 통일한다. """ now = time.monotonic() if ( self._legacy_ctx is not None and self._legacy_owner_confirmed and (now - self._legacy_checked_at) < _LEGACY_OWNER_TTL_SEC ): return self._legacy_ctx async with self._legacy_lock: now = time.monotonic() if ( self._legacy_ctx is not None and self._legacy_owner_confirmed and (now - self._legacy_checked_at) < _LEGACY_OWNER_TTL_SEC ): return self._legacy_ctx owner_id: int | None = None if db is not None: try: owner = await get_legacy_owner(db) if owner is not None: owner_id = owner.id except Exception as e: logger.warning(f"[UserContextRegistry] legacy owner 조회 실패: {e}") from app.services.korea_investment_api import korea_investment_api from app.services.mdd_manager import mdd_manager from app.services.loss_tracker import loss_tracker self._legacy_ctx = UserContext( user_id=owner_id, username="legacy", mode=settings.KIS_MODE, kis=korea_investment_api, mdd=mdd_manager, loss_tracker=loss_tracker, auto_trade_enabled=settings.AUTO_TRADE_ENABLED, is_legacy=True, ) self._legacy_owner_id = owner_id self._legacy_owner_confirmed = owner_id is not None self._legacy_checked_at = now return self._legacy_ctx # ── 인증 사용자 → 컨텍스트 해석 ─────────────────────────────────────── async def resolve(self, current_user: CurrentUser, db: AsyncSession) -> UserContext: """`CurrentUser` → `UserContext`. 실패 시 400 HTTPException.""" legacy_ctx = await self.legacy(db) use_legacy = ( not settings.MULTI_TENANT_ENABLED or current_user.is_api_key or (current_user.id is not None and current_user.id == legacy_ctx.user_id) ) if use_legacy: return legacy_ctx if current_user.id is None: raise HTTPException(status_code=400, detail="KIS 계좌가 등록·검증되지 않았습니다. 내 계좌 설정에서 등록하세요") ctx = await self.get(current_user.id, db) if ctx is None: raise HTTPException(status_code=400, detail="KIS 계좌가 등록·검증되지 않았습니다. 내 계좌 설정에서 등록하세요") return ctx # ── 2단계 스케줄러용: 활성 사용자 전체 순회 ─────────────────────────── async def iter_active(self, db: AsyncSession) -> list[UserContext]: """`is_enabled and auto_trade_enabled` 인 모든 사용자의 컨텍스트 (60초 TTL 캐시). 지금은 조립 로직만 제공 — 실제 2단계 스케줄러 연결은 이번 범위 밖. """ now = time.monotonic() if self._active_cache is not None and (now - self._active_cached_at) < _ACTIVE_LIST_TTL_SEC: return self._active_cache result = await db.execute( select(UserKisCredential).where( UserKisCredential.is_enabled.is_(True), UserKisCredential.auto_trade_enabled.is_(True), ) ) creds = result.scalars().all() contexts: list[UserContext] = [] for cred in creds: try: ctx = await self.get(cred.user_id, db) except Exception as e: logger.error(f"[UserContextRegistry] iter_active: user={cred.user_id} 조립 실패: {e}") ctx = None if ctx is not None: contexts.append(ctx) self._active_cache = contexts self._active_cached_at = now return contexts # ── 캐시 무효화 ─────────────────────────────────────────────────────── async def invalidate(self, user_id: int) -> None: """자격증명 변경/삭제 시 credentials API가 호출 — 컨텍스트 캐시 + KIS 풀 정리.""" async with self._lock: self._contexts.pop(user_id, None) await kis_client_pool.invalidate(user_id) self._active_cache = None async def invalidate_all(self) -> None: async with self._lock: self._contexts.clear() self._active_cache = None # 모듈 싱글톤 user_context_registry = UserContextRegistry() async def get_user_context( user: CurrentUser = Depends(get_current_user), db: AsyncSession = Depends(get_db), ) -> UserContext: """FastAPI 의존성 — 세션/API 키 인증 사용자를 `UserContext`로 해석.""" return await user_context_registry.resolve(user, db)