"""SQLite 数据库连接与初始化(分库) - assets.db :Provider / Asset / VPSDetail / DomainDetail / AIAccount / Account - metrics.db:MetricPoint / ServerInfo / SecurityCheck / EventLog """ import logging from pathlib import Path from typing import Generator from sqlmodel import Session, SQLModel, create_engine logger = logging.getLogger("vps-manager.db") BASE_DIR = Path(__file__).resolve().parent.parent DATA_DIR = BASE_DIR / "data" DATA_DIR.mkdir(parents=True, exist_ok=True) ASSETS_DB_URL = f"sqlite:///{DATA_DIR / 'assets.db'}" METRICS_DB_URL = f"sqlite:///{DATA_DIR / 'metrics.db'}" # timeout:sqlite3 内置 busy 等待(秒),随连接创建生效。 # SQLite 默认使用 NullPool(每请求新建连接),PRAGMA busy_timeout 只在单连接有效, # 必须通过 connect_args 传递,否则并发写入(Agent 上报 + 清理任务)会报 database is locked。 assets_engine = create_engine( ASSETS_DB_URL, echo=False, connect_args={"check_same_thread": False, "timeout": 30}, ) metrics_engine = create_engine( METRICS_DB_URL, echo=False, connect_args={"check_same_thread": False, "timeout": 30}, ) def _enable_wal(engine) -> None: """启用 WAL 模式,提升 SQLite 并发读写能力(journal_mode 持久化到库文件,设置一次即可)""" import sqlalchemy as sa with engine.connect() as conn: conn.execute(sa.text("PRAGMA journal_mode=WAL")) _enable_wal(assets_engine) _enable_wal(metrics_engine) # 兼容旧代码:默认 engine 指向资产库 engine = assets_engine def init_db() -> None: """分库建表(幂等,可重复调用)""" from app.models.asset import Account, AIAccount, Asset, CloudflareDetail, DomainDetail, VPSDetail # noqa: F401 from app.models.credential import Credential # noqa: F401 from app.models.monitor import ( # noqa: F401 EventLog, MetricPoint, SecurityCheck, ServerInfo, ) from app.models.provider import Provider # noqa: F401 from app.models.ssl import SiteCert, Subdomain # noqa: F401 asset_models = [Provider, Asset, VPSDetail, DomainDetail, AIAccount, CloudflareDetail, Subdomain, SiteCert, Credential, Account] metric_models = [MetricPoint, ServerInfo, SecurityCheck, EventLog] SQLModel.metadata.create_all( assets_engine, tables=[m.__table__ for m in asset_models] ) SQLModel.metadata.create_all( metrics_engine, tables=[m.__table__ for m in metric_models] ) _migrate_assets_db() _migrate_indexes() def _migrate_assets_db() -> None: """轻量迁移:为已有表补充新增列(SQLite create_all 不会修改已有表结构)""" import sqlalchemy as sa with assets_engine.begin() as conn: insp = sa.inspect(conn) if insp.has_table("assets"): cols = {c["name"] for c in insp.get_columns("assets")} if "external_id" not in cols: conn.execute(sa.text("ALTER TABLE assets ADD COLUMN external_id VARCHAR")) if "renew_url" not in cols: conn.execute(sa.text("ALTER TABLE assets ADD COLUMN renew_url VARCHAR")) if "cancel_url" not in cols: conn.execute(sa.text("ALTER TABLE assets ADD COLUMN cancel_url VARCHAR")) if "account_id" not in cols: conn.execute(sa.text("ALTER TABLE assets ADD COLUMN account_id INTEGER")) # 资产账号引用从 name 字符串迁移到 account_id 外键(幂等) _backfill_asset_account_id(conn) if insp.has_table("providers"): cols = {c["name"] for c in insp.get_columns("providers")} if "last_synced_at" not in cols: conn.execute(sa.text("ALTER TABLE providers ADD COLUMN last_synced_at DATETIME")) if "services" not in cols: conn.execute(sa.text("ALTER TABLE providers ADD COLUMN services VARCHAR DEFAULT ''")) # 预设平台的 sdk_type / services 对齐当前代码能力(幂等,用户自建行不动) _align_preset_providers(conn) if insp.has_table("ai_accounts"): cols = {c["name"] for c in insp.get_columns("ai_accounts")} if "api_key_encrypted" not in cols: conn.execute(sa.text("ALTER TABLE ai_accounts ADD COLUMN api_key_encrypted VARCHAR")) # 迁移:将明文 api_key 加密后存入 api_key_encrypted,并清空原字段 _migrate_plaintext_api_keys(conn) if insp.has_table("accounts"): cols = {c["name"] for c in insp.get_columns("accounts")} if "login_user" not in cols: conn.execute(sa.text("ALTER TABLE accounts ADD COLUMN login_user VARCHAR")) if "login_password_encrypted" not in cols: conn.execute(sa.text("ALTER TABLE accounts ADD COLUMN login_password_encrypted VARCHAR")) if "api_config_encrypted" not in cols: conn.execute(sa.text("ALTER TABLE accounts ADD COLUMN api_config_encrypted VARCHAR")) if "last_synced_at" not in cols: conn.execute(sa.text("ALTER TABLE accounts ADD COLUMN last_synced_at DATETIME")) if "credential_id" not in cols: conn.execute(sa.text("ALTER TABLE accounts ADD COLUMN credential_id INTEGER")) # 凭证下沉:平台 api_config 迁到默认账号;AI 资产自动挂靠账号 _backfill_account_credentials(conn) # 合并存量 (platform,name) 重复账号,为建唯一索引做准备 _dedupe_accounts(conn) # 凭据中心化:账号登录密码搬入 credentials(唯一事实源) _migrate_account_credentials_to_vault(conn) if insp.has_table("credentials"): cols = {c["name"] for c in insp.get_columns("credentials")} if "otp_verified_at" not in cols: conn.execute( sa.text("ALTER TABLE credentials ADD COLUMN otp_verified_at DATETIME") ) # 旧版绑定强制校验过动态码,回填为「已验证」,避免升级后全量误报 conn.execute( sa.text( "UPDATE credentials SET otp_verified_at = COALESCE(updated_at, created_at) " "WHERE otp_secret_encrypted IS NOT NULL AND otp_secret_encrypted != ''" ) ) def _migrate_indexes() -> None: """为已有数据库补充复合索引(create_all 不会为已存在的表补索引,IF NOT EXISTS 幂等)""" import sqlalchemy as sa # 旧的全局 name 唯一索引(早期 Account.name unique=True 产物)会阻止跨平台同名, # 与新的 (platform,name) 联合唯一冲突,需先删除 with assets_engine.begin() as conn: conn.execute(sa.text("DROP INDEX IF EXISTS ix_accounts_name")) stmts = [ (assets_engine, "CREATE INDEX IF NOT EXISTS ix_assets_provider_ext_type ON assets (provider_id, external_id, asset_type)"), (assets_engine, "CREATE INDEX IF NOT EXISTS ix_assets_account_id ON assets (account_id)"), # (platform, name) 联合唯一:同邮箱可跨平台复用,同平台内不重名 (assets_engine, "CREATE UNIQUE INDEX IF NOT EXISTS uq_accounts_platform_name ON accounts (platform, name)"), # 已有 accounts 表不会因模型加 index=True 自动补索引,显式创建 (assets_engine, "CREATE INDEX IF NOT EXISTS ix_accounts_credential_id ON accounts (credential_id)"), (metrics_engine, "CREATE INDEX IF NOT EXISTS ix_metric_points_asset_ts ON metric_points (asset_id, ts)"), (metrics_engine, "CREATE INDEX IF NOT EXISTS ix_security_checks_asset_ts ON security_checks (asset_id, ts)"), ] for eng, sql in stmts: with eng.begin() as conn: conn.execute(sa.text(sql)) def _migrate_plaintext_api_keys(conn) -> None: """一次性迁移:将 ai_accounts 中残留的明文 api_key 加密后存入 api_key_encrypted,并清空原字段""" import sqlalchemy as sa from app.core.crypto import encrypt rows = conn.execute( sa.text("SELECT id, api_key FROM ai_accounts WHERE api_key IS NOT NULL AND api_key != ''") ).fetchall() if not rows: return for row in rows: encrypted = encrypt(row[1]) if encrypted: conn.execute( sa.text("UPDATE ai_accounts SET api_key_encrypted = :enc, api_key = NULL WHERE id = :id"), {"enc": encrypted, "id": row[0]}, ) def _align_preset_providers(conn) -> None: """把预设平台的 sdk_type / services 对齐到当前代码能力(幂等) seed_providers 只插入缺失 slug、不更新已有行,所以代码侧的能力变化 (新增适配器、废弃从未实现的 sdk_type)必须在这里落地,否则线上会长期 停留在旧声明上。 - sdk_type:库里声明的值不在适配器注册表内就清空。留着会让平台卡片显示 「API 对接: xxx」,用户建好账号点同步才撞「暂不支持的 SDK 类型」—— 声明支持却不支持,比明说不支持更糟。 - services:只增量补标签,不删除。用户可能手动加过标签(线上 aliyun-intl 就多了 ai_agent/other),按种子值整体覆盖会丢用户配置。 只处理 slug 命中 PRESET_PROVIDERS 的行,用户自建平台(如 tothost)不动。 """ import sqlalchemy as sa from app.adapters import registry from app.core.seed import PRESET_PROVIDERS supported = set(registry.supported_types()) preset = {d["slug"]: d for d in PRESET_PROVIDERS} rows = conn.execute( sa.text("SELECT id, slug, sdk_type, services FROM providers") ).fetchall() for row in rows: data = preset.get(row[1]) if not data: continue declared = row[2] if declared and declared not in supported: conn.execute( sa.text("UPDATE providers SET sdk_type = NULL WHERE id = :id"), {"id": row[0]}, ) seed_services = [s.strip() for s in (data.get("services") or "").split(",") if s.strip()] if not seed_services: continue current = [s.strip() for s in (row[3] or "").split(",") if s.strip()] missing = [s for s in seed_services if s not in current] if missing: conn.execute( sa.text("UPDATE providers SET services = :services WHERE id = :id"), {"services": ",".join(current + missing), "id": row[0]}, ) def _dedupe_accounts(conn) -> None: """建 (platform,name) 唯一索引前的防护:合并存量重复账号。 platform 统一按 COALESCE(platform,'') 规范化比较(避免 NULL 绕过)。 重复时保留 id 最小者,将其余行的资产引用(account_id/account)与凭证并入后删除。 当前数据量小,预期无重复;此函数仅在发现重复时产生写操作。 """ import sqlalchemy as sa rows = conn.execute( sa.text( "SELECT COALESCE(platform,'') AS p, name, COUNT(*) AS c, MIN(id) AS keep_id " "FROM accounts GROUP BY p, name HAVING c > 1" ) ).fetchall() for p, name, _c, keep_id in rows: dups = conn.execute( sa.text( "SELECT id FROM accounts WHERE COALESCE(platform,'') = :p AND name = :n AND id != :keep" ), {"p": p, "n": name, "keep": keep_id}, ).fetchall() for (dup_id,) in dups: # 资产引用并入保留行 conn.execute( sa.text("UPDATE assets SET account_id = :keep WHERE account_id = :dup"), {"keep": keep_id, "dup": dup_id}, ) # 凭证:保留行为空时才从重复行拷贝 conn.execute( sa.text( "UPDATE accounts SET " "api_config_encrypted = COALESCE(api_config_encrypted, (SELECT api_config_encrypted FROM accounts WHERE id = :dup)), " "login_password_encrypted = COALESCE(login_password_encrypted, (SELECT login_password_encrypted FROM accounts WHERE id = :dup)), " "login_user = COALESCE(login_user, (SELECT login_user FROM accounts WHERE id = :dup)) " "WHERE id = :keep" ), {"keep": keep_id, "dup": dup_id}, ) conn.execute(sa.text("DELETE FROM accounts WHERE id = :dup"), {"dup": dup_id}) def _migrate_account_credentials_to_vault(conn) -> None: """凭据中心化迁移(幂等):把账号登录密码搬入 credentials 表。 规则(spec §5): - 仅迁移 login_password_encrypted 非空且 credential_id 为空的账号(重跑不重复建) - site = providers.name(platform=slug 匹配)→ 否则 platform 原文 → 否则 '未分类' - username = COALESCE(login_user, name) - login_type 写入 enum 成员名 'PASSWORD'(SQLAlchemy Enum 列存名不存值, 与现有 assets.asset_type='VPS' 等形态一致) - 密文原样搬入(不解密再加密,避免中间态明文暴露) - 搬入成功后置空 accounts.login_password_encrypted(唯一事实源,防双份漂移) """ import sqlalchemy as sa from datetime import datetime, timezone rows = conn.execute( sa.text( "SELECT a.id, a.name, a.login_user, a.login_password_encrypted, " "COALESCE(a.platform, ''), COALESCE(p.name, '') " "FROM accounts a LEFT JOIN providers p ON p.slug = COALESCE(a.platform, '') " "WHERE a.login_password_encrypted IS NOT NULL AND a.login_password_encrypted != '' " "AND a.credential_id IS NULL" ) ).fetchall() for acc_id, name, login_user, pwd_enc, platform, provider_name in rows: site = provider_name or platform or "未分类" username = login_user or name now = datetime.now(timezone.utc).replace(tzinfo=None) conn.execute( sa.text( "INSERT INTO credentials (site, username, login_type, password_encrypted, " "created_at, updated_at) VALUES (:site, :username, 'PASSWORD', :pwd, :now, :now)" ), {"site": site, "username": username, "pwd": pwd_enc, "now": now}, ) cred_id = conn.execute(sa.text("SELECT last_insert_rowid()")).scalar() conn.execute( sa.text( "UPDATE accounts SET credential_id = :cid, login_password_encrypted = NULL " "WHERE id = :id" ), {"cid": cred_id, "id": acc_id}, ) def _backfill_asset_account_id(conn) -> None: """资产账号引用迁移(幂等):把 Asset.account 的 name 字符串引用回填为 account_id 外键。 匹配规则:按 accounts.name 匹配;同名多个时优先 platform 与 asset.provider 一致者。 匹配不到(历史自由文本/未登记账号)则跳过并保留 account 字符串,供人工处理。 """ import sqlalchemy as sa rows = conn.execute( sa.text( "SELECT id, account, provider FROM assets " "WHERE account_id IS NULL AND account IS NOT NULL AND account != ''" ) ).fetchall() if not rows: return for asset_id, acc_name, provider in rows: # 同名账号可能多个:优先 platform 与资产 provider 一致的 cand = conn.execute( sa.text( "SELECT id, COALESCE(platform,'') FROM accounts WHERE name = :n " "ORDER BY (COALESCE(platform,'') = :prov) DESC, id ASC" ), {"n": acc_name, "prov": provider or ""}, ).fetchall() if cand: conn.execute( sa.text("UPDATE assets SET account_id = :aid WHERE id = :id"), {"aid": cand[0][0], "id": asset_id}, ) def _backfill_account_credentials(conn) -> None: """凭证下沉迁移(幂等,仅在首次加列后产生效果): 1. 平台凭证下沉:api_config_encrypted 非空的 Provider → 确保存在 「{slug}-默认」账号并拷入凭证(不删平台原值,保留可回滚)。 2. AI 资产挂靠:account 为空的 ai_agent 资产 → 按 ai_accounts.provider 找/建账号并关联;账号无 API 配置时把该资产的 api_key 写入账号配置。 """ import json from datetime import datetime, timezone import sqlalchemy as sa from app.core.crypto import decrypt, encrypt def _ensure_account(name: str, platform: str) -> int: # 按 (platform, name) 查,避免同名跨平台账号混淆 row = conn.execute( sa.text("SELECT id FROM accounts WHERE name = :name AND COALESCE(platform,'') = :platform"), {"name": name, "platform": platform}, ).first() if row: return row[0] conn.execute( sa.text("INSERT INTO accounts (name, platform, created_at) VALUES (:name, :platform, :ts)"), {"name": name, "platform": platform, "ts": datetime.now(timezone.utc).replace(tzinfo=None)}, ) return conn.execute( sa.text("SELECT id FROM accounts WHERE name = :name AND COALESCE(platform,'') = :platform"), {"name": name, "platform": platform}, ).first()[0] # 1. 平台凭证下沉到默认账号 for pid, slug, cfg in conn.execute( sa.text( "SELECT id, slug, api_config_encrypted FROM providers " "WHERE api_config_encrypted IS NOT NULL AND api_config_encrypted != ''" ) ).fetchall(): acc_name = f"{slug}-默认" _ensure_account(acc_name, slug) conn.execute( sa.text( "UPDATE accounts SET api_config_encrypted = :cfg " "WHERE name = :name AND COALESCE(platform,'') = :platform " "AND (api_config_encrypted IS NULL OR api_config_encrypted = '')" ), {"cfg": cfg, "name": acc_name, "platform": slug}, ) # 2. 无账号的 AI 资产按 provider 挂靠(account_id 外键 + account 字符串兼容) # 注意:SQLAlchemy Enum 列存的是成员名('AI_AGENT')而不是枚举值('ai_agent'), # 写值会让这个分支静默零命中(同文件 _migrate_account_credentials_to_vault 用的 'PASSWORD' 即成员名) rows = conn.execute( sa.text( "SELECT a.id, COALESCE(ai.provider, '') FROM assets a " "JOIN ai_accounts ai ON ai.asset_id = a.id " "WHERE a.asset_type = 'AI_AGENT' AND a.account_id IS NULL AND (a.account IS NULL OR a.account = '')" ) ).fetchall() attached = 0 for asset_id, provider in rows: if not provider: continue # 沿用步骤 1 的「{slug}-默认」命名:同平台已有账号时直接复用。 # 用裸 provider 名会建出「deepseek」与「deepseek-默认」两个语义重复的账号, # 而新建那个没凭证,在 UI 上点同步/测试会直接报「缺少凭证配置」 acc_id = _ensure_account(f"{provider}-默认", provider) conn.execute( sa.text("UPDATE assets SET account = :acc, account_id = :aid WHERE id = :id"), {"acc": provider, "aid": acc_id, "id": asset_id}, ) attached += 1 # 账号尚无 API 配置时,把该资产的 api_key 写入账号配置(解密后重组 JSON 再加密) key_row = conn.execute( sa.text("SELECT api_key_encrypted FROM ai_accounts WHERE asset_id = :id"), {"id": asset_id}, ).first() if key_row and key_row[0]: api_key = decrypt(key_row[0]) if api_key: cfg_json = json.dumps({"api_key": api_key}) cfg_enc = encrypt(cfg_json) if cfg_enc: conn.execute( sa.text( "UPDATE accounts SET api_config_encrypted = :cfg " "WHERE id = :aid " "AND (api_config_encrypted IS NULL OR api_config_encrypted = '')" ), {"cfg": cfg_enc, "aid": acc_id}, ) if attached: # 这段迁移每次启动都会重跑(幂等),有实际改动时必须留痕迹便于事后核对 logger.info("凭证下沉迁移:本次将 %s 个 AI 资产挂靠到账号", attached) def get_session() -> Generator[Session, None, None]: """资产库会话(默认) expire_on_commit=False:commit 后不失效对象属性,避免后续访问触发隐式 重新加载查询;需要最新值的场景由调用方显式 session.refresh()。 """ with Session(assets_engine, expire_on_commit=False) as session: yield session def get_metrics_session() -> Generator[Session, None, None]: """监控 / 日志库会话""" with Session(metrics_engine, expire_on_commit=False) as session: yield session