258 lines
11 KiB
Python
258 lines
11 KiB
Python
"""SQLite 数据库连接与初始化（分库）
|
||
|
||
- assets.db ：Provider / Asset / VPSDetail / DomainDetail / AIAccount / Account
|
||
- metrics.db：MetricPoint / ServerInfo / SecurityCheck / EventLog
|
||
"""
|
||
|
||
from pathlib import Path
|
||
from typing import Generator
|
||
|
||
from sqlmodel import Session, SQLModel, create_engine
|
||
|
||
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.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, 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 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 ''"))
|
||
# 为已有预设平台补充 services（仅补空值，用户自定义行不动）
|
||
_backfill_provider_services(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"))
|
||
# 凭证下沉：平台 api_config 迁到默认账号；AI 资产自动挂靠账号
|
||
_backfill_account_credentials(conn)
|
||
|
||
|
||
def _migrate_indexes() -> None:
|
||
"""为已有数据库补充复合索引（create_all 不会为已存在的表补索引，IF NOT EXISTS 幂等）"""
|
||
import sqlalchemy as sa
|
||
|
||
stmts = [
|
||
(assets_engine, "CREATE INDEX IF NOT EXISTS ix_assets_provider_ext_type ON assets (provider_id, external_id, asset_type)"),
|
||
(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 _backfill_provider_services(conn) -> None:
|
||
"""为已有预设平台补充 services 服务列表（只更新 services 为空的预设 slug）"""
|
||
import sqlalchemy as sa
|
||
from app.core.seed import PRESET_PROVIDERS
|
||
|
||
for data in PRESET_PROVIDERS:
|
||
services = data.get("services")
|
||
if not services:
|
||
continue
|
||
conn.execute(
|
||
sa.text(
|
||
"UPDATE providers SET services = :services "
|
||
"WHERE slug = :slug AND (services IS NULL OR services = '')"
|
||
),
|
||
{"services": services, "slug": data["slug"]},
|
||
)
|
||
|
||
|
||
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:
|
||
row = conn.execute(sa.text("SELECT id FROM accounts WHERE name = :name"), {"name": name}).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"), {"name": name}).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 (api_config_encrypted IS NULL OR api_config_encrypted = '')"
|
||
),
|
||
{"cfg": cfg, "name": acc_name},
|
||
)
|
||
|
||
# 2. 无账号的 AI 资产按 provider 挂靠
|
||
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 IS NULL OR a.account = '')"
|
||
)
|
||
).fetchall()
|
||
for asset_id, provider in rows:
|
||
if not provider:
|
||
continue
|
||
_ensure_account(provider, provider)
|
||
conn.execute(
|
||
sa.text("UPDATE assets SET account = :acc WHERE id = :id"),
|
||
{"acc": provider, "id": asset_id},
|
||
)
|
||
# 账号尚无 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 name = :name AND (api_config_encrypted IS NULL OR api_config_encrypted = '')"
|
||
),
|
||
{"cfg": cfg_enc, "name": provider},
|
||
)
|
||
|
||
|
||
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
|