Upload 3 files

app/cache.py

@@ -0,0 +1,133 @@
+"""内存缓存与限流计数，减少对数据库的依赖。
+
+key 缓存：首次校验后把 key 的校验信息与配额缓存在内存，
+后续请求命中缓存即不查库；后台改 key 时清缓存。
+RPM：用内存滑动窗口计数，不再查 request_logs 表。
+"""
+from __future__ import annotations
+
+import asyncio
+import time
+from dataclasses import dataclass
+from typing import Dict, List, Set
+
+import hashlib
+
+
+@dataclass
+class CachedKey:
+    """key 的缓存校验信息（不含敏感的完整 key）。"""
+    id: int
+    name: str
+    prefix: str
+    key_hash: str
+    is_active: bool
+    expires_at: float | None  # epoch 秒
+    allowed_models: list[str] | None
+    rpm_limit: int | None
+    tpm_limit: int | None
+    token_quota: int | None
+    tokens_used: int
+    # 负载快照时间，用于判断是否需要回库刷新 tokens_used
+    loaded_at: float
+
+
+class KeyCache:
+    """按 prefix 索引的 key 内存缓存。"""
+
+    def __init__(self) -> None:
+        self._by_prefix: Dict[str, CachedKey] = {}
+        self._by_id: Dict[int, str] = {}  # id -> prefix，用于按 id 失效
+        self._lock = asyncio.Lock()
+
+    def get(self, prefix: str) -> CachedKey | None:
+        return self._by_prefix.get(prefix)
+
+    async def put(self, k: CachedKey) -> None:
+        async with self._lock:
+            self._by_prefix[k.prefix] = k
+            self._by_id[k.id] = k.prefix
+
+    async def invalidate_id(self, key_id: int) -> None:
+        """按 id 失效（后台改/删 key 时调用）。"""
+        async with self._lock:
+            prefix = self._by_id.pop(key_id, None)
+            if prefix:
+                self._by_prefix.pop(prefix, None)
+
+    async def invalidate_prefix(self, prefix: str) -> None:
+        async with self._lock:
+            ck = self._by_prefix.pop(prefix, None)
+            if ck:
+                self._by_id.pop(ck.id, None)
+
+    async def clear(self) -> None:
+        async with self._lock:
+            self._by_prefix.clear()
+            self._by_id.clear()
+
+    async def add_usage(self, key_id: int, tokens: int) -> None:
+        """成功转发后把用量累加到内存缓存（不回库），保持配额判断近实时。"""
+        async with self._lock:
+            prefix = self._by_id.get(key_id)
+            if prefix:
+                ck = self._by_prefix.get(prefix)
+                if ck:
+                    ck.tokens_used += tokens
+
+
+class RPMCounter:
+    """内存滑动窗口 RPM 计数，按 key_id 维度。"""
+
+    def __init__(self, window: float = 60.0) -> None:
+        self._window = window
+        self._hits: Dict[int, List[float]] = {}
+        self._lock = asyncio.Lock()
+
+    async def hit_and_check(self, key_id: int, limit: int | None) -> bool:
+        """记录一次命中并返回是否在限额内（True=放行，False=超限）。"""
+        if not limit:
+            return True
+        now = time.time()
+        cutoff = now - self._window
+        async with self._lock:
+            hits = self._hits.get(key_id)
+            if hits is None:
+                hits = []
+                self._hits[key_id] = hits
+            # 清掉过期命中
+            hits[:] = [t for t in hits if t > cutoff]
+            if len(hits) >= limit:
+                return False
+            hits.append(now)
+            return True
+
+
+key_cache = KeyCache()
+rpm_counter = RPMCounter()
+
+
+def verify_key_hash(raw: str, digest: str) -> bool:
+    return hashlib.sha256(raw.encode()).hexdigest() == digest
+
+
+async def cached_key_from(raw: str, db_loader) -> CachedKey | None:
+    """先查缓存，未命中则调 db_loader(prefix) 回库加载并缓存。"""
+    prefix = raw[:20]
+    ck = key_cache.get(prefix)
+    if ck is not None and verify_key_hash(raw, ck.key_hash):
+        return ck
+    # 未命中或哈希不匹配，回库
+    db_key = await db_loader(prefix)
+    if db_key is None:
+        return None
+    exp = db_key.expires_at.timestamp() if db_key.expires_at else None
+    ck = CachedKey(
+        id=db_key.id, name=db_key.name, prefix=db_key.prefix, key_hash=db_key.key_hash,
+        is_active=db_key.is_active, expires_at=exp, allowed_models=db_key.allowed_models,
+        rpm_limit=db_key.rpm_limit, tpm_limit=db_key.tpm_limit,
+        token_quota=db_key.token_quota, tokens_used=db_key.tokens_used or 0,
+        loaded_at=time.time(),
+    )
+    await key_cache.put(ck)
+    return ck

