192 lines
7.2 KiB
Python
192 lines
7.2 KiB
Python
# -*- 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)
|
|
|
|
|
|
def assert_single_table(sql: str) -> str:
|
|
"""静态检查: 一条 SQL 只能碰一张表, 且不得出现 JOIN / 逗号连表 / 跨表子查询。"""
|
|
s = re.sub(r"--[^\n]*", " ", str(sql))
|
|
s = re.sub(r"/\*.*?\*/", " ", s, flags=re.S)
|
|
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
|