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() # 回滚事务