tradingSystem/app/db/session.py

# -*- coding: utf-8 -*-
"""
数据库连接与「严格单表访问」守卫
==================================
开发约定 (README): 153 代理侧数据库**严格单表访问** —— 代理不支持多表联查,
写错了要到线上才报错。本模块在 SQL 执行入口做静态检查, 把这条纪律钉死在代码里。

三个库 (键名对齐 bionic_trader, 同一 .env 可共用):
  proxy   PROXY_DB_URL        153 代理: pms_* 全部表 + trading_* + strategy_daily_results
  factor  SOURCE_DB_EXT_DSN   因子分表 (自算参考位/MA/ATR)
  index   DB_MYSQL_URL        大盘指数 zs_day_data (页面区制提示)

连接失败一律抛 DBUnavailable, 由上层 (API/任务) 捕获并降级 —— 页面照常打开、
调度任务照常守成不产生新指令 (设计 §13 三条铁律之「故障即守成」)。
"""
from __future__ import annotations

import logging
import re
import time
import threading

from sqlalchemy import create_engine, text
from sqlalchemy.exc import SQLAlchemyError

from config.settings import settings

logger = logging.getLogger("pms.db")

_engines = {}
_lock = threading.Lock()

# 短路冷却: 库刚连挂过就直接快速失败, 不再逐次等 TCP 超时。
# 目的 —— 库不可达时页面仍能秒开 (每个接口只吃一次超时), 调度任务也不会被拖死。
FAIL_COOLDOWN_SEC = 10.0
_last_fail = {}

_DSN = {
    "proxy": lambda: settings.PROXY_DB_URL,
    "factor": lambda: settings.SOURCE_DB_EXT_DSN,
    "index": lambda: settings.DB_MYSQL_URL,
}


class DBUnavailable(RuntimeError):
    """连库/执行失败。上层必须捕获并降级, 不得让异常冒泡成 500。"""


class MultiTableSQL(ValueError):
    """违反严格单表访问纪律。"""


# ---------------------------------------------------------------- 单表守卫
_TABLE_RE = re.compile(r"\b(?:from|join|into|update)\s+`?([a-zA-Z_][a-zA-Z0-9_]*)`?", re.I)
_JOIN_RE = re.compile(r"\bjoin\b", re.I)
# FROM 子句 (到下一个子句关键字为止) —— 用于识别 "FROM a, b" 这种逗号连表
_FROM_CLAUSE_RE = re.compile(
    r"\bfrom\s+(.+?)(?=\bwhere\b|\bgroup\b|\border\b|\blimit\b|\bhaving\b|\bunion\b|\bon\b|\)|;|$)",
    re.I | re.S)
# `ON DUPLICATE KEY UPDATE` 里的 UPDATE 后面跟的是**列名**不是表名。不先摘掉这一段,
# 下面的 _TABLE_RE 会把第一个被赋值的列当成第二张表, 于是所有 upsert 一律误判成多表。
# 2026-07-28 修: 四处 upsert 全部命中该误判 —— set_param (页面改参数)、ensure_position
# (建仓落账本第一步)、upsert_report (日报落库)、upsert_industry (行业映射导入),
# 一执行就抛 MultiTableSQL。单测走内存桩不经过本函数, 所以一直没暴露。
# 见 scripts/test_batch6_units.py 的「单表守卫」一组。
_ODKU_RE = re.compile(r"\bon\s+duplicate\s+key\s+update\b", re.I)


def assert_single_table(sql: str) -> str:
    """静态检查: 一条 SQL 只能碰一张表, 且不得出现 JOIN / 逗号连表 / 跨表子查询。"""
    s = re.sub(r"--[^\n]*", " ", str(sql))
    s = re.sub(r"/\*.*?\*/", " ", s, flags=re.S)
    s = _ODKU_RE.sub(" ", s)            # 见上: 摘掉 upsert 尾巴, 它不引入新表
    if _JOIN_RE.search(s):
        raise MultiTableSQL(f"严格单表访问: SQL 含 JOIN —— {s[:120]}")
    for clause in _FROM_CLAUSE_RE.findall(s):
        parts = [x.strip() for x in clause.split(",") if x.strip()]
        if len(parts) > 1:
            raise MultiTableSQL(f"严格单表访问: FROM 子句逗号连表 {parts} —— {s[:120]}")
    tables = {m.lower() for m in _TABLE_RE.findall(s)}
    if len(tables) > 1:
        raise MultiTableSQL(f"严格单表访问: SQL 涉及多表 {sorted(tables)} —— {s[:120]}")
    return sql


