butubb
f22263ab45
为把数据库从单文件 SQLite 迁到 MySQL 5.7 铺路。本期不改后端： DB_HOST 为空时仍走 SQLite，本地开发无感。 - config.py: 新增 DEPLOY_ENV（默认 dev）与 DB_HOST/PORT/USER/PASSWORD/NAME/ CHARSET/COLLATION、DATABASE_URL、两个逃生阀（DB_ALLOW_ENV_MISMATCH / DB_ALLOW_SQLITE_FALLBACK） - core/db_config.py（新增）: URI 组装；按方言分叉的引擎参数（utf8mb4、 pool_pre_ping、pool_recycle=1800、READ COMMITTED、STRICT_TRANS_TABLES）； 连接探活；app_meta 方言中立读写（MySQL 里 key 是保留字，需反引号） - 防混库三层: ①库名与环境绑定（dev→auto_control_dev / prod→auto_control） ②库标签 app_meta.deployment_env 与 .env 声明比对 ③启动横幅打印当前库 （生产用 WARNING 级）。不符直接拒绝启动并说明两边分别是什么 - web_server.py: 硬编码 sqlite URI → db_config；配置错在装配期就 exit 2； init_db 之后跑库标签校验 + 横幅 - core/models.py: PRAGMA 监听器加 sqlite 类型守卫——它挂在 Engine 基类上， MySQL 连接执行 PRAGMA 会直接导致建连失败 - requirements.txt 加 PyMySQL；scripts/start.sh 依赖守卫加 pymysql， 并在 exec 前打印 DEPLOY_ENV/DB_NAME/DB_HOST - .env.example 新增「数据库」段；DEVELOPMENT.md §4.1/4.2、DEPLOY.md §2.2 同步 验证: 用 DATABASE_URL 指向 users.db 的一致快照副本跑通主要只读接口 （health/devices/jobs/pool/groups/discovery/summary 全 200，app_meta 读写正常）； DEPLOY_ENV=prod 且无 DB_HOST 时退出码 2；开 DB_ALLOW_SQLITE_FALLBACK 后可回退。
270 lines
12 KiB
Python
270 lines
12 KiB
Python
"""数据库连接装配：URI 组装、引擎参数、环境防呆校验、启动横幅。
|
||
|
||
为什么单独一个模块：
|
||
web_server 要装配 app、迁移脚本/回归脚本要复用同一套 URI 与校验逻辑；
|
||
塞进 config.py 会把它撑肥，塞进 models.py 会引入循环依赖。
|
||
|
||
支持两种后端：
|
||
- **MySQL**（DB_HOST 非空）：正式用法，dev 库 / 正式库由 .env 的 DEPLOY_ENV 决定
|
||
- **SQLite**（DB_HOST 为空）：仅作迁移期回退；生产禁止静默回退，需 DB_ALLOW_SQLITE_FALLBACK=1
|
||
|
||
防呆设计（曾经把生产数据搞混过，这层是刚需）：
|
||
1. 库名与环境绑定：dev → auto_control_dev，prod → auto_control，对不上拒绝启动
|
||
2. 库标签：app_meta.deployment_env 记录"这个库属于哪个环境"，不匹配硬停
|
||
3. 启动横幅：每次启动都打印目标库，prod 用 WARNING 级（tail 日志必然可见）
|
||
"""
|
||
import os
|
||
import uuid
|
||
|
||
from sqlalchemy import create_engine, text
|
||
from sqlalchemy.engine import make_url
|
||
|
||
from config import (DATABASE_URL, DB_ALLOW_ENV_MISMATCH, DB_ALLOW_SQLITE_FALLBACK,
|
||
DB_CHARSET, DB_COLLATION, DB_HOST, DB_NAME, DB_PASSWORD,
|
||
DB_PORT, DB_USER, DATA_DIR, DEPLOY_ENV)
|
||
from core.logger import get_logger
|
||
|
||
_log = get_logger("core.db")
|
||
|
||
# 环境 ↔ 库名绑定：防止「dev 的 .env 抄错成正式库名」这类事故
|
||
_EXPECTED_DB_NAME = {"dev": "auto_control_dev", "prod": "auto_control"}
|
||
|
||
_REQUIRED_MYSQL_VERSION = "5.7"
|
||
|
||
# app_meta 里记录库身份的键
|
||
_K_ENV, _K_ID, _K_CLAIMED = "deployment_env", "deployment_id", "deployment_claimed_at"
|
||
|
||
|
||
class DBConfigError(Exception):
|
||
"""数据库配置错误（消息直接面向运维，要说清怎么改）。"""
|
||
|
||
|
||
# ================== URI 与引擎参数 ==================
|
||
def is_sqlite(uri):
|
||
return str(uri).startswith("sqlite")
|
||
|
||
|
||
def build_db_uri():
|
||
"""按 .env 组装连接串。DATABASE_URL 优先级最高（脚本临时指向别的库用）。"""
|
||
if DATABASE_URL:
|
||
return DATABASE_URL
|
||
if DB_HOST:
|
||
from urllib.parse import quote_plus
|
||
return "mysql+pymysql://{}:{}@{}:{}/{}?charset={}".format(
|
||
quote_plus(DB_USER), quote_plus(DB_PASSWORD),
|
||
DB_HOST, DB_PORT, DB_NAME, DB_CHARSET)
|
||
# 回退 SQLite：生产环境必须显式开逃生阀，否则拒绝启动（静默降级本身就是混库温床）
|
||
if DEPLOY_ENV == "prod" and not DB_ALLOW_SQLITE_FALLBACK:
|
||
raise DBConfigError(
|
||
"生产环境（DEPLOY_ENV=prod）未配置 DB_HOST，且未显式允许回退 SQLite。\n"
|
||
" 如确需临时回退到 sqlite，请在 .env 里加 DB_ALLOW_SQLITE_FALLBACK=1\n"
|
||
" 否则请补全 DB_HOST/DB_USER/DB_PASSWORD/DB_NAME。")
|
||
return "sqlite:///" + os.path.join(DATA_DIR, "users.db")
|
||
|
||
|
||
def engine_options(uri):
|
||
"""create_engine 参数。SQLite 与 MySQL 需求不同，必须分叉。"""
|
||
if is_sqlite(uri):
|
||
return {"pool_pre_ping": True}
|
||
return {
|
||
# 长驻后台线程（TaskManager/device_discovery）长时间空闲后，
|
||
# 连接可能已被 MySQL 的 wait_timeout 掐断 / 中间链路掉线
|
||
"pool_pre_ping": True,
|
||
"pool_recycle": 1800,
|
||
"pool_size": 10,
|
||
"max_overflow": 20,
|
||
"pool_timeout": 30,
|
||
# 接近 SQLite「逐语句读最新提交」的语义，并显著降低间隙锁/死锁概率。
|
||
# 备份导出需要一致性快照时会单独把那个连接提到 REPEATABLE READ。
|
||
"isolation_level": "READ COMMITTED",
|
||
"connect_args": {
|
||
"charset": DB_CHARSET,
|
||
# 保留 STRICT_TRANS_TABLES（超长写入报错而不是静默截断）；
|
||
# 去掉 5.7 默认带的 ONLY_FULL_GROUP_BY，避免给未来埋雷
|
||
"init_command": "SET sql_mode='STRICT_TRANS_TABLES,NO_ENGINE_SUBSTITUTION'",
|
||
},
|
||
}
|
||
|
||
|
||
def create_engine_from_uri(uri=None):
|
||
"""独立建一个 engine（不经过 Flask-SQLAlchemy）。脚本与启动前探活用。"""
|
||
uri = uri or build_db_uri()
|
||
return create_engine(uri, **engine_options(uri))
|
||
|
||
|
||
def describe_target(uri=None):
|
||
"""人类可读的目标描述（**不含密码**），用于日志与横幅。"""
|
||
uri = uri or build_db_uri()
|
||
try:
|
||
u = make_url(uri)
|
||
except Exception:
|
||
return str(uri)
|
||
if u.get_backend_name() == "sqlite":
|
||
return u.database
|
||
return "{}@{}:{}/{}".format(u.username or "", u.host or "", u.port or 3306,
|
||
u.database or "")
|
||
|
||
|
||
# ================== app_meta 读写（方言中立）==================
|
||
def _q(name):
|
||
"""按当前方言给标识符加引号：MySQL 里 key 是保留字，必须反引号。"""
|
||
from core.models import db
|
||
return db.engine.dialect.identifier_preparer.quote(name)
|
||
|
||
|
||
def meta_get(key, default=None):
|
||
from core.models import db
|
||
row = db.session.execute(
|
||
text("SELECT value FROM app_meta WHERE {} = :k".format(_q("key"))),
|
||
{"k": key}).fetchone()
|
||
return row[0] if row else default
|
||
|
||
|
||
def meta_set(key, value):
|
||
"""方言中立 upsert：SQLite 用 INSERT OR REPLACE，MySQL 用 ON DUPLICATE KEY。"""
|
||
from core.models import db
|
||
k = _q("key")
|
||
if db.engine.dialect.name == "mysql":
|
||
sql = ("INSERT INTO app_meta ({k}, value) VALUES (:k, :v) "
|
||
"ON DUPLICATE KEY UPDATE value = VALUES(value)").format(k=k)
|
||
else:
|
||
sql = "INSERT OR REPLACE INTO app_meta ({k}, value) VALUES (:k, :v)".format(k=k)
|
||
db.session.execute(text(sql), {"k": key, "v": value})
|
||
db.session.commit()
|
||
|
||
|
||
# ================== 启动前校验（不需要业务表）==================
|
||
def check_connection(uri=None):
|
||
"""连接探活 + 环境/库名一致性校验。失败抛 DBConfigError，由调用方决定是否退出。
|
||
|
||
不需要业务表，因此可以在 init_db 之前跑，做到"配置错就早点死"。
|
||
"""
|
||
uri = uri or build_db_uri()
|
||
|
||
if is_sqlite(uri):
|
||
_log.warning("数据库为 SQLite（%s）—— 回退模式，非正式用法", describe_target(uri))
|
||
return {"backend": "sqlite", "target": describe_target(uri)}
|
||
|
||
if DEPLOY_ENV not in _EXPECTED_DB_NAME:
|
||
raise DBConfigError(
|
||
"DEPLOY_ENV 取值非法: {!r}（只允许 dev / prod）".format(DEPLOY_ENV))
|
||
|
||
try:
|
||
eng = create_engine_from_uri(uri)
|
||
with eng.connect() as conn:
|
||
version = conn.execute(text("SELECT VERSION()")).scalar() or ""
|
||
db_name = conn.execute(text("SELECT DATABASE()")).scalar() or ""
|
||
eng.dispose()
|
||
except DBConfigError:
|
||
raise
|
||
except Exception as e:
|
||
raise DBConfigError(
|
||
"连接数据库失败: {}\n 目标: {}\n 请检查 .env 的 DB_HOST/DB_PORT/"
|
||
"DB_USER/DB_PASSWORD 与本机到该主机的网络。/".format(e, describe_target(uri)))
|
||
|
||
if not version.startswith(_REQUIRED_MYSQL_VERSION):
|
||
_log.warning("MySQL 版本为 %s，本项目按 %s.x 验证（更高版本通常兼容，"
|
||
"但请留意差异）", version, _REQUIRED_MYSQL_VERSION)
|
||
|
||
expected = _EXPECTED_DB_NAME[DEPLOY_ENV]
|
||
if db_name != expected and not DB_ALLOW_ENV_MISMATCH:
|
||
raise DBConfigError(
|
||
"环境与库名不匹配，已拒绝启动：\n"
|
||
" .env 声明 DEPLOY_ENV={} → 期望库名 {}\n"
|
||
" 实际连接的库 → {}\n"
|
||
" 目标: {}\n"
|
||
"如确认无误（例如临时指向别的库），可在 .env 加 DB_ALLOW_ENV_MISMATCH=1。"
|
||
.format(DEPLOY_ENV, expected, db_name, describe_target(uri)))
|
||
|
||
return {"backend": "mysql", "version": version, "database": db_name,
|
||
"target": describe_target(uri)}
|
||
|
||
|
||
# ================== 库标签校验（需要 app_meta 表，init_db 之后调用）==================
|
||
def verify_deployment_label():
|
||
"""核对 app_meta 里登记的库环境，与 .env 声明的环境是否一致。
|
||
|
||
首次见到某个库时会自动"认领"（写入标签）——生产库的标签由迁移脚本显式写入，
|
||
所以正常情况下不会走到自动认领分支。
|
||
"""
|
||
from core.models import db
|
||
if is_sqlite(db.engine.url):
|
||
return
|
||
|
||
recorded = meta_get(_K_ENV)
|
||
if recorded is None:
|
||
meta_set(_K_ENV, DEPLOY_ENV)
|
||
meta_set(_K_ID, uuid.uuid4().hex)
|
||
meta_set(_K_CLAIMED, _now())
|
||
_log.warning("首次见到该库，已登记为 %s 环境（部署标识 %s）",
|
||
DEPLOY_ENV, meta_get(_K_ID))
|
||
return
|
||
|
||
if recorded != DEPLOY_ENV and not DB_ALLOW_ENV_MISMATCH:
|
||
raise DBConfigError(
|
||
"目标库登记的环境与 .env 声明不符，已拒绝启动：\n"
|
||
" .env 声明 → {}\n"
|
||
" 库中登记（{}）→ {}\n"
|
||
" 目标: {}\n"
|
||
"这通常意味着你正把「{} 的配置」连到「{} 的库」上。\n"
|
||
"如确需这样操作，请在 .env 加 DB_ALLOW_ENV_MISMATCH=1。"
|
||
.format(DEPLOY_ENV, meta_get(_K_CLAIMED) or "认领时间未知", recorded,
|
||
describe_target(), DEPLOY_ENV, recorded))
|
||
elif recorded != DEPLOY_ENV:
|
||
_log.warning("库环境标签不符已按 DB_ALLOW_ENV_MISMATCH=1 放行：声明 %s / 库登记 %s",
|
||
DEPLOY_ENV, recorded)
|
||
|
||
|
||
# ================== 启动横幅 ==================
|
||
def print_banner():
|
||
"""每次启动打印一次「我现在连的是哪个库」。prod 用 WARNING 级。"""
|
||
from core.models import db
|
||
eng = db.engine
|
||
uri = eng.url
|
||
lines = []
|
||
if is_sqlite(uri):
|
||
lines.append(" 数据库 : SQLite 回退模式 <<< 非正式用法 >>>")
|
||
lines.append(" 目标 : " + describe_target())
|
||
else:
|
||
try:
|
||
with eng.connect() as conn:
|
||
version = conn.execute(text("SELECT VERSION()")).scalar() or ""
|
||
db_name = conn.execute(text("SELECT DATABASE()")).scalar() or ""
|
||
coll = conn.execute(text(
|
||
"SELECT @@collation_database")).scalar() or ""
|
||
n_tables = conn.execute(text(
|
||
"SELECT COUNT(*) FROM information_schema.tables "
|
||
"WHERE table_schema = DATABASE()")).scalar() or 0
|
||
counts = []
|
||
for t in ("device", "task_job", "device_group", "agent_conversation"):
|
||
try:
|
||
n = db.session.execute(
|
||
text("SELECT COUNT(*) FROM {}".format(t))).scalar()
|
||
counts.append("{}={}".format(t, n))
|
||
except Exception:
|
||
db.session.rollback()
|
||
lines.append(" MySQL : {} 库 {} ({})".format(version, db_name, coll))
|
||
lines.append(" 目标 : " + describe_target())
|
||
lines.append(" 表 : {} 张".format(n_tables))
|
||
lines.append(" 数据量 : " + " ".join(counts))
|
||
except Exception as e:
|
||
lines.append(" 连接信息采集失败: {}".format(e))
|
||
lines.append(" 库标签 : {}（{} 认领）".format(
|
||
meta_get(_K_ENV) or "未登记", meta_get(_K_CLAIMED) or "—"))
|
||
|
||
tag = "生产库 PROD" if DEPLOY_ENV == "prod" else "开发库 DEV"
|
||
bar = "=" * 64
|
||
head = ("{} 数据库已连接 [ {} ]{}".format(
|
||
bar, tag, " <<< 这是正式数据，操作前想清楚 >>>" if DEPLOY_ENV == "prod" else ""))
|
||
body = "\n".join(lines)
|
||
text_block = "\n{}\n{}\n{}".format(head, body, bar)
|
||
if DEPLOY_ENV == "prod":
|
||
_log.warning(text_block)
|
||
else:
|
||
_log.info(text_block)
|
||
|
||
|
||
def _now():
|
||
import time
|
||
return time.strftime("%Y-%m-%d %H:%M:%S")
|