259 lines
9.8 KiB
Python
259 lines
9.8 KiB
Python
"""数据库连接在运行时初始化（与 EasySub 相同的做法）。
|
||
|
||
来源优先级：环境变量 DATABASE_URL > 网页安装向导保存的 data/db_config.json。
|
||
都没有时进入安装向导模式：流量照常采集到内存，数据库配置好后才落库。
|
||
"""
|
||
import threading
|
||
import time
|
||
|
||
from fastapi import HTTPException
|
||
from sqlalchemy import create_engine, event, inspect
|
||
from sqlalchemy.dialects.mysql import insert as mysql_insert
|
||
from sqlalchemy.dialects.sqlite import insert as sqlite_insert
|
||
from sqlalchemy.orm import DeclarativeBase, sessionmaker
|
||
|
||
|
||
class Base(DeclarativeBase):
|
||
pass
|
||
|
||
|
||
# 由 init_engine() 赋值；未配置前为 None
|
||
engine = None
|
||
SessionLocal = None
|
||
|
||
_ready_hooks: list = []
|
||
_connect_lock = threading.Lock()
|
||
_last_reconnect_attempt = 0.0
|
||
_RECONNECT_INTERVAL = 15.0
|
||
|
||
|
||
class NotConfigured(RuntimeError):
|
||
pass
|
||
|
||
|
||
def _make_engine(url: str):
|
||
if url.startswith("sqlite"):
|
||
import os
|
||
path = url.split("///", 1)[-1]
|
||
if path and path != ":memory:":
|
||
os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)
|
||
# 锁等待 30 秒（默认 5 秒）：历史重算持有写锁时，每分钟落库和 API 不至于直接报 database is locked
|
||
eng = create_engine(url, connect_args={"check_same_thread": False, "timeout": 30})
|
||
|
||
@event.listens_for(eng, "connect")
|
||
def _pragma(conn, _):
|
||
cur = conn.cursor()
|
||
cur.execute("PRAGMA journal_mode=WAL")
|
||
cur.execute("PRAGMA synchronous=NORMAL")
|
||
cur.close()
|
||
|
||
return eng
|
||
return create_engine(url, pool_pre_ping=True, pool_recycle=1800, pool_size=10, max_overflow=10,
|
||
connect_args={"connect_timeout": 10} if url.startswith("mysql") else {})
|
||
|
||
|
||
def init_engine(url: str):
|
||
global engine, SessionLocal
|
||
reset_engine()
|
||
engine = _make_engine(url)
|
||
SessionLocal = sessionmaker(bind=engine, autoflush=False, expire_on_commit=False)
|
||
return engine
|
||
|
||
|
||
def reset_engine():
|
||
global engine, SessionLocal
|
||
if engine is not None:
|
||
engine.dispose()
|
||
engine = None
|
||
SessionLocal = None
|
||
|
||
|
||
def is_configured() -> bool:
|
||
return SessionLocal is not None
|
||
|
||
|
||
def session():
|
||
"""后台任务用：未配置数据库时抛 NotConfigured。"""
|
||
if SessionLocal is None:
|
||
raise NotConfigured("数据库尚未配置")
|
||
return SessionLocal()
|
||
|
||
|
||
def is_mysql() -> bool:
|
||
return engine is not None and engine.dialect.name == "mysql"
|
||
|
||
|
||
_str_lens: dict = {}
|
||
|
||
|
||
def fit(table, rows: list[dict]) -> list[dict]:
|
||
"""按列定义截断字符串：网络来的主机名、域名、规则名可能超长，MySQL 严格模式下一行超长会让整批写入失败
|
||
（0.6.8 的 VARCHAR(45)、全站审计的 DHCP 主机名 >128）。SQLite 不校验长度，测试里发现不了"""
|
||
lens = _str_lens.get(table.name)
|
||
if lens is None:
|
||
lens = _str_lens[table.name] = {c.name: c.type.length for c in table.columns
|
||
if getattr(c.type, "length", None)}
|
||
if not lens:
|
||
return rows
|
||
out = []
|
||
for r in rows:
|
||
bad = [k for k, n in lens.items() if isinstance(r.get(k), str) and len(r[k]) > n]
|
||
out.append({**r, **{k: r[k][:lens[k]] for k in bad}} if bad else r)
|
||
return out
|
||
|
||
|
||
def upsert_add(session, table, rows: list[dict], keys: list[str], add_cols: list[str], set_cols: list[str] = ()):
|
||
"""批量 upsert：主键冲突时 add_cols 累加、set_cols 覆盖。MySQL / SQLite 通用。字符串按列长截断"""
|
||
if not rows:
|
||
return
|
||
rows = fit(table, rows)
|
||
if is_mysql():
|
||
stmt = mysql_insert(table).values(rows)
|
||
upd = {c: table.c[c] + stmt.inserted[c] for c in add_cols}
|
||
upd.update({c: stmt.inserted[c] for c in set_cols})
|
||
stmt = stmt.on_duplicate_key_update(**upd)
|
||
else:
|
||
stmt = sqlite_insert(table).values(rows)
|
||
upd = {c: table.c[c] + stmt.excluded[c] for c in add_cols}
|
||
upd.update({c: stmt.excluded[c] for c in set_cols})
|
||
stmt = stmt.on_conflict_do_update(index_elements=keys, set_=upd)
|
||
session.execute(stmt)
|
||
|
||
|
||
def _create_database_if_missing(err: Exception) -> bool:
|
||
"""MySQL 报 1049（库不存在）时，用同一账号连到服务器建库。成功返回 True。"""
|
||
if not is_mysql() or "1049" not in str(err):
|
||
return False
|
||
url = engine.url
|
||
name = url.database
|
||
# 注意：url.set(database=None) 在 SQLAlchemy 2.0 里表示"不修改"，必须用 _replace 真正去掉库名
|
||
server = create_engine(url._replace(database=None), pool_pre_ping=True)
|
||
try:
|
||
with server.connect() as conn:
|
||
conn.exec_driver_sql(
|
||
f"CREATE DATABASE IF NOT EXISTS `{name}` CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci")
|
||
print(f"[db] 已自动创建数据库 {name}", flush=True)
|
||
return True
|
||
finally:
|
||
server.dispose()
|
||
|
||
|
||
def init_db(retries: int = 20, delay: float = 3.0):
|
||
"""建库建表。容器启动时 MySQL 往往还没就绪，这里重试。"""
|
||
from app import models # noqa: F401 注册模型
|
||
|
||
for i in range(1, retries + 1):
|
||
try:
|
||
Base.metadata.create_all(engine)
|
||
add_missing_columns()
|
||
return
|
||
except Exception as e: # noqa: BLE001
|
||
try:
|
||
if _create_database_if_missing(e):
|
||
Base.metadata.create_all(engine)
|
||
add_missing_columns()
|
||
return
|
||
except Exception as ce: # noqa: BLE001
|
||
print(f"[db] 数据库不存在，且自动创建失败（账号可能没有建库权限）：{ce}", flush=True)
|
||
print(f"[db] 数据库未就绪（{i}/{retries}）：{type(e).__name__}: {str(e)[:160]}", flush=True)
|
||
if i == retries:
|
||
raise
|
||
time.sleep(delay)
|
||
|
||
|
||
def add_missing_columns():
|
||
"""升级时补上新版本新增的列（create_all 只建新表，不改已有的表）。新增列一律可为空，不需要回填。
|
||
同时把代码里加宽过的列在存量库上加宽（0.6.8：access_log.remote 45→255；只加宽不缩窄，SQLite 不校验长度跳过）。"""
|
||
insp = inspect(engine)
|
||
with engine.begin() as conn:
|
||
for table in Base.metadata.sorted_tables:
|
||
if not insp.has_table(table.name):
|
||
continue
|
||
have = {c["name"]: c for c in insp.get_columns(table.name)}
|
||
for col in table.columns:
|
||
if col.name not in have:
|
||
if col.nullable:
|
||
ctype = col.type.compile(dialect=engine.dialect)
|
||
conn.exec_driver_sql(f"ALTER TABLE `{table.name}` ADD COLUMN `{col.name}` {ctype} NULL")
|
||
print(f"[db] 已为 {table.name} 添加列 {col.name}", flush=True)
|
||
continue
|
||
length = getattr(col.type, "length", None)
|
||
# 注意：inspector 的列字典没有 "length" 键，长度在反射出来的类型对象上
|
||
db_len = getattr(have[col.name].get("type"), "length", None)
|
||
if insp.dialect.name == "mysql" and length and db_len and length > int(db_len):
|
||
ctype = col.type.compile(dialect=engine.dialect)
|
||
conn.exec_driver_sql(f"ALTER TABLE `{table.name}` MODIFY `{col.name}` {ctype}")
|
||
print(f"[db] 已把 {table.name}.{col.name} 从 {db_len} 加宽到 {length}", flush=True)
|
||
# 新版本加的索引：create_all 不会给已有的表补，这里补上
|
||
have_idx = {i["name"] for i in insp.get_indexes(table.name)}
|
||
for idx in table.indexes:
|
||
if idx.name and idx.name not in have_idx:
|
||
idx.create(conn)
|
||
print(f"[db] 已为 {table.name} 添加索引 {idx.name}", flush=True)
|
||
|
||
|
||
def add_ready_hook(fn):
|
||
"""数据库连接成功后要执行的初始化（创建管理员、恢复状态等）。"""
|
||
if fn not in _ready_hooks:
|
||
_ready_hooks.append(fn)
|
||
|
||
|
||
def connect(url: str, retries: int = 1, delay: float = 3.0):
|
||
"""连接并建表，成功后执行 ready hooks；失败时恢复为未配置状态并抛出异常。"""
|
||
init_engine(url)
|
||
try:
|
||
init_db(retries=retries, delay=delay)
|
||
except Exception:
|
||
reset_engine()
|
||
raise
|
||
for fn in _ready_hooks:
|
||
fn()
|
||
|
||
|
||
def configured_url() -> str | None:
|
||
from app import bootstrap
|
||
from app.config import get_settings
|
||
|
||
if get_settings().database_url:
|
||
return get_settings().database_url
|
||
if bootstrap.config_exists():
|
||
return bootstrap.build_url(bootstrap.load_config())
|
||
return None
|
||
|
||
|
||
def connect_from_saved(force: bool = False, retries: int = 1) -> bool:
|
||
"""用环境变量或已保存的配置连接（幂等）。非 force 调用有节流，避免数据库故障时每个请求都阻塞。"""
|
||
global _last_reconnect_attempt
|
||
if SessionLocal is not None:
|
||
return True
|
||
now = time.monotonic()
|
||
if not force and now - _last_reconnect_attempt < _RECONNECT_INTERVAL:
|
||
return False
|
||
_last_reconnect_attempt = now
|
||
url = configured_url()
|
||
if not url:
|
||
return False
|
||
# 启动时的后台连接和定时重连可能同时触发，只让一个去连
|
||
if not _connect_lock.acquire(blocking=False):
|
||
return False
|
||
try:
|
||
if SessionLocal is not None:
|
||
return True
|
||
connect(url, retries=retries)
|
||
return True
|
||
except Exception as e: # noqa: BLE001
|
||
print(f"[db] 连接失败：{type(e).__name__}: {str(e)[:200]}", flush=True)
|
||
return False
|
||
finally:
|
||
_connect_lock.release()
|
||
|
||
|
||
def get_db():
|
||
if SessionLocal is None and not connect_from_saved():
|
||
raise HTTPException(503, "数据库尚未配置或暂时不可用；首次使用请在安装向导中填写 MySQL 信息")
|
||
db = SessionLocal()
|
||
try:
|
||
yield db
|
||
finally:
|
||
db.close()
|