File size: 29,377 Bytes
e82bac2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513 514 515 516 517 518 519 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547 548 549 550 551 552 | # -*- coding: utf-8 -*-
"""
数据库工具函数。
包含数据库初始化、连接管理、以及针对数据库模型的 CRUD (创建、读取、更新、删除) 操作。
"""
import logging # 导入日志模块
import aiosqlite # 导入异步 SQLite 驱动
from typing import Dict, Any, List, Tuple, Optional, AsyncGenerator # 导入类型提示
from datetime import datetime, timezone # 导入日期时间处理
from sqlalchemy.orm import sessionmaker, Session # 导入 SQLAlchemy 同步会话 (可能用于某些旧代码或特定场景)
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession # 导入 SQLAlchemy 异步引擎和会话
from app.core.database.models import Base, UserKeyAssociation, KeyScore, Setting, ApiKey # 导入数据库模型
from sqlalchemy import select, update, delete # 导入 SQLAlchemy 查询、更新、删除构造器
from sqlalchemy.dialects import sqlite # 导入 SQLite 方言 (可能用于特定查询)
from sqlalchemy.sql import text # 导入用于执行原生 SQL 的 text 函数
from app.core.tracking import key_scores_cache, cache_lock # 导入 Key 分数缓存和锁 (路径不变)
from contextlib import asynccontextmanager # 导入异步上下文管理器
import sqlalchemy # 导入 SQLAlchemy 核心库
import sqlalchemy.exc # 导入 SQLAlchemy 异常
# 配置日志记录
logger = logging.getLogger("my_logger") # 获取日志记录器实例
# --- 数据库路径和 URL 配置 ---
# 使用 os 模块动态确定数据库文件的绝对路径,增强可移植性
import os
# _current_dir = os.path.dirname(os.path.abspath(__file__)) # 获取当前文件所在目录的绝对路径
# _project_root = os.path.abspath(os.path.join(_current_dir, '..', '..')) # 获取项目根目录 (假设 utils.py 在 app/core/database/ 下)
# DATABASE_PATH = os.path.join(_project_root, 'app', 'data', 'context_store.db') # 拼接数据库文件的绝对路径
# --- 数据库路径和 URL 配置 (Hugging Face Spaces 部署优化) ---
# 根据 .clinerules/deployment-preferences.md,默认使用内存数据库
# 检查环境变量 HF_SPACE_ID 来判断是否在 Hugging Face 环境中
# 或者引入自定义环境变量 APP_DB_MODE 来控制数据库模式
# 同时考虑 app.config.KEY_STORAGE_MODE 的设置
from app import config as app_config # 导入应用配置
# --- 调试日志:打印相关环境变量和配置值 ---
logger.debug(f"环境变量 HF_SPACE_ID: {os.getenv('HF_SPACE_ID')}")
logger.debug(f"环境变量 APP_DB_MODE: {os.getenv('APP_DB_MODE')}")
logger.debug(f"环境变量 KEY_STORAGE_MODE: {os.getenv('KEY_STORAGE_MODE')}")
logger.debug(f"从 app.config 导入的 KEY_STORAGE_MODE: {app_config.KEY_STORAGE_MODE}")
# --- 结束调试日志 ---
IS_HF_ENV = bool(os.getenv("HF_SPACE_ID")) # 简单判断是否在 Hugging Face 环境
use_memory_db_reason = []
if IS_HF_ENV:
use_memory_db_reason.append("Hugging Face Environment (HF_SPACE_ID set)")
if os.getenv("APP_DB_MODE") == "memory":
use_memory_db_reason.append("APP_DB_MODE=memory")
# 检查从 config 模块解析后的 KEY_STORAGE_MODE
if app_config.KEY_STORAGE_MODE == "memory":
use_memory_db_reason.append(f"config.KEY_STORAGE_MODE is '{app_config.KEY_STORAGE_MODE}' (可能是默认值或环境变量设置)")
if use_memory_db_reason:
DATABASE_PATH = ":memory:" # 内存数据库
logger.info(f"最终决定使用内存数据库。原因: {'; '.join(use_memory_db_reason)}")
else:
# 文件数据库逻辑
logger.info(f"尝试使用文件数据库。内存条件评估: HF Env: {IS_HF_ENV}, APP_DB_MODE env: '{os.getenv('APP_DB_MODE')}', config.KEY_STORAGE_MODE: '{app_config.KEY_STORAGE_MODE}'")
custom_db_path_from_env = app_config.CONTEXT_DB_PATH # 从配置中获取用户指定的路径
if custom_db_path_from_env:
logger.info(f"检测到环境变量 CONTEXT_DB_PATH 设置为: '{custom_db_path_from_env}'")
# 如果用户指定了路径,尝试使用它
# 需要确保路径的目录存在
db_dir = os.path.dirname(custom_db_path_from_env)
if db_dir: # 如果路径包含目录部分
try:
os.makedirs(db_dir, exist_ok=True)
DATABASE_PATH = custom_db_path_from_env
logger.info(f"使用 CONTEXT_DB_PATH 指定的数据库路径: {DATABASE_PATH}")
except PermissionError as e:
logger.error(f"根据 CONTEXT_DB_PATH 创建目录 '{db_dir}' 时发生权限错误: {e}。将回退到默认路径或内存数据库。")
DATABASE_PATH = None # 标记为 None,以便后续逻辑回退
except Exception as e:
logger.error(f"根据 CONTEXT_DB_PATH 创建目录 '{db_dir}' 时发生未知错误: {e}。将回退到默认路径或内存数据库。")
DATABASE_PATH = None # 标记为 None
else: # 如果路径只是文件名,则在当前工作目录创建 (或需要进一步定义行为)
DATABASE_PATH = custom_db_path_from_env
logger.info(f"使用 CONTEXT_DB_PATH 指定的数据库文件名 (将在当前目录创建): {DATABASE_PATH}")
# 注意:直接在当前目录创建可能不是最佳实践,但遵循了优先使用环境变量的原则。
# 更好的做法是,如果只提供文件名,则在默认应用数据目录中使用该文件名。
else: # 如果 CONTEXT_DB_PATH 未设置,则使用默认逻辑
DATABASE_PATH = None # 确保在进入默认逻辑前为 None
if DATABASE_PATH is None: # 如果 CONTEXT_DB_PATH 尝试失败或未设置
logger.info("CONTEXT_DB_PATH 未设置或使用失败,尝试默认应用数据目录。")
_home_dir = os.path.expanduser("~")
_app_data_dir = os.path.join(_home_dir, '.gemini_api_proxy', 'data')
try:
os.makedirs(_app_data_dir, exist_ok=True)
DATABASE_PATH = os.path.join(_app_data_dir, 'context_store.db')
logger.info(f"默认数据库路径设置为: {DATABASE_PATH}")
except PermissionError as e:
logger.error(f"创建默认应用数据目录 '{_app_data_dir}' 时发生权限错误: {e}。将回退到内存数据库。")
DATABASE_PATH = ":memory:"
except Exception as e:
logger.error(f"创建默认应用数据目录时发生未知错误: {e}。将回退到内存数据库。")
DATABASE_PATH = ":memory:"
if DATABASE_PATH == ":memory:":
DATABASE_URL = "sqlite+aiosqlite:///:memory:" # 内存数据库的 URL
else:
DATABASE_URL = f"sqlite+aiosqlite:///{DATABASE_PATH}" # 文件数据库的 URL
logger.info(f"最终数据库路径: {DATABASE_PATH}")
logger.info(f"最终数据库 URL: {DATABASE_URL}")
DEFAULT_CONTEXT_TTL_DAYS = 30 # 定义默认的上下文生存时间 (天数)
IS_MEMORY_DB = DATABASE_PATH == ':memory:' # 判断是否使用内存数据库的布尔常量
# --- 数据库连接管理 ---
@asynccontextmanager # 使用异步上下文管理器装饰器
async def get_db_connection() -> AsyncGenerator[aiosqlite.Connection, None]:
"""
(可能已废弃/特定用途) 获取一个原生的 aiosqlite 数据库连接。
使用 asynccontextmanager 确保连接在使用后能被正确关闭。
注意:项目主要应使用 SQLAlchemy 的 AsyncSession。
"""
conn = None # 初始化连接变量
try:
# 连接到数据库文件
conn = await aiosqlite.connect(DATABASE_PATH) # 建立异步连接
logger.info("数据库连接已获取 (类型: aiosqlite.Connection)。") # 记录日志
yield conn # 返回连接对象供上下文使用
except Exception as e:
# 记录连接失败的错误
logger.error(f"获取 aiosqlite 数据库连接失败: {e}", exc_info=True)
raise # 重新抛出异常
finally:
# 无论是否发生异常,只要连接成功建立,就关闭连接
if conn is not None:
await conn.close() # 关闭连接
logger.debug("aiosqlite 数据库连接已关闭。") # 记录关闭日志
# --- 数据库初始化 ---
async def initialize_db_tables() -> None:
"""
初始化数据库表。如果表已存在,则不会重复创建。
此函数应在应用启动时调用。
"""
try:
# 调用内部的初始化函数
await _initialize_db_tables()
except Exception as e:
# 记录初始化失败的错误
logger.error(f"初始化数据库表失败: {e}", exc_info=True)
raise # 重新抛出异常,以便应用启动时能感知到问题
async def _initialize_db_tables() -> None:
"""
内部函数:使用 SQLAlchemy 的异步引擎和 Base.metadata 来创建所有定义的表。
"""
# 创建 SQLAlchemy 异步引擎实例
# echo=False 关闭 SQL 语句的日志输出,生产环境应保持关闭
engine = create_async_engine(DATABASE_URL, echo=False)
try:
# 使用引擎的 begin() 方法创建一个异步事务块
async with engine.begin() as conn:
# 在异步事务中同步运行 Base.metadata.create_all
# 这会检查数据库中是否存在 Base 中定义的所有表,如果不存在则创建
await conn.run_sync(Base.metadata.create_all)
# 记录初始化成功的日志
logger.info("所有通过 SQLAlchemy Base 定义的数据库表已成功初始化/验证。")
except Exception as e:
# 记录使用 SQLAlchemy 初始化失败的错误
logger.error(f"使用 SQLAlchemy 初始化数据库表失败: {e}", exc_info=True)
raise # 重新抛出异常
finally:
# 无论成功或失败,最后都关闭引擎及其连接池
await engine.dispose()
logger.debug("SQLAlchemy 异步引擎已关闭。") # 记录引擎关闭日志
# --- ApiKey CRUD (创建、读取、更新、删除) 操作 ---
async def add_api_key(
db: AsyncSession, # 异步数据库会话
key_string: str, # 要添加的 Key 字符串
description: Optional[str] = None, # Key 的描述 (可选)
expires_at: Optional[datetime] = None, # Key 的过期时间 (可选, 应为 UTC datetime)
is_active: bool = True, # Key 是否激活 (默认 True)
enable_context_completion: bool = True, # 是否启用上下文补全 (默认 True)
user_id: Optional[str] = None # 关联的用户 ID (可选)
) -> Optional[ApiKey]:
"""
向数据库异步添加一个新的 API Key 记录。
Args:
db (AsyncSession): SQLAlchemy 异步数据库会话。
key_string (str): 要添加的 API Key 字符串。
description (Optional[str]): Key 的描述信息。
expires_at (Optional[datetime]): Key 的过期时间 (UTC)。如果提供,请确保是 aware datetime。
is_active (bool): Key 是否激活。
enable_context_completion (bool): 此 Key 是否启用上下文补全。
user_id (Optional[str]): 与此 Key 关联的用户 ID。
Returns:
Optional[ApiKey]: 如果成功创建,返回包含数据库生成信息的 ApiKey 对象;
如果 Key 已存在或发生其他错误,则返回 None。
"""
try:
# 1. 检查 Key 是否已存在,避免重复添加
stmt_check = select(ApiKey).where(ApiKey.key_string == key_string) # 构建查询语句
result_check = await db.execute(stmt_check) # 执行查询
existing_key = result_check.scalar_one_or_none() # 获取单个结果,不存在则为 None
if existing_key: # 如果 Key 已存在
logger.warning(f"尝试添加已存在的 API Key: {key_string[:8]}...") # 记录警告
return None # 返回 None 表示未添加
# 2. 创建新的 ApiKey ORM 对象
new_api_key = ApiKey(
key_string=key_string, # 设置 Key 字符串
description=description, # 设置描述
# 如果提供了 expires_at,确保它是 UTC 时间;否则为 None
expires_at=expires_at.replace(tzinfo=timezone.utc) if expires_at else None,
is_active=is_active, # 设置激活状态
enable_context_completion=enable_context_completion, # 设置上下文补全状态
user_id=user_id # 设置关联用户 ID
# created_at 字段由数据库自动设置 (server_default=func.now())
)
# 3. 将新对象添加到会话中
db.add(new_api_key)
# 4. 提交事务以保存到数据库
await db.commit()
# 5. 刷新对象以获取数据库生成的 ID 和默认值 (如 created_at)
await db.refresh(new_api_key)
# 记录成功添加的日志
logger.info(f"成功添加 API Key: {new_api_key.key_string[:8]}... (ID: {new_api_key.id})")
return new_api_key # 返回创建的对象
except sqlalchemy.exc.IntegrityError as e: # 捕获唯一约束冲突错误 (通常是 key_string 重复)
await db.rollback() # 回滚事务
logger.error(f"添加 API Key 时发生唯一约束冲突 (Key 可能已存在): {key_string[:8]}... - {e}", exc_info=False) # 记录错误 (不包含堆栈)
return None # 返回 None
except Exception as e: # 捕获其他可能的数据库异常
await db.rollback() # 回滚事务
logger.error(f"添加 API Key {key_string[:8]}... 失败: {e}", exc_info=True) # 记录详细错误日志
return None # 返回 None
async def get_all_api_keys_from_db(db: AsyncSession) -> List[ApiKey]:
"""
从数据库异步获取所有 API 密钥对象。
Args:
db (AsyncSession): SQLAlchemy 异步数据库会话。
Returns:
List[ApiKey]: 包含所有 ApiKey 对象的列表,按创建时间排序。如果出错则返回空列表。
"""
try:
# 构建查询语句,选择所有 ApiKey 记录,并按创建时间降序排序
stmt = select(ApiKey).order_by(ApiKey.created_at.desc())
# 执行查询
result = await db.execute(stmt)
# 获取所有结果并将标量(ApiKey 对象)转换为列表
api_keys = result.scalars().all()
# 记录成功获取的 Key 数量
logger.info(f"成功从数据库获取 {len(api_keys)} 个 API Key。")
return list(api_keys) # 确保返回的是列表类型
except Exception as e: # 捕获可能的数据库异常
# 记录获取失败的错误日志
logger.error(f"从数据库获取所有 API Key 失败: {e}", exc_info=True)
return [] # 出错时返回空列表
async def get_api_key_by_string(db: AsyncSession, key_string: str) -> Optional[ApiKey]:
"""
根据 Key 字符串从数据库异步获取单个 API Key 对象。
Args:
db (AsyncSession): SQLAlchemy 异步数据库会话。
key_string (str): 要查询的 API Key 字符串。
Returns:
Optional[ApiKey]: 如果找到匹配的 Key,则返回 ApiKey 对象;否则返回 None。
"""
if not key_string: return None # 如果 key_string 为空,直接返回 None
try:
# 构建查询语句,根据 key_string 查找 ApiKey
stmt = select(ApiKey).where(ApiKey.key_string == key_string)
# 执行查询
result = await db.execute(stmt)
# 获取单个结果,如果不存在则为 None
api_key = result.scalar_one_or_none()
if api_key: # 如果找到了 Key
logger.debug(f"成功获取 API Key: {key_string[:8]}...") # 记录调试日志
else: # 如果未找到 Key
logger.debug(f"未找到 API Key: {key_string[:8]}...") # 记录调试日志
return api_key # 返回找到的对象或 None
except Exception as e: # 捕获可能的数据库异常
# 记录获取失败的错误日志
logger.error(f"获取 API Key {key_string[:8]}... 失败: {e}", exc_info=True)
return None # 出错时返回 None
async def update_api_key(db: AsyncSession, key_string: str, updates: Dict[str, Any]) -> Optional[ApiKey]:
"""
异步更新数据库中指定 API Key 的信息。
Args:
db (AsyncSession): SQLAlchemy 异步数据库会话。
key_string (str): 要更新的 API Key 字符串。
updates (Dict[str, Any]): 一个字典,包含要更新的字段名和对应的新值。
例如: {'description': '新描述', 'is_active': False}。
Returns:
Optional[ApiKey]: 如果成功更新,返回更新后的 ApiKey 对象;
如果 Key 不存在或发生错误,则返回 None。
"""
if not key_string or not updates: return None # 如果 key_string 或 updates 为空,直接返回 None
try:
# 1. 先查询要更新的 Key 是否存在
stmt_select = select(ApiKey).where(ApiKey.key_string == key_string)
result_select = await db.execute(stmt_select)
api_key_to_update = result_select.scalar_one_or_none() # 获取要更新的对象
if not api_key_to_update: # 如果 Key 不存在
logger.warning(f"尝试更新不存在的 API Key: {key_string[:8]}...") # 记录警告
return None # 返回 None
# 2. 过滤掉不允许直接更新的字段 (如 id, key_string, created_at)
# 确保只更新模型中存在的、且允许修改的字段
allowed_updates = {
k: v for k, v in updates.items()
if hasattr(ApiKey, k) and k not in ['id', 'key_string', 'created_at']
}
# 3. 特殊处理 expires_at 字段,确保存入数据库的是 UTC 时间
if 'expires_at' in allowed_updates:
expiry_value = allowed_updates['expires_at']
if isinstance(expiry_value, datetime): # 如果是 datetime 对象
# 确保它是 aware datetime (带时区信息),并转换为 UTC
if expiry_value.tzinfo is None:
logger.warning(f"更新 Key {key_string[:8]}... 的 expires_at 时提供了 naive datetime,假设为本地时间并转换为 UTC。建议提供 aware datetime。")
# 假设本地时区,转换为 UTC (这里可能需要更健壮的时区处理)
# 或者直接要求调用者提供 aware datetime
allowed_updates['expires_at'] = expiry_value.astimezone(timezone.utc)
else:
allowed_updates['expires_at'] = expiry_value.astimezone(timezone.utc)
elif expiry_value is None: # 允许将过期时间设置为 None (永不过期)
pass # 保留 None 值
else: # 如果提供的 expires_at 不是 datetime 或 None
logger.warning(f"更新 Key {key_string[:8]}... 时提供的 expires_at ('{expiry_value}') 不是有效的 datetime 对象或 None,已忽略此字段。") # 记录警告
del allowed_updates['expires_at'] # 从更新字典中移除无效的 expires_at
# 4. 检查是否还有有效的更新字段
if not allowed_updates:
logger.warning(f"没有有效的字段需要为 Key {key_string[:8]}... 更新。") # 记录警告
return api_key_to_update # 返回未修改的对象
# 5. 执行更新操作
stmt_update = (
update(ApiKey) # 构建 UPDATE 语句
.where(ApiKey.key_string == key_string) # 指定更新条件
.values(**allowed_updates) # 设置要更新的值
.execution_options(synchronize_session="fetch") # 选项:在执行后同步会话状态
)
await db.execute(stmt_update) # 执行更新语句
await db.commit() # 提交事务
await db.refresh(api_key_to_update) # 刷新对象以获取数据库中更新后的值
logger.info(f"成功更新 API Key: {key_string[:8]}... 更新内容: {allowed_updates}") # 记录成功更新日志
return api_key_to_update # 返回更新后的对象
except Exception as e: # 捕获可能的数据库异常
await db.rollback() # 回滚事务
logger.error(f"更新 API Key {key_string[:8]}... 失败: {e}", exc_info=True) # 记录错误日志
return None # 返回 None
async def delete_api_key(db: AsyncSession, key_string: str) -> bool:
"""
从数据库异步删除指定的 API Key。
Args:
db (AsyncSession): SQLAlchemy 异步数据库会话。
key_string (str): 要删除的 API Key 字符串。
Returns:
bool: 如果成功删除 Key (或 Key 原本就不存在),返回 True;如果发生错误,返回 False。
注意:即使 Key 不存在,rowcount 也可能为 0,但操作本身没有错误。
可以根据需求调整返回值逻辑,例如严格要求 rowcount > 0 才返回 True。
"""
if not key_string: return False # 如果 key_string 为空,返回 False
try:
# 构建 DELETE 语句
stmt = delete(ApiKey).where(ApiKey.key_string == key_string)
# 执行删除语句
result = await db.execute(stmt)
# 提交事务
await db.commit()
# 检查受影响的行数
if result.rowcount > 0: # 如果删除了至少一行
logger.info(f"成功删除 API Key: {key_string[:8]}...") # 记录成功日志
return True
else: # 如果没有删除任何行 (Key 可能原本就不存在)
logger.warning(f"尝试删除 API Key: {key_string[:8]}... 时未找到匹配项。") # 记录警告
return True # 也可以认为操作成功完成(目标状态已达成)
except Exception as e: # 捕获可能的数据库异常
await db.rollback() # 回滚事务
logger.error(f"删除 API Key {key_string[:8]}... 失败: {e}", exc_info=True) # 记录错误日志
return False # 返回 False 表示删除失败
# --- 其他数据库函数 (部分可能需要审查或完善) ---
async def is_valid_proxy_key(db: AsyncSession, key: str) -> bool:
"""
异步检查提供的代理 Key (字符串) 是否在数据库中存在且处于活动状态。
Args:
db (AsyncSession): SQLAlchemy 异步数据库会话。
key (str): 要检查的代理 Key 字符串。
Returns:
bool: 如果 Key 有效且激活,返回 True;否则返回 False。
"""
# 调用 get_api_key_by_string 获取 Key 对象
api_key_obj = await get_api_key_by_string(db, key)
# 检查对象是否存在,并且其 is_active 属性为 True
is_valid = bool(api_key_obj and api_key_obj.is_active)
# 记录调试日志
logger.debug(f"检查 Key '{key[:8]}...' 有效性: {is_valid}")
return is_valid # 返回检查结果
async def get_key_id_by_cached_content_id(db: AsyncSession, cached_content_id: str) -> Optional[int]:
"""
(需要实现) 根据缓存内容 ID 获取关联的 Key ID。
目前是模拟实现。
Args:
db (AsyncSession): SQLAlchemy 异步数据库会话。
cached_content_id (str): 缓存内容的唯一 ID。
Returns:
Optional[int]: 关联的 Key ID,如果找不到或出错则返回 None。
"""
# logger.warning("get_key_id_by_cached_content_id 函数尚未完全实现,返回模拟数据。") # 移除警告
from app.core.database.models import CachedContent # 导入模型
try:
# 假设 cached_content_id 是 CachedContent.content_id (通常是哈希或 Gemini cache name)
# 或者,如果它可能是整数 CachedContent.id,则需要类型检查
if isinstance(cached_content_id, int): # 如果传入的是整数 ID
stmt = select(CachedContent.key_id).where(CachedContent.id == cached_content_id)
else: # 否则按 content_id (字符串) 查询
stmt = select(CachedContent.key_id).where(CachedContent.content_id == cached_content_id)
result = await db.execute(stmt)
key_id = result.scalar_one_or_none()
if key_id is not None:
logger.info(f"成功从缓存标识符 '{cached_content_id}' 获取到 key_id: {key_id}")
else:
logger.info(f"未从缓存标识符 '{cached_content_id}' 找到关联的 key_id。")
return key_id
except Exception as e:
logger.error(f"根据缓存内容标识符 '{cached_content_id}' 获取 Key ID 失败: {e}", exc_info=True)
return None
async def get_key_string_by_id(db: AsyncSession, key_id: int) -> Optional[str]:
"""
根据 Key ID 获取 Key 字符串。
"""
# logger.warning("get_key_string_by_id 函数尚未完全实现,返回模拟数据。") # 移除警告
from app.core.database.models import ApiKey # 导入模型
try:
stmt = select(ApiKey.key_string).where(ApiKey.id == key_id)
result = await db.execute(stmt)
key_string = result.scalar_one_or_none()
if key_string:
logger.info(f"成功为 Key ID {key_id} 获取到 Key 字符串: {key_string[:8]}...")
else:
logger.info(f"未为 Key ID {key_id} 找到对应的 Key 字符串。")
return key_string
except Exception as e:
logger.error(f"根据 Key ID {key_id} 获取 Key 字符串失败: {e}", exc_info=True)
return None
async def get_user_last_used_key_id(db: AsyncSession, user_id: str) -> Optional[int]:
"""
获取指定用户上次成功使用的 Key ID (用于粘性会话)。
"""
# logger.warning("get_user_last_used_key_id 函数尚未完全实现,返回模拟数据。") # 移除警告
from app.core.database.models import UserKeyAssociation # 导入模型
try:
stmt = select(UserKeyAssociation.key_id)\
.where(UserKeyAssociation.user_id == user_id)\
.order_by(UserKeyAssociation.last_used_timestamp.desc())\
.limit(1)
result = await db.execute(stmt)
key_id = result.scalar_one_or_none()
if key_id is not None:
logger.info(f"用户 '{user_id}' 上次使用的 Key ID: {key_id}")
else:
logger.info(f"未找到用户 '{user_id}' 上次使用的 Key ID。")
return key_id
except Exception as e:
logger.error(f"获取用户 '{user_id}' 上次成功使用的 Key ID 失败: {e}", exc_info=True)
return None
async def get_key_scores(model_name: str) -> Dict[str, float]:
"""
(需要实现/可能已废弃) 获取指定模型的 Key 分数。
目前从内存缓存 key_scores_cache 获取 (如果存在)。
Args:
model_name (str): 模型名称。
Returns:
Dict[str, float]: 包含 Key 字符串和对应分数的字典。
"""
logger.debug("get_key_scores 函数可能依赖于内存缓存或未完全实现的数据库逻辑。") # 记录警告改为 debug
try:
# TODO: 实现从 KeyScore 表获取分数的逻辑,或者确认是否依赖内存缓存
# 暂时从内存缓存获取
with cache_lock: # 使用锁保护缓存访问
scores = key_scores_cache.get(model_name, {}) # 获取模型对应的分数,默认为空字典
return scores.copy() # 返回缓存的副本
except Exception as e:
logger.error(f"获取模型 '{model_name}' 的 Key 分数失败: {e}", exc_info=True) # 记录错误
return {} # 出错时返回空字典
async def update_setting(db: AsyncSession, key: str, value: str):
"""
(重复/可能已废弃) 更新或插入设置项。
此函数与 set_setting 功能重复,且使用了原生 SQL。建议统一使用 set_setting。
Args:
db (AsyncSession): SQLAlchemy 异步数据库会话。
key (str): 设置项的键。
value (str): 设置项的值。
"""
logger.warning("调用了可能重复或已废弃的 update_setting 函数,建议使用 set_setting。") # 记录警告
try:
# 使用 SQLAlchemy 的 text() 构造原生 SQL 更新语句 (不推荐,优先使用 ORM 或 Core API)
# 注意:直接使用 text() 可能存在 SQL 注入风险,如果 key 或 value 来自外部输入。
# 在这个特定场景下,key 和 value 通常是内部定义的,风险较低,但仍不推荐。
# 推荐使用 SQLAlchemy 的 update() 构造器,如 update_api_key 函数所示。
stmt = text("UPDATE settings SET value = :value WHERE key = :key") # 构造 SQL 语句
parameters = {"key": key, "value": value} # 创建参数字典
result = await db.execute(stmt, parameters) # 执行 SQL 语句
if result.rowcount == 0: # 如果没有行被更新 (说明 key 不存在)
# 如果 key 不存在,则插入新行 (同样使用原生 SQL)
stmt_insert = text("INSERT INTO settings (key, value) VALUES (:key, :value)")
await db.execute(stmt_insert, parameters)
logger.info(f"设置 '{key}' 不存在,已插入新值 '{value}'") # 记录插入日志
else:
logger.info(f"设置 '{key}' 已更新为 '{value}' (通过 update_setting)") # 记录更新日志
await db.commit() # 提交事务
except Exception as e:
logger.error(f"使用 update_setting 更新设置 '{key}' 时发生错误: {e}", exc_info=True) # 记录错误日志
await db.rollback() # 回滚事务
|