87 lines
3.1 KiB
Python
87 lines
3.1 KiB
Python
"""SQLite 数据库连接与初始化（分库）
|
||
|
||
- assets.db ：Provider / Asset / VPSDetail / DomainDetail / AIAccount
|
||
- 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'}"
|
||
|
||
assets_engine = create_engine(
|
||
ASSETS_DB_URL, echo=False, connect_args={"check_same_thread": False}
|
||
)
|
||
metrics_engine = create_engine(
|
||
METRICS_DB_URL, echo=False, connect_args={"check_same_thread": False}
|
||
)
|
||
|
||
# 兼容旧代码：默认 engine 指向资产库
|
||
engine = assets_engine
|
||
|
||
|
||
def init_db() -> None:
|
||
"""分库建表（幂等，可重复调用）"""
|
||
from app.models.asset import 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
|
||
|
||
asset_models = [Provider, Asset, VPSDetail, DomainDetail, AIAccount, CloudflareDetail]
|
||
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()
|
||
|
||
|
||
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 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"))
|
||
|
||
|
||
def get_session() -> Generator[Session, None, None]:
|
||
"""资产库会话（默认）"""
|
||
with Session(assets_engine) as session:
|
||
yield session
|
||
|
||
|
||
def get_metrics_session() -> Generator[Session, None, None]:
|
||
"""监控 / 日志库会话"""
|
||
with Session(metrics_engine) as session:
|
||
yield session
|