"""数据库连接装配：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": str(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")