crawler-plugin/backend/utils/db.py

"""
数据库连接与初始化

请求级连接复用：pymysql 每次 connect 都要做一次完整的鉴权握手，而本进程
（15124）只服务 /api/version、/api/version/latest 这类轻量查询，握手开销
占了单次请求的大头。连接因此挂在 flask.g 上按请求复用，请求结束由
close_db（注册在 app.teardown_appcontext）统一释放。
"""
import re

import pymysql

try:
    from flask import g, has_app_context
except ImportError:  # 非 Flask 环境（脚本/离线核查）下退化为每次新建连接
    g = None

    def has_app_context() -> bool:
        return False

# 请求级连接挂在 flask.g 上的属性名（避免与业务字段冲突）
_CONN_ATTR = "_shufu_db_conn"

try:
    from config import mysql_host as config_mysql_host
    from config import mysql_user as config_mysql_user
    from config import mysql_password as config_mysql_password
    from config import mysql_database as config_mysql_database
    from config import mysql_host_source as config_mysql_host_source
    from config import mysql_user_source as config_mysql_user_source
    from config import mysql_database_source as config_mysql_database_source
except ImportError:
    config_mysql_host = 'localhost'
    config_mysql_user = 'root'
    config_mysql_password = ''
    config_mysql_database = 'maixiang_ai'
    config_mysql_host_source = 'fallback.default'
    config_mysql_user_source = 'fallback.default'
    config_mysql_database_source = 'fallback.default'

mysql_host = config_mysql_host
mysql_user = config_mysql_user
mysql_password = config_mysql_password
mysql_database = config_mysql_database
mysql_host_source = config_mysql_host_source
mysql_user_source = config_mysql_user_source
mysql_database_source = config_mysql_database_source


def _safe_identifier(name):
    """仅允许字母、数字、下划线的数据库/表名片段，防止注入 DDL 片段。"""
    if not re.fullmatch(r'[A-Za-z0-9_]+', name or ''):
        raise ValueError(f'非法标识符: {name!r}')
    return name


def describe_db_target():
    return (
        f"{mysql_user}@{mysql_host}/{mysql_database} "
        f"(host={mysql_host_source}, user={mysql_user_source}, database={mysql_database_source})"
    )


def _connect():
    return pymysql.connect(
        host=mysql_host,
        user=mysql_user,
        password=mysql_password,
        database=mysql_database,
        charset='utf8mb4',
        cursorclass=pymysql.cursors.DictCursor
    )


def _is_alive(conn) -> bool:
    """连接可用性探测：已关闭直接判死，否则发一次 COM_PING。

    COM_PING 一次往返远小于一次完整鉴权握手；不用 ping(reconnect=True)
    （该参数在新版 pymysql 已废弃，且失败时会自行重连、语义不可控）。
    """
    try:
        if not conn.open:
            return False
        conn.ping(reconnect=False)
        return True
    except Exception:
        return False


def _close_quietly(conn) -> None:
    """静默关闭连接（可能为 None / 已关闭 / 断连）。"""
    if conn is None:
        return
    try:
        conn.close()
    except Exception:
        pass


def get_db():
    """取当前请求的数据库连接（同一请求内复用，避免每次请求都做鉴权握手）。

    连接挂在 flask.g 上，请求结束由 close_db 释放；被服务端 wait_timeout
    断开、或调用方提前 close 过，这里都会静默换成新连接，调用方无感知。
    无 Flask 应用上下文时（脚本调用）退回每次新建，保持旧行为。
    """
    ctx = g if has_app_context() else None
    conn = getattr(ctx, _CONN_ATTR, None) if ctx is not None else None
    if conn is not None:
        if _is_alive(conn):
            return conn
        # 断连/已被关闭：丢弃旧连接后重建（不把坏连接交给业务查询）
        _close_quietly(conn)
        try:
            delattr(ctx, _CONN_ATTR)
        except Exception:
            pass
    conn = _connect()
    if ctx is not None:
        setattr(ctx, _CONN_ATTR, conn)
    return conn


def close_db(exc=None) -> None:
    """释放当前请求的连接（注册为 app.teardown_appcontext 回调）。

    异常路径同样会走 teardown，因此查询抛错时连接不会泄漏到下一次请求。
    """
    ctx = g if has_app_context() else None
    if ctx is None:
        return
    conn = getattr(ctx, _CONN_ATTR, None)
    if conn is None:
        return
    try:
        delattr(ctx, _CONN_ATTR)
    except Exception:
        pass
    _close_quietly(conn)


def init_db():
    """确保版本公开 API 依赖的最小表结构存在。

    本进程（15124）只服务 /api/version、/api/version/latest，仅依赖 web_config 表。

    历史遗留说明（2026-09 全维度审查后删除）：这里原本还会创建 users 表、执行 role 角色
    迁移、并在查不到 super_admin 时用 ADMIN_PASSWORD（默认 admin123）插入一个超管。
    管理后台早已迁到 Java，users 表结构由 Flyway 迁移管理，本进程再写会与之冲突；
    其中「无超管即用默认密码建超管」在生产等同于后门（删除超管或改 role 枚举后本进程
    重启就会静默重建一个密码已知的超管），故整段移除。
    """
    conn = pymysql.connect(
        host=mysql_host,
        user=mysql_user,
        password=mysql_password,
        charset='utf8mb4'
    )
    try:
        db_name = _safe_identifier(mysql_database)
        with conn.cursor() as cur:
            cur.execute(f"CREATE DATABASE IF NOT EXISTS `{db_name}` DEFAULT CHARSET utf8mb4")
            cur.execute(f"USE `{db_name}`")
            cur.execute("""
                CREATE TABLE IF NOT EXISTS web_config (
                    id INT AUTO_INCREMENT PRIMARY KEY,
                    version VARCHAR(64) NOT NULL,
                    file_url VARCHAR(1024) NOT NULL,
                    created_at DATETIME DEFAULT CURRENT_TIMESTAMP
                )
            """)
        conn.commit()
    finally:
        conn.close()