app/db.py

@@ -42,7 +42,10 @@ def _normalize_db_url() -> tuple[object, dict]:
 
     # 兼容 Neon/Supabase 的 PgBouncer 连接池：关闭预编译语句缓存，
     # 否则跨事务会报 "prepared statement ... does not exist"
-    connect_args.setdefault("statement_cache_size", 0)
+    # asyncpg 专属：兼容 Neon/Supabase 的 PgBouncer 连接池，否则跨事务会报
+    # "prepared statement ... does not exist"（仅 asyncpg 驱动可设）
+    if url.drivername == "postgresql+asyncpg":
+        connect_args.setdefault("statement_cache_size", 0)
 
     return url, connect_args
 

app/security.py

@@ -12,6 +12,7 @@ from passlib.context import CryptContext
 from sqlalchemy import select
 from sqlalchemy.ext.asyncio import AsyncSession
 
+from . import cache
 from .config import get_settings
 from .db import get_db
 from .models import AdminUser, ApiKey
@@ -69,11 +70,22 @@ async def get_admin(
     return user
 
 
+async def _load_key_from_db(prefix: str, db: AsyncSession) -> ApiKey | None:
+    """从数据库按前缀加载 key。"""
+    return (
+        await db.execute(select(ApiKey).where(ApiKey.prefix == prefix))
+    ).scalar_one_or_none()
+
+
 async def get_valid_api_key(
     credentials: HTTPAuthorizationCredentials | None = Depends(bearer),
     db: AsyncSession = Depends(get_db),
-) -> ApiKey:
-    """代理接口依赖：校验客户端 relay key。"""
+) -> cache.CachedKey:
+    """代理接口依赖：校验客户端 relay key，优先走内存缓存。
+
+    缓存未命中时回库加载，加载失败则抛 401（无效 key）。
+    返回 CachedKey（内存缓存对象），供后续配额/日志使用，避免再查库。
+    """
     if not credentials or not credentials.credentials:
         raise HTTPException(
             status_code=401,
@@ -81,10 +93,39 @@ async def get_valid_api_key(
             headers={"WWW-Authenticate": "Bearer"},
         )
     raw = credentials.credentials.strip()
-    prefix = raw[:20]
-    key = (await db.execute(select(ApiKey).where(ApiKey.prefix == prefix))).scalar_one_or_none()
-    if not key or not key.is_active or not verify_api_key(raw, key.key_hash):
+
+    # 先查内存缓存
+    ck = cache.key_cache.get(raw[:20])
+    if ck is not None and verify_api_key(raw, ck.key_hash):
+        if not ck.is_active:
+            raise HTTPException(status_code=401, detail="无效的 API Key")
+        if ck.expires_at and ck.expires_at < datetime.now(timezone.utc).timestamp():
+            raise HTTPException(status_code=401, detail="API Key 已过期")
+        return ck
+
+    # 缓存未命中或哈希不匹配，回库加载
+    # 缓存未命中或哈希不匹配，回库加载（数据库异常时降级：缓存可用则继续，否则 503）
+    try:
+        db_key = await _load_key_from_db(raw[:20], db)
+    except Exception:
+        # 数据库挂了：若缓存里还有这个 key（哈希匹配但已过期/停用等分支没进），仍可用
+        ck = cache.key_cache.get(raw[:20])
+        if ck is not None and verify_api_key(raw, ck.key_hash):
+            return ck
+        raise HTTPException(status_code=503, detail="数据库暂时不可用，请稍后再试")
+    if not db_key or not db_key.is_active or not verify_api_key(raw, db_key.key_hash):
         raise HTTPException(status_code=401, detail="无效的 API Key")
-    if key.expires_at and key.expires_at < datetime.now(timezone.utc):
+    if db_key.expires_at and db_key.expires_at < datetime.now(timezone.utc):
         raise HTTPException(status_code=401, detail="API Key 已过期")
-    return key
+
+    # 缓存校验信息
+    exp = db_key.expires_at.timestamp() if db_key.expires_at else None
+    ck = cache.CachedKey(
+        id=db_key.id, name=db_key.name, prefix=db_key.prefix, key_hash=db_key.key_hash,
+        is_active=db_key.is_active, expires_at=exp, allowed_models=db_key.allowed_models,
+        rpm_limit=db_key.rpm_limit, tpm_limit=db_key.tpm_limit,
+        token_quota=db_key.token_quota, tokens_used=db_key.tokens_used or 0,
+        loaded_at=datetime.now(timezone.utc).timestamp(),
+    )
+    await cache.key_cache.put(ck)
+    return ck