|
|
""" |
|
|
SQLite 存储管理器 |
|
|
""" |
|
|
|
|
|
import asyncio |
|
|
import json |
|
|
import os |
|
|
import time |
|
|
from typing import Any, Dict, List, Optional, Tuple |
|
|
|
|
|
import aiosqlite |
|
|
|
|
|
from log import log |
|
|
|
|
|
|
|
|
class SQLiteManager: |
|
|
"""SQLite 数据库管理器""" |
|
|
|
|
|
|
|
|
STATE_FIELDS = { |
|
|
"error_codes", |
|
|
"disabled", |
|
|
"last_success", |
|
|
"user_email", |
|
|
"model_cooldowns", |
|
|
} |
|
|
|
|
|
|
|
|
REQUIRED_COLUMNS = { |
|
|
"credentials": [ |
|
|
("disabled", "INTEGER DEFAULT 0"), |
|
|
("error_codes", "TEXT DEFAULT '[]'"), |
|
|
("last_success", "REAL"), |
|
|
("user_email", "TEXT"), |
|
|
("model_cooldowns", "TEXT DEFAULT '{}'"), |
|
|
("rotation_order", "INTEGER DEFAULT 0"), |
|
|
("call_count", "INTEGER DEFAULT 0"), |
|
|
("created_at", "REAL DEFAULT (unixepoch())"), |
|
|
("updated_at", "REAL DEFAULT (unixepoch())") |
|
|
], |
|
|
"antigravity_credentials": [ |
|
|
("disabled", "INTEGER DEFAULT 0"), |
|
|
("error_codes", "TEXT DEFAULT '[]'"), |
|
|
("last_success", "REAL"), |
|
|
("user_email", "TEXT"), |
|
|
("model_cooldowns", "TEXT DEFAULT '{}'"), |
|
|
("rotation_order", "INTEGER DEFAULT 0"), |
|
|
("call_count", "INTEGER DEFAULT 0"), |
|
|
("created_at", "REAL DEFAULT (unixepoch())"), |
|
|
("updated_at", "REAL DEFAULT (unixepoch())") |
|
|
] |
|
|
} |
|
|
|
|
|
def __init__(self): |
|
|
self._db_path = None |
|
|
self._credentials_dir = None |
|
|
self._initialized = False |
|
|
self._lock = asyncio.Lock() |
|
|
|
|
|
|
|
|
self._config_cache: Dict[str, Any] = {} |
|
|
self._config_loaded = False |
|
|
|
|
|
async def initialize(self) -> None: |
|
|
"""初始化 SQLite 数据库""" |
|
|
if self._initialized: |
|
|
return |
|
|
|
|
|
async with self._lock: |
|
|
if self._initialized: |
|
|
return |
|
|
|
|
|
try: |
|
|
|
|
|
self._credentials_dir = os.getenv("CREDENTIALS_DIR", "./creds") |
|
|
self._db_path = os.path.join(self._credentials_dir, "credentials.db") |
|
|
|
|
|
|
|
|
os.makedirs(self._credentials_dir, exist_ok=True) |
|
|
|
|
|
|
|
|
async with aiosqlite.connect(self._db_path) as db: |
|
|
|
|
|
await db.execute("PRAGMA journal_mode=WAL") |
|
|
await db.execute("PRAGMA foreign_keys=ON") |
|
|
|
|
|
|
|
|
await self._ensure_schema_compatibility(db) |
|
|
|
|
|
|
|
|
await self._create_tables(db) |
|
|
|
|
|
await db.commit() |
|
|
|
|
|
|
|
|
await self._load_config_cache() |
|
|
|
|
|
self._initialized = True |
|
|
log.info(f"SQLite storage initialized at {self._db_path}") |
|
|
|
|
|
except Exception as e: |
|
|
log.error(f"Error initializing SQLite: {e}") |
|
|
raise |
|
|
|
|
|
async def _ensure_schema_compatibility(self, db: aiosqlite.Connection) -> None: |
|
|
""" |
|
|
确保数据库结构兼容,自动修复缺失的列 |
|
|
""" |
|
|
try: |
|
|
|
|
|
for table_name, columns in self.REQUIRED_COLUMNS.items(): |
|
|
|
|
|
async with db.execute( |
|
|
"SELECT name FROM sqlite_master WHERE type='table' AND name=?", |
|
|
(table_name,) |
|
|
) as cursor: |
|
|
if not await cursor.fetchone(): |
|
|
log.debug(f"Table {table_name} does not exist, will be created") |
|
|
continue |
|
|
|
|
|
|
|
|
async with db.execute(f"PRAGMA table_info({table_name})") as cursor: |
|
|
existing_columns = {row[1] for row in await cursor.fetchall()} |
|
|
|
|
|
|
|
|
added_count = 0 |
|
|
for col_name, col_def in columns: |
|
|
if col_name not in existing_columns: |
|
|
try: |
|
|
await db.execute(f"ALTER TABLE {table_name} ADD COLUMN {col_name} {col_def}") |
|
|
log.info(f"Added missing column {table_name}.{col_name}") |
|
|
added_count += 1 |
|
|
except Exception as e: |
|
|
log.error(f"Failed to add column {table_name}.{col_name}: {e}") |
|
|
|
|
|
if added_count > 0: |
|
|
log.info(f"Table {table_name}: added {added_count} missing column(s)") |
|
|
|
|
|
except Exception as e: |
|
|
log.error(f"Error ensuring schema compatibility: {e}") |
|
|
|
|
|
|
|
|
async def _create_tables(self, db: aiosqlite.Connection): |
|
|
"""创建数据库表和索引""" |
|
|
|
|
|
await db.execute(""" |
|
|
CREATE TABLE IF NOT EXISTS credentials ( |
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT, |
|
|
filename TEXT UNIQUE NOT NULL, |
|
|
credential_data TEXT NOT NULL, |
|
|
|
|
|
-- 状态字段 |
|
|
disabled INTEGER DEFAULT 0, |
|
|
error_codes TEXT DEFAULT '[]', |
|
|
last_success REAL, |
|
|
user_email TEXT, |
|
|
|
|
|
-- 模型级 CD 支持 (JSON: {model_key: cooldown_timestamp}) |
|
|
model_cooldowns TEXT DEFAULT '{}', |
|
|
|
|
|
-- 轮换相关 |
|
|
rotation_order INTEGER DEFAULT 0, |
|
|
call_count INTEGER DEFAULT 0, |
|
|
|
|
|
-- 时间戳 |
|
|
created_at REAL DEFAULT (unixepoch()), |
|
|
updated_at REAL DEFAULT (unixepoch()) |
|
|
) |
|
|
""") |
|
|
|
|
|
|
|
|
await db.execute(""" |
|
|
CREATE TABLE IF NOT EXISTS antigravity_credentials ( |
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT, |
|
|
filename TEXT UNIQUE NOT NULL, |
|
|
credential_data TEXT NOT NULL, |
|
|
|
|
|
-- 状态字段 |
|
|
disabled INTEGER DEFAULT 0, |
|
|
error_codes TEXT DEFAULT '[]', |
|
|
last_success REAL, |
|
|
user_email TEXT, |
|
|
|
|
|
-- 模型级 CD 支持 (JSON: {model_name: cooldown_timestamp}) |
|
|
model_cooldowns TEXT DEFAULT '{}', |
|
|
|
|
|
-- 轮换相关 |
|
|
rotation_order INTEGER DEFAULT 0, |
|
|
call_count INTEGER DEFAULT 0, |
|
|
|
|
|
-- 时间戳 |
|
|
created_at REAL DEFAULT (unixepoch()), |
|
|
updated_at REAL DEFAULT (unixepoch()) |
|
|
) |
|
|
""") |
|
|
|
|
|
|
|
|
await db.execute(""" |
|
|
CREATE INDEX IF NOT EXISTS idx_disabled |
|
|
ON credentials(disabled) |
|
|
""") |
|
|
await db.execute(""" |
|
|
CREATE INDEX IF NOT EXISTS idx_rotation_order |
|
|
ON credentials(rotation_order) |
|
|
""") |
|
|
|
|
|
|
|
|
await db.execute(""" |
|
|
CREATE INDEX IF NOT EXISTS idx_ag_disabled |
|
|
ON antigravity_credentials(disabled) |
|
|
""") |
|
|
await db.execute(""" |
|
|
CREATE INDEX IF NOT EXISTS idx_ag_rotation_order |
|
|
ON antigravity_credentials(rotation_order) |
|
|
""") |
|
|
|
|
|
|
|
|
await db.execute(""" |
|
|
CREATE TABLE IF NOT EXISTS config ( |
|
|
key TEXT PRIMARY KEY, |
|
|
value TEXT NOT NULL, |
|
|
updated_at REAL DEFAULT (unixepoch()) |
|
|
) |
|
|
""") |
|
|
|
|
|
log.debug("SQLite tables and indexes created") |
|
|
|
|
|
async def _load_config_cache(self): |
|
|
"""加载配置到内存缓存(仅在初始化时调用一次)""" |
|
|
if self._config_loaded: |
|
|
return |
|
|
|
|
|
try: |
|
|
async with aiosqlite.connect(self._db_path) as db: |
|
|
async with db.execute("SELECT key, value FROM config") as cursor: |
|
|
rows = await cursor.fetchall() |
|
|
|
|
|
for key, value in rows: |
|
|
try: |
|
|
self._config_cache[key] = json.loads(value) |
|
|
except json.JSONDecodeError: |
|
|
self._config_cache[key] = value |
|
|
|
|
|
self._config_loaded = True |
|
|
log.debug(f"Loaded {len(self._config_cache)} config items into cache") |
|
|
|
|
|
except Exception as e: |
|
|
log.error(f"Error loading config cache: {e}") |
|
|
self._config_cache = {} |
|
|
|
|
|
async def close(self) -> None: |
|
|
"""关闭数据库连接""" |
|
|
self._initialized = False |
|
|
log.debug("SQLite storage closed") |
|
|
|
|
|
def _ensure_initialized(self): |
|
|
"""确保已初始化""" |
|
|
if not self._initialized: |
|
|
raise RuntimeError("SQLite manager not initialized") |
|
|
|
|
|
def _get_table_name(self, mode: str) -> str: |
|
|
"""根据 mode 获取对应的表名""" |
|
|
if mode == "antigravity": |
|
|
return "antigravity_credentials" |
|
|
elif mode == "geminicli": |
|
|
return "credentials" |
|
|
else: |
|
|
raise ValueError(f"Invalid mode: {mode}. Must be 'geminicli' or 'antigravity'") |
|
|
|
|
|
|
|
|
|
|
|
async def get_next_available_credential( |
|
|
self, mode: str = "geminicli", model_key: Optional[str] = None |
|
|
) -> Optional[Tuple[str, Dict[str, Any]]]: |
|
|
""" |
|
|
随机获取一个可用凭证(负载均衡) |
|
|
- 未禁用 |
|
|
- 如果提供了 model_key,还会检查模型级冷却 |
|
|
- 随机选择 |
|
|
|
|
|
Args: |
|
|
mode: 凭证模式 ("geminicli" 或 "antigravity") |
|
|
model_key: 模型键(用于模型级冷却检查,antigravity 用模型名,gcli 用 pro/flash) |
|
|
|
|
|
Note: |
|
|
- 对于 antigravity: model_key 是具体模型名(如 "gemini-2.0-flash-exp") |
|
|
- 对于 gcli: model_key 是 "pro" 或 "flash" |
|
|
""" |
|
|
self._ensure_initialized() |
|
|
|
|
|
try: |
|
|
table_name = self._get_table_name(mode) |
|
|
async with aiosqlite.connect(self._db_path) as db: |
|
|
current_time = time.time() |
|
|
|
|
|
|
|
|
async with db.execute(f""" |
|
|
SELECT filename, credential_data, model_cooldowns |
|
|
FROM {table_name} |
|
|
WHERE disabled = 0 |
|
|
ORDER BY RANDOM() |
|
|
""") as cursor: |
|
|
rows = await cursor.fetchall() |
|
|
|
|
|
|
|
|
if not model_key: |
|
|
if rows: |
|
|
filename, credential_json, _ = rows[0] |
|
|
credential_data = json.loads(credential_json) |
|
|
return filename, credential_data |
|
|
return None |
|
|
|
|
|
|
|
|
for filename, credential_json, model_cooldowns_json in rows: |
|
|
model_cooldowns = json.loads(model_cooldowns_json or '{}') |
|
|
|
|
|
|
|
|
model_cooldown = model_cooldowns.get(model_key) |
|
|
if model_cooldown is None or current_time >= model_cooldown: |
|
|
|
|
|
credential_data = json.loads(credential_json) |
|
|
return filename, credential_data |
|
|
|
|
|
return None |
|
|
|
|
|
except Exception as e: |
|
|
log.error(f"Error getting next available credential (mode={mode}, model_key={model_key}): {e}") |
|
|
return None |
|
|
|
|
|
async def get_available_credentials_list(self) -> List[str]: |
|
|
""" |
|
|
获取所有可用凭证列表 |
|
|
- 未禁用 |
|
|
- 按轮换顺序排序 |
|
|
""" |
|
|
self._ensure_initialized() |
|
|
|
|
|
try: |
|
|
async with aiosqlite.connect(self._db_path) as db: |
|
|
async with db.execute(""" |
|
|
SELECT filename |
|
|
FROM credentials |
|
|
WHERE disabled = 0 |
|
|
ORDER BY rotation_order ASC |
|
|
""") as cursor: |
|
|
rows = await cursor.fetchall() |
|
|
return [row[0] for row in rows] |
|
|
|
|
|
except Exception as e: |
|
|
log.error(f"Error getting available credentials list: {e}") |
|
|
return [] |
|
|
|
|
|
|
|
|
|
|
|
async def store_credential(self, filename: str, credential_data: Dict[str, Any], mode: str = "geminicli") -> bool: |
|
|
"""存储或更新凭证""" |
|
|
self._ensure_initialized() |
|
|
|
|
|
try: |
|
|
table_name = self._get_table_name(mode) |
|
|
async with aiosqlite.connect(self._db_path) as db: |
|
|
|
|
|
async with db.execute(f""" |
|
|
SELECT disabled, error_codes, last_success, user_email, |
|
|
rotation_order, call_count |
|
|
FROM {table_name} WHERE filename = ? |
|
|
""", (filename,)) as cursor: |
|
|
existing = await cursor.fetchone() |
|
|
|
|
|
if existing: |
|
|
|
|
|
await db.execute(f""" |
|
|
UPDATE {table_name} |
|
|
SET credential_data = ?, |
|
|
updated_at = unixepoch() |
|
|
WHERE filename = ? |
|
|
""", (json.dumps(credential_data), filename)) |
|
|
else: |
|
|
|
|
|
async with db.execute(f""" |
|
|
SELECT COALESCE(MAX(rotation_order), -1) + 1 FROM {table_name} |
|
|
""") as cursor: |
|
|
row = await cursor.fetchone() |
|
|
next_order = row[0] |
|
|
|
|
|
await db.execute(f""" |
|
|
INSERT INTO {table_name} |
|
|
(filename, credential_data, rotation_order, last_success) |
|
|
VALUES (?, ?, ?, ?) |
|
|
""", (filename, json.dumps(credential_data), next_order, time.time())) |
|
|
|
|
|
await db.commit() |
|
|
log.debug(f"Stored credential: {filename} (mode={mode})") |
|
|
return True |
|
|
|
|
|
except Exception as e: |
|
|
log.error(f"Error storing credential {filename}: {e}") |
|
|
return False |
|
|
|
|
|
async def get_credential(self, filename: str, mode: str = "geminicli") -> Optional[Dict[str, Any]]: |
|
|
"""获取凭证数据,支持basename匹配以兼容旧数据""" |
|
|
self._ensure_initialized() |
|
|
|
|
|
try: |
|
|
table_name = self._get_table_name(mode) |
|
|
async with aiosqlite.connect(self._db_path) as db: |
|
|
|
|
|
async with db.execute(f""" |
|
|
SELECT credential_data FROM {table_name} WHERE filename = ? |
|
|
""", (filename,)) as cursor: |
|
|
row = await cursor.fetchone() |
|
|
if row: |
|
|
return json.loads(row[0]) |
|
|
|
|
|
|
|
|
async with db.execute(f""" |
|
|
SELECT credential_data FROM {table_name} |
|
|
WHERE filename LIKE '%' || ? OR filename = ? |
|
|
""", (filename, filename)) as cursor: |
|
|
rows = await cursor.fetchall() |
|
|
|
|
|
for row in rows: |
|
|
return json.loads(row[0]) |
|
|
|
|
|
return None |
|
|
|
|
|
except Exception as e: |
|
|
log.error(f"Error getting credential {filename}: {e}") |
|
|
return None |
|
|
|
|
|
async def list_credentials(self, mode: str = "geminicli") -> List[str]: |
|
|
"""列出所有凭证文件名(包括禁用的)""" |
|
|
self._ensure_initialized() |
|
|
|
|
|
try: |
|
|
table_name = self._get_table_name(mode) |
|
|
async with aiosqlite.connect(self._db_path) as db: |
|
|
async with db.execute(f""" |
|
|
SELECT filename FROM {table_name} ORDER BY rotation_order |
|
|
""") as cursor: |
|
|
rows = await cursor.fetchall() |
|
|
return [row[0] for row in rows] |
|
|
|
|
|
except Exception as e: |
|
|
log.error(f"Error listing credentials: {e}") |
|
|
return [] |
|
|
|
|
|
async def delete_credential(self, filename: str, mode: str = "geminicli") -> bool: |
|
|
"""删除凭证,支持basename匹配以兼容旧数据""" |
|
|
self._ensure_initialized() |
|
|
|
|
|
try: |
|
|
table_name = self._get_table_name(mode) |
|
|
async with aiosqlite.connect(self._db_path) as db: |
|
|
|
|
|
result = await db.execute(f""" |
|
|
DELETE FROM {table_name} WHERE filename = ? |
|
|
""", (filename,)) |
|
|
deleted_count = result.rowcount |
|
|
|
|
|
|
|
|
if deleted_count == 0: |
|
|
result = await db.execute(f""" |
|
|
DELETE FROM {table_name} WHERE filename LIKE '%' || ? |
|
|
""", (filename,)) |
|
|
deleted_count = result.rowcount |
|
|
|
|
|
await db.commit() |
|
|
|
|
|
if deleted_count > 0: |
|
|
log.debug(f"Deleted {deleted_count} credential(s): {filename} (mode={mode})") |
|
|
return True |
|
|
else: |
|
|
log.warning(f"No credential found to delete: {filename} (mode={mode})") |
|
|
return False |
|
|
|
|
|
except Exception as e: |
|
|
log.error(f"Error deleting credential {filename}: {e}") |
|
|
return False |
|
|
|
|
|
async def update_credential_state(self, filename: str, state_updates: Dict[str, Any], mode: str = "geminicli") -> bool: |
|
|
"""更新凭证状态,支持basename匹配以兼容旧数据""" |
|
|
self._ensure_initialized() |
|
|
|
|
|
try: |
|
|
table_name = self._get_table_name(mode) |
|
|
log.debug(f"[DB] update_credential_state 开始: filename={filename}, state_updates={state_updates}, mode={mode}, table={table_name}") |
|
|
|
|
|
|
|
|
set_clauses = [] |
|
|
values = [] |
|
|
|
|
|
for key, value in state_updates.items(): |
|
|
if key in self.STATE_FIELDS: |
|
|
if key == "error_codes": |
|
|
set_clauses.append(f"{key} = ?") |
|
|
values.append(json.dumps(value)) |
|
|
elif key == "model_cooldowns": |
|
|
set_clauses.append(f"{key} = ?") |
|
|
values.append(json.dumps(value)) |
|
|
else: |
|
|
set_clauses.append(f"{key} = ?") |
|
|
values.append(value) |
|
|
|
|
|
if not set_clauses: |
|
|
log.info(f"[DB] 没有需要更新的状态字段") |
|
|
return True |
|
|
|
|
|
set_clauses.append("updated_at = unixepoch()") |
|
|
values.append(filename) |
|
|
|
|
|
log.debug(f"[DB] SQL参数: set_clauses={set_clauses}, values={values}") |
|
|
|
|
|
async with aiosqlite.connect(self._db_path) as db: |
|
|
|
|
|
sql_exact = f""" |
|
|
UPDATE {table_name} |
|
|
SET {', '.join(set_clauses)} |
|
|
WHERE filename = ? |
|
|
""" |
|
|
log.debug(f"[DB] 执行精确匹配SQL: {sql_exact}") |
|
|
log.debug(f"[DB] SQL参数值: {values}") |
|
|
|
|
|
result = await db.execute(sql_exact, values) |
|
|
updated_count = result.rowcount |
|
|
log.debug(f"[DB] 精确匹配 rowcount={updated_count}") |
|
|
|
|
|
|
|
|
if updated_count == 0: |
|
|
sql_basename = f""" |
|
|
UPDATE {table_name} |
|
|
SET {', '.join(set_clauses)} |
|
|
WHERE filename LIKE '%' || ? |
|
|
""" |
|
|
log.debug(f"[DB] 精确匹配失败,尝试basename匹配SQL: {sql_basename}") |
|
|
result = await db.execute(sql_basename, values) |
|
|
updated_count = result.rowcount |
|
|
log.info(f"[DB] basename匹配 rowcount={updated_count}") |
|
|
|
|
|
|
|
|
log.debug(f"[DB] 准备commit,总更新行数={updated_count}") |
|
|
await db.commit() |
|
|
log.debug(f"[DB] commit完成") |
|
|
|
|
|
success = updated_count > 0 |
|
|
log.debug(f"[DB] update_credential_state 结束: success={success}, updated_count={updated_count}") |
|
|
return success |
|
|
|
|
|
except Exception as e: |
|
|
log.error(f"[DB] Error updating credential state {filename}: {e}", exc_info=True) |
|
|
return False |
|
|
|
|
|
async def get_credential_state(self, filename: str, mode: str = "geminicli") -> Dict[str, Any]: |
|
|
"""获取凭证状态,支持basename匹配以兼容旧数据""" |
|
|
self._ensure_initialized() |
|
|
|
|
|
try: |
|
|
table_name = self._get_table_name(mode) |
|
|
async with aiosqlite.connect(self._db_path) as db: |
|
|
|
|
|
async with db.execute(f""" |
|
|
SELECT disabled, error_codes, last_success, user_email, model_cooldowns |
|
|
FROM {table_name} WHERE filename = ? |
|
|
""", (filename,)) as cursor: |
|
|
row = await cursor.fetchone() |
|
|
|
|
|
if row: |
|
|
error_codes_json = row[1] or '[]' |
|
|
model_cooldowns_json = row[4] or '{}' |
|
|
return { |
|
|
"disabled": bool(row[0]), |
|
|
"error_codes": json.loads(error_codes_json), |
|
|
"last_success": row[2] or time.time(), |
|
|
"user_email": row[3], |
|
|
"model_cooldowns": json.loads(model_cooldowns_json), |
|
|
} |
|
|
|
|
|
|
|
|
async with db.execute(f""" |
|
|
SELECT disabled, error_codes, last_success, user_email, model_cooldowns |
|
|
FROM {table_name} WHERE filename LIKE '%' || ? |
|
|
""", (filename,)) as cursor: |
|
|
row = await cursor.fetchone() |
|
|
|
|
|
if row: |
|
|
error_codes_json = row[1] or '[]' |
|
|
model_cooldowns_json = row[4] or '{}' |
|
|
return { |
|
|
"disabled": bool(row[0]), |
|
|
"error_codes": json.loads(error_codes_json), |
|
|
"last_success": row[2] or time.time(), |
|
|
"user_email": row[3], |
|
|
"model_cooldowns": json.loads(model_cooldowns_json), |
|
|
} |
|
|
|
|
|
|
|
|
return { |
|
|
"disabled": False, |
|
|
"error_codes": [], |
|
|
"last_success": time.time(), |
|
|
"user_email": None, |
|
|
"model_cooldowns": {}, |
|
|
} |
|
|
|
|
|
except Exception as e: |
|
|
log.error(f"Error getting credential state {filename}: {e}") |
|
|
return {} |
|
|
|
|
|
async def get_all_credential_states(self, mode: str = "geminicli") -> Dict[str, Dict[str, Any]]: |
|
|
"""获取所有凭证状态""" |
|
|
self._ensure_initialized() |
|
|
|
|
|
try: |
|
|
table_name = self._get_table_name(mode) |
|
|
async with aiosqlite.connect(self._db_path) as db: |
|
|
async with db.execute(f""" |
|
|
SELECT filename, disabled, error_codes, last_success, |
|
|
user_email, model_cooldowns |
|
|
FROM {table_name} |
|
|
""") as cursor: |
|
|
rows = await cursor.fetchall() |
|
|
|
|
|
states = {} |
|
|
current_time = time.time() |
|
|
|
|
|
for row in rows: |
|
|
filename = row[0] |
|
|
error_codes_json = row[2] or '[]' |
|
|
model_cooldowns_json = row[5] or '{}' |
|
|
model_cooldowns = json.loads(model_cooldowns_json) |
|
|
|
|
|
|
|
|
if model_cooldowns: |
|
|
model_cooldowns = { |
|
|
k: v for k, v in model_cooldowns.items() |
|
|
if v > current_time |
|
|
} |
|
|
|
|
|
states[filename] = { |
|
|
"disabled": bool(row[1]), |
|
|
"error_codes": json.loads(error_codes_json), |
|
|
"last_success": row[3] or time.time(), |
|
|
"user_email": row[4], |
|
|
"model_cooldowns": model_cooldowns, |
|
|
} |
|
|
|
|
|
return states |
|
|
|
|
|
except Exception as e: |
|
|
log.error(f"Error getting all credential states: {e}") |
|
|
return {} |
|
|
|
|
|
async def get_credentials_summary( |
|
|
self, |
|
|
offset: int = 0, |
|
|
limit: Optional[int] = None, |
|
|
status_filter: str = "all", |
|
|
mode: str = "geminicli", |
|
|
error_code_filter: Optional[str] = None, |
|
|
cooldown_filter: Optional[str] = None |
|
|
) -> Dict[str, Any]: |
|
|
""" |
|
|
获取凭证的摘要信息(不包含完整凭证数据)- 支持分页和状态筛选 |
|
|
|
|
|
Args: |
|
|
offset: 跳过的记录数(默认0) |
|
|
limit: 返回的最大记录数(None表示返回所有) |
|
|
status_filter: 状态筛选(all=全部, enabled=仅启用, disabled=仅禁用) |
|
|
mode: 凭证模式 ("geminicli" 或 "antigravity") |
|
|
error_code_filter: 错误码筛选(格式如"400"或"403",筛选包含该错误码的凭证) |
|
|
cooldown_filter: 冷却状态筛选("in_cooldown"=冷却中, "no_cooldown"=未冷却) |
|
|
|
|
|
Returns: |
|
|
包含 items(凭证列表)、total(总数)、offset、limit 的字典 |
|
|
""" |
|
|
self._ensure_initialized() |
|
|
|
|
|
try: |
|
|
|
|
|
table_name = self._get_table_name(mode) |
|
|
|
|
|
async with aiosqlite.connect(self._db_path) as db: |
|
|
|
|
|
global_stats = {"total": 0, "normal": 0, "disabled": 0} |
|
|
async with db.execute(f""" |
|
|
SELECT disabled, COUNT(*) FROM {table_name} GROUP BY disabled |
|
|
""") as stats_cursor: |
|
|
stats_rows = await stats_cursor.fetchall() |
|
|
for disabled, count in stats_rows: |
|
|
global_stats["total"] += count |
|
|
if disabled: |
|
|
global_stats["disabled"] = count |
|
|
else: |
|
|
global_stats["normal"] = count |
|
|
|
|
|
|
|
|
where_clauses = [] |
|
|
count_params = [] |
|
|
|
|
|
if status_filter == "enabled": |
|
|
where_clauses.append("disabled = 0") |
|
|
elif status_filter == "disabled": |
|
|
where_clauses.append("disabled = 1") |
|
|
|
|
|
filter_value = None |
|
|
filter_int = None |
|
|
if error_code_filter and str(error_code_filter).strip().lower() != "all": |
|
|
filter_value = str(error_code_filter).strip() |
|
|
try: |
|
|
filter_int = int(filter_value) |
|
|
except ValueError: |
|
|
filter_int = None |
|
|
|
|
|
|
|
|
where_clause = "" |
|
|
if where_clauses: |
|
|
where_clause = "WHERE " + " AND ".join(where_clauses) |
|
|
|
|
|
|
|
|
all_query = f""" |
|
|
SELECT filename, disabled, error_codes, last_success, |
|
|
user_email, rotation_order, model_cooldowns |
|
|
FROM {table_name} |
|
|
{where_clause} |
|
|
ORDER BY rotation_order |
|
|
""" |
|
|
|
|
|
async with db.execute(all_query, count_params) as cursor: |
|
|
all_rows = await cursor.fetchall() |
|
|
|
|
|
current_time = time.time() |
|
|
all_summaries = [] |
|
|
|
|
|
for row in all_rows: |
|
|
filename = row[0] |
|
|
error_codes_json = row[2] or '[]' |
|
|
model_cooldowns_json = row[6] or '{}' |
|
|
model_cooldowns = json.loads(model_cooldowns_json) |
|
|
|
|
|
|
|
|
active_cooldowns = {} |
|
|
if model_cooldowns: |
|
|
active_cooldowns = { |
|
|
k: v for k, v in model_cooldowns.items() |
|
|
if v > current_time |
|
|
} |
|
|
|
|
|
error_codes = json.loads(error_codes_json) |
|
|
if filter_value: |
|
|
match = False |
|
|
for code in error_codes: |
|
|
if code == filter_value or code == filter_int: |
|
|
match = True |
|
|
break |
|
|
if isinstance(code, str) and filter_int is not None: |
|
|
try: |
|
|
if int(code) == filter_int: |
|
|
match = True |
|
|
break |
|
|
except ValueError: |
|
|
pass |
|
|
if not match: |
|
|
continue |
|
|
|
|
|
summary = { |
|
|
"filename": filename, |
|
|
"disabled": bool(row[1]), |
|
|
"error_codes": error_codes, |
|
|
"last_success": row[3] or current_time, |
|
|
"user_email": row[4], |
|
|
"rotation_order": row[5], |
|
|
"model_cooldowns": active_cooldowns, |
|
|
} |
|
|
|
|
|
|
|
|
if cooldown_filter == "in_cooldown": |
|
|
|
|
|
if active_cooldowns: |
|
|
all_summaries.append(summary) |
|
|
elif cooldown_filter == "no_cooldown": |
|
|
|
|
|
if not active_cooldowns: |
|
|
all_summaries.append(summary) |
|
|
else: |
|
|
|
|
|
all_summaries.append(summary) |
|
|
|
|
|
|
|
|
total_count = len(all_summaries) |
|
|
if limit is not None: |
|
|
summaries = all_summaries[offset:offset + limit] |
|
|
else: |
|
|
summaries = all_summaries[offset:] |
|
|
|
|
|
return { |
|
|
"items": summaries, |
|
|
"total": total_count, |
|
|
"offset": offset, |
|
|
"limit": limit, |
|
|
"stats": global_stats, |
|
|
} |
|
|
|
|
|
except Exception as e: |
|
|
log.error(f"Error getting credentials summary: {e}") |
|
|
return { |
|
|
"items": [], |
|
|
"total": 0, |
|
|
"offset": offset, |
|
|
"limit": limit, |
|
|
"stats": {"total": 0, "normal": 0, "disabled": 0}, |
|
|
} |
|
|
|
|
|
async def get_duplicate_credentials_by_email(self, mode: str = "geminicli") -> Dict[str, Any]: |
|
|
""" |
|
|
获取按邮箱分组的重复凭证信息(只查询邮箱和文件名,不加载完整凭证数据) |
|
|
用于去重操作 |
|
|
|
|
|
Args: |
|
|
mode: 凭证模式 ("geminicli" 或 "antigravity") |
|
|
|
|
|
Returns: |
|
|
包含 email_groups(邮箱分组)、duplicate_count(重复数量)、no_email_count(无邮箱数量)的字典 |
|
|
""" |
|
|
self._ensure_initialized() |
|
|
|
|
|
try: |
|
|
|
|
|
table_name = self._get_table_name(mode) |
|
|
|
|
|
async with aiosqlite.connect(self._db_path) as db: |
|
|
|
|
|
query = f""" |
|
|
SELECT filename, user_email |
|
|
FROM {table_name} |
|
|
ORDER BY filename |
|
|
""" |
|
|
|
|
|
async with db.execute(query) as cursor: |
|
|
rows = await cursor.fetchall() |
|
|
|
|
|
|
|
|
email_to_files = {} |
|
|
no_email_files = [] |
|
|
|
|
|
for filename, user_email in rows: |
|
|
if user_email: |
|
|
if user_email not in email_to_files: |
|
|
email_to_files[user_email] = [] |
|
|
email_to_files[user_email].append(filename) |
|
|
else: |
|
|
no_email_files.append(filename) |
|
|
|
|
|
|
|
|
duplicate_groups = [] |
|
|
total_duplicate_count = 0 |
|
|
|
|
|
for email, files in email_to_files.items(): |
|
|
if len(files) > 1: |
|
|
|
|
|
duplicate_groups.append({ |
|
|
"email": email, |
|
|
"kept_file": files[0], |
|
|
"duplicate_files": files[1:], |
|
|
"duplicate_count": len(files) - 1, |
|
|
}) |
|
|
total_duplicate_count += len(files) - 1 |
|
|
|
|
|
return { |
|
|
"email_groups": email_to_files, |
|
|
"duplicate_groups": duplicate_groups, |
|
|
"duplicate_count": total_duplicate_count, |
|
|
"no_email_files": no_email_files, |
|
|
"no_email_count": len(no_email_files), |
|
|
"unique_email_count": len(email_to_files), |
|
|
"total_count": len(rows), |
|
|
} |
|
|
|
|
|
except Exception as e: |
|
|
log.error(f"Error getting duplicate credentials by email: {e}") |
|
|
return { |
|
|
"email_groups": {}, |
|
|
"duplicate_groups": [], |
|
|
"duplicate_count": 0, |
|
|
"no_email_files": [], |
|
|
"no_email_count": 0, |
|
|
"unique_email_count": 0, |
|
|
"total_count": 0, |
|
|
} |
|
|
|
|
|
|
|
|
|
|
|
async def set_config(self, key: str, value: Any) -> bool: |
|
|
"""设置配置(写入数据库 + 更新内存缓存)""" |
|
|
self._ensure_initialized() |
|
|
|
|
|
try: |
|
|
async with aiosqlite.connect(self._db_path) as db: |
|
|
await db.execute(""" |
|
|
INSERT INTO config (key, value, updated_at) |
|
|
VALUES (?, ?, unixepoch()) |
|
|
ON CONFLICT(key) DO UPDATE SET |
|
|
value = excluded.value, |
|
|
updated_at = excluded.updated_at |
|
|
""", (key, json.dumps(value))) |
|
|
await db.commit() |
|
|
|
|
|
|
|
|
self._config_cache[key] = value |
|
|
return True |
|
|
|
|
|
except Exception as e: |
|
|
log.error(f"Error setting config {key}: {e}") |
|
|
return False |
|
|
|
|
|
async def reload_config_cache(self): |
|
|
"""重新加载配置缓存(在批量修改配置后调用)""" |
|
|
self._ensure_initialized() |
|
|
self._config_loaded = False |
|
|
await self._load_config_cache() |
|
|
log.info("Config cache reloaded from database") |
|
|
|
|
|
async def get_config(self, key: str, default: Any = None) -> Any: |
|
|
"""获取配置(从内存缓存)""" |
|
|
self._ensure_initialized() |
|
|
return self._config_cache.get(key, default) |
|
|
|
|
|
async def get_all_config(self) -> Dict[str, Any]: |
|
|
"""获取所有配置(从内存缓存)""" |
|
|
self._ensure_initialized() |
|
|
return self._config_cache.copy() |
|
|
|
|
|
async def delete_config(self, key: str) -> bool: |
|
|
"""删除配置""" |
|
|
self._ensure_initialized() |
|
|
|
|
|
try: |
|
|
async with aiosqlite.connect(self._db_path) as db: |
|
|
await db.execute("DELETE FROM config WHERE key = ?", (key,)) |
|
|
await db.commit() |
|
|
|
|
|
|
|
|
self._config_cache.pop(key, None) |
|
|
return True |
|
|
|
|
|
except Exception as e: |
|
|
log.error(f"Error deleting config {key}: {e}") |
|
|
return False |
|
|
|
|
|
|
|
|
|
|
|
async def set_model_cooldown( |
|
|
self, |
|
|
filename: str, |
|
|
model_key: str, |
|
|
cooldown_until: Optional[float], |
|
|
mode: str = "geminicli" |
|
|
) -> bool: |
|
|
""" |
|
|
设置特定模型的冷却时间 |
|
|
|
|
|
Args: |
|
|
filename: 凭证文件名 |
|
|
model_key: 模型键(antigravity 用模型名,gcli 用 pro/flash) |
|
|
cooldown_until: 冷却截止时间戳(None 表示清除冷却) |
|
|
mode: 凭证模式 ("geminicli" 或 "antigravity") |
|
|
|
|
|
Returns: |
|
|
是否成功 |
|
|
""" |
|
|
self._ensure_initialized() |
|
|
|
|
|
try: |
|
|
table_name = self._get_table_name(mode) |
|
|
async with aiosqlite.connect(self._db_path) as db: |
|
|
|
|
|
async with db.execute(f""" |
|
|
SELECT model_cooldowns FROM {table_name} WHERE filename = ? |
|
|
""", (filename,)) as cursor: |
|
|
row = await cursor.fetchone() |
|
|
|
|
|
if not row: |
|
|
log.warning(f"Credential {filename} not found") |
|
|
return False |
|
|
|
|
|
model_cooldowns = json.loads(row[0] or '{}') |
|
|
|
|
|
|
|
|
if cooldown_until is None: |
|
|
model_cooldowns.pop(model_key, None) |
|
|
else: |
|
|
model_cooldowns[model_key] = cooldown_until |
|
|
|
|
|
|
|
|
await db.execute(f""" |
|
|
UPDATE {table_name} |
|
|
SET model_cooldowns = ?, |
|
|
updated_at = unixepoch() |
|
|
WHERE filename = ? |
|
|
""", (json.dumps(model_cooldowns), filename)) |
|
|
await db.commit() |
|
|
|
|
|
log.debug(f"Set model cooldown: {filename}, model_key={model_key}, cooldown_until={cooldown_until}") |
|
|
return True |
|
|
|
|
|
except Exception as e: |
|
|
log.error(f"Error setting model cooldown for {filename}: {e}") |
|
|
return False |
|
|
|