# ---------------------------------------------------------------- 引擎
def get_engine(name: str = "proxy"):
    if name in _engines:
        return _engines[name]
    with _lock:
        if name in _engines:
            return _engines[name]
        dsn_fn = _DSN.get(name)
        if not dsn_fn:
            raise DBUnavailable(f"未知数据源 {name}")
        try:
            eng = create_engine(dsn_fn(), pool_pre_ping=True, pool_size=5, max_overflow=5,
                                pool_recycle=3600, future=True,
                                # 连接超时兜底: 库不可达时快速失败, 不拖死页面与调度任务
                                connect_args={"connect_timeout": 3})
        except Exception as e:  # DSN 非法等
            raise DBUnavailable(f"数据源 {name} 初始化失败: {e}") from e
        _engines[name] = eng
        return eng


def _check_cooldown(name: str):
    ts = _last_fail.get(name)
    if ts and (time.monotonic() - ts) < FAIL_COOLDOWN_SEC:
        raise DBUnavailable(
            f"数据源 {name} 刚刚连接失败, {FAIL_COOLDOWN_SEC:.0f} 秒内快速失败 "
            f"(避免逐次等待超时); 恢复后自动重连")


def _mark(name: str, ok: bool):
    if ok:
        _last_fail.pop(name, None)
    else:
        _last_fail[name] = time.monotonic()


def ping(name: str = "proxy") -> dict:
    try:
        with get_engine(name).connect() as c:
            c.execute(text("SELECT 1"))
        _mark(name, True)
        return {"ok": True, "source": name}
    except Exception as e:
        _mark(name, False)
        return {"ok": False, "source": name, "error": f"{type(e).__name__}: {e}"}


# ---------------------------------------------------------------- 执行助手
def fetch_all(sql: str, params=None, source: str = "proxy") -> list:
    assert_single_table(sql)
    _check_cooldown(source)
    try:
        with get_engine(source).connect() as c:
            rows = c.execute(text(sql), params or {}).mappings().fetchall()
        _mark(source, True)
        return [dict(r) for r in rows]
    except MultiTableSQL:
        raise
    except SQLAlchemyError as e:
        _mark(source, _is_sql_error(e))
        raise DBUnavailable(f"查询失败({source}): {type(e).__name__}: {e}") from e
    except Exception as e:
        _mark(source, False)
        raise DBUnavailable(f"查询异常({source}): {type(e).__name__}: {e}") from e


def _is_sql_error(e) -> bool:
    """区分「SQL 本身错」(库是通的, 不进冷却) 与「连不上库」(进冷却)。"""
    msg = str(e).lower()
    return not any(k in msg for k in ("can't connect", "connection refused", "timed out",
                                      "lost connection", "gone away", "no route to host",
                                      "name or service not known"))


def fetch_one(sql: str, params=None, source: str = "proxy"):
    rows = fetch_all(sql, params, source)
    return rows[0] if rows else None


def execute(sql: str, params=None, source: str = "proxy") -> int:
    """写操作 (INSERT/UPDATE/DELETE)。返回受影响行数。"""
    assert_single_table(sql)
    _check_cooldown(source)
    try:
        with get_engine(source).begin() as c:
            r = c.execute(text(sql), params or {})
            _mark(source, True)
            return int(r.rowcount or 0)
    except MultiTableSQL:
        raise
    except SQLAlchemyError as e:
        _mark(source, _is_sql_error(e))
        raise DBUnavailable(f"写入失败({source}): {type(e).__name__}: {e}") from e
    except Exception as e:
        _mark(source, False)
        raise DBUnavailable(f"写入异常({source}): {type(e).__name__}: {e}") from e


def execute_many(sql: str, seq_params: list, source: str = "proxy") -> int:
    assert_single_table(sql)
    if not seq_params:
        return 0
    _check_cooldown(source)
    try:
        with get_engine(source).begin() as c:
            r = c.execute(text(sql), seq_params)
            _mark(source, True)
            return int(r.rowcount or 0)
    except SQLAlchemyError as e:
        _mark(source, _is_sql_error(e))
        raise DBUnavailable(f"批量写入失败({source}): {type(e).__name__}: {e}") from e
    except Exception as e:
        _mark(source, False)
        raise DBUnavailable(f"批量写入异常({source}): {type(e).__name__}: {e}") from e