403 lines
20 KiB
Python
403 lines
20 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)
|
||
# `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
|
||
|
||
|
||
# ---------------------------------------------------------------- 双实例表名前缀
|
||
# (2026-08-28 双系统定稿) 模拟与正式两套 PMS **共用 my_quant_db**, 靠表名区分:
|
||
# 正式实例在 .env 里设 PMS_TABLE_PREFIX (如 "real_"), 这里在 SQL 执行的**唯一入口**
|
||
# (fetch_all / fetch_one / execute / execute_many) 把每个 pms_* 表名统一改写成
|
||
# {前缀}pms_* —— 全仓库 129 处表名引用一处改写, 不逐条改 SQL。默认空串 = 不改写,
|
||
# 模拟环境行为逐字节不变。
|
||
#
|
||
# 改写按词边界匹配**小写 pms_ 开头的完整标识符**:
|
||
# · 19 张自有表全部命中; 已带前缀的 real_pms_* 不会二次加前缀 (前面无词边界);
|
||
# · 绑定参数的值不在 SQL 文本里, 不受影响; 大写参数键名 ('PMS_TOTAL_SCALE') 不匹配;
|
||
# · 字符串字面量里的表名 (如 information_schema 按 TABLE_NAME 查列) 一并改写,
|
||
# 恰好是想要的语义 —— migrate 类脚本查的就该是本实例的表;
|
||
# · 下游表 (trading_* 等) 不带 pms_ 前缀, 一个字不动 (它们另走 PMS_DS_TABLE_* 配置)。
|
||
# 两个绕过本入口直接拿引擎执行 DDL 的脚本 (init_db / migrate_archived_at) 已显式调
|
||
# map_tables; test_batch19 的旁路扫描钉住"不许再有第三个"。
|
||
_PREFIX_RE = re.compile(r"^[A-Za-z][A-Za-z0-9_]*$")
|
||
_PMS_TBL_RE = re.compile(r"\b(pms_[a-z0-9_]+)\b")
|
||
|
||
|
||
def _load_prefix() -> str:
|
||
p = (settings.PMS_TABLE_PREFIX or "").strip()
|
||
if p and not _PREFIX_RE.match(p):
|
||
raise ValueError(f"PMS_TABLE_PREFIX 不合法 (只许字母数字下划线): {p!r}")
|
||
return p
|
||
|
||
|
||
TABLE_PREFIX = _load_prefix()
|
||
|
||
|
||
def map_tables(sql: str) -> str:
|
||
"""按实例表名前缀改写 SQL 里的 pms_* 表名; 前缀为空原样返回。"""
|
||
if not TABLE_PREFIX:
|
||
return str(sql)
|
||
return _PMS_TBL_RE.sub(TABLE_PREFIX + r"\1", str(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:
|
||
sql = map_tables(sql)
|
||
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)。返回受影响行数。"""
|
||
sql = map_tables(sql)
|
||
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:
|
||
sql = map_tables(sql)
|
||
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
|
||
|
||
|
||
def apply_audited_repair(plan: dict, guards: list, *, rollback=False) -> dict:
|
||
"""维护专用：核验原行、逐条单表写入和审计在同一事务。绝不调下单或回放服务。"""
|
||
import json
|
||
from datetime import datetime
|
||
from app.core.ledger_repair import canonical
|
||
allowed = {
|
||
'pms_lot': {'qty','closed_qty','close_avg_price','realized_pnl','status','note'},
|
||
'pms_position': {'total_qty','avail_qty','base_qty','fill_qty','add_qty','dca_qty','t0_qty',
|
||
'avg_cost','realized_t_profit','cushion_pct','cushion_state','pct_of_scale','status'},
|
||
'pms_instruction': {'exec_qty','exec_avg_price','status','progress_json'},
|
||
'pms_strategy': {'state_json'},
|
||
}
|
||
if TABLE_PREFIX:
|
||
raise ValueError('历史修复仅允许模拟仓')
|
||
run_id = plan['run_id'] + ('_ROLLBACK' if rollback else '')
|
||
_check_cooldown('proxy')
|
||
with get_engine('proxy').begin() as c:
|
||
def query(sql, params=None):
|
||
sql = map_tables(sql)
|
||
assert_single_table(sql)
|
||
return c.execute(text(sql), params or {})
|
||
done = query("SELECT id FROM pms_op_log WHERE op='history_repair_summary' AND ref=:r FOR UPDATE",
|
||
{'r':run_id}).mappings().fetchall()
|
||
if done:
|
||
return {'ok':True, 'already_applied':True, 'run_id':run_id}
|
||
for g in guards:
|
||
rows = [dict(r) for r in query(g['sql'],g.get('params')).mappings().fetchall()]
|
||
if canonical(rows) != canonical(g['expected']):
|
||
raise ValueError('修复前置数据已变化，拒绝写入: '+g['name'])
|
||
for item in plan['changes']:
|
||
table, fields = item['table'], item['fields']
|
||
if table not in allowed or set(fields) - allowed[table] - {'updated_at'}:
|
||
raise ValueError('修复列越界')
|
||
current = [dict(r) for r in query(f'SELECT * FROM {table} WHERE id=:id FOR UPDATE',item['key']).mappings().fetchall()]
|
||
if len(current) != 1 or canonical(current[0]) != canonical(item['before']):
|
||
raise ValueError('目标行已变: '+table+':'+str(item['key']))
|
||
now = datetime.now().replace(microsecond=0)
|
||
for item in plan['changes']:
|
||
fields = dict(item['fields'])
|
||
fields.setdefault('updated_at',now)
|
||
query(f"UPDATE {item['table']} SET "+', '.join(f'{k}=:{k}' for k in fields)+' WHERE id=:row_id',
|
||
dict(fields,row_id=item['key']['id']))
|
||
after = dict(item['before'],**fields)
|
||
query('INSERT INTO pms_op_log (at,op,by_user,ts_code,result,reason,params_json,ref) '
|
||
'VALUES (:at,:op,:by,:code,:result,:reason,:params,:ref)',
|
||
dict(at=now,op='history_repair_row',by='codex',code=after.get('ts_code'),result='OK',
|
||
reason='逐笔账本修复'+('回退' if rollback else ''),
|
||
params=canonical(dict(item,after=after)),ref=run_id))
|
||
query('INSERT INTO pms_op_log (at,op,by_user,result,reason,params_json,ref) '
|
||
'VALUES (:at,:op,:by,:result,:reason,:params,:ref)',
|
||
dict(at=now,op='history_repair_summary',by='codex',result='OK',reason='历史修复事务完成',
|
||
params=canonical({k:v for k,v in plan.items() if k not in ('changes','allocations','daily')}),ref=run_id))
|
||
return {'ok':True,'run_id':run_id,'changed_rows':len(plan['changes'])}
|
||
|
||
|
||
def withdraw_rejected_proposal(old, hard):
|
||
"""提议生命周期专用事务。每条语句单表，原状态与 JSON 一起防并发。"""
|
||
import json
|
||
from datetime import datetime, timezone, timedelta
|
||
check = hard['last_accept_check']
|
||
stage = check['stage']
|
||
if stage not in ('judge', 'system', 'qualification') or check['status'] != 'withdrawn':
|
||
raise ValueError('只允许已确认的否决收尾')
|
||
encode = lambda v: json.dumps(v, ensure_ascii=False, default=str)
|
||
rejected = (check.get('details') or {}).get('rejected') or []
|
||
# 历史失败没有保存当时行情时，留空。不得把原提议旧价当作新否决的基线。
|
||
evidence = next((r.get('hard_numbers') for r in rejected if r.get('by') == stage and r.get('hard_numbers')), {})
|
||
ledger_hard = dict(evidence, proposal_withdrawal={'proposal_id': old['proposal_id'],
|
||
'check_at': check['at'], 'reason': check['message']})
|
||
now = datetime.now()
|
||
at = datetime.fromisoformat(check['at'])
|
||
if at.tzinfo is not None:
|
||
at = at.astimezone(timezone(timedelta(hours=8))).replace(tzinfo=None)
|
||
_check_cooldown('proxy')
|
||
with get_engine('proxy').begin() as c:
|
||
def query(sql, params):
|
||
sql = map_tables(sql)
|
||
assert_single_table(sql)
|
||
return c.execute(text(sql), params)
|
||
n = query("UPDATE pms_proposal SET status='EXPIRED', decided_at=:at, hard_numbers_json=:hn "
|
||
"WHERE proposal_id=:pid AND status='WAIT_USER' AND hard_numbers_json=:old",
|
||
{'at':now, 'hn':encode(hard), 'pid':old['proposal_id'], 'old':old['hard_numbers_json']}).rowcount
|
||
if not n:
|
||
return 0
|
||
query('INSERT INTO pms_action_ledger (ts_code,decided_at,action,arbiter,verdict,price_at,'
|
||
'hard_numbers_json,failed_checks_json,reason,ref_id,outcome_scored) '
|
||
'VALUES (:code,:at,:action,:arbiter,:verdict,:price,:hard,:failed,:reason,:ref,0)',
|
||
{'code':old['ts_code'], 'at':at, 'action':old['action'],
|
||
'arbiter':stage if stage in ('judge','system') else 'rule',
|
||
'verdict':'REJECT' if stage in ('judge','system') else 'NOTE',
|
||
'price':evidence.get('price') or 0, 'hard':encode(ledger_hard),
|
||
'failed':encode(check.get('reasons') or []),
|
||
'reason':('提议已撤下：'+check['message'])[:500], 'ref':old['proposal_id']})
|
||
return 1
|
||
|
||
|
||
def publish_prepared_proposal(old, candidate, verdict):
|
||
"""准备结果与变更审计同成同败。旧确认与后台发布只有一个能赢。"""
|
||
import json
|
||
from datetime import datetime
|
||
encode = lambda v: json.dumps(v, ensure_ascii=False, default=str)
|
||
hard = candidate['hard_numbers']
|
||
if (hard.get('preparation') or {}).get('status') != 'ready' or not hard.get('consumption_id'):
|
||
raise ValueError('准备结果缺少完整依据或消费凭据')
|
||
old_hard = old.get('hard_numbers') or {}
|
||
audit = dict(hard, previous_proposal={'qty': old['qty'], 'decision': old_hard.get('decision'),
|
||
'consumption_id': old_hard.get('consumption_id'), 'proposal_revision': old_hard.get('proposal_revision'),
|
||
'judge_verdict': old.get('judge_verdict'), 'judge_reason': old.get('judge_reason')})
|
||
_check_cooldown('proxy')
|
||
with get_engine('proxy').begin() as c:
|
||
def query(sql, params):
|
||
sql = map_tables(sql)
|
||
assert_single_table(sql)
|
||
return c.execute(text(sql), params)
|
||
n = query("UPDATE pms_proposal SET qty=:qty,hard_numbers_json=:hn,judge_verdict=:jv,judge_reason=:jr "
|
||
"WHERE proposal_id=:pid AND status='WAIT_USER' AND hard_numbers_json=:old",
|
||
{'qty':candidate['qty'],'hn':encode(hard),'jv':verdict.get('verdict'),
|
||
'jr':(verdict.get('reason') or '')[:500],'pid':old['proposal_id'],'old':old['hard_numbers_json']}).rowcount
|
||
if not n:
|
||
return 0
|
||
query('INSERT INTO pms_action_ledger (ts_code,decided_at,action,arbiter,verdict,price_at,'
|
||
'hard_numbers_json,failed_checks_json,reason,ref_id,outcome_scored) '
|
||
'VALUES (:code,:at,:action,:arbiter,:verdict,:price,:hard,:failed,:reason,:ref,0)',
|
||
{'code':old['ts_code'],'at':datetime.now(),'action':old['action'],'arbiter':'rule',
|
||
'verdict':'NOTE','price':float(hard.get('price') or 0),'hard':encode(audit),'failed':'[]',
|
||
'reason':'后台完成建议更新，等待人工确认','ref':old['proposal_id']})
|
||
return 1
|
||
|
||
|
||
def accept_open_proposal(old, hard, instruction, reason):
|
||
"""普通建仓的确认、动作审计、待执行指令同成同败；每条SQL仍严格单表。"""
|
||
import json
|
||
from datetime import datetime
|
||
if old['action'] != 'OPEN' or hard.get('kind') == 'strategy' or int(old['qty']) <= 0:
|
||
raise ValueError('只接受有效的普通建仓建议')
|
||
encode = lambda v: json.dumps(v, ensure_ascii=False, default=str)
|
||
now = datetime.now()
|
||
_check_cooldown('proxy')
|
||
with get_engine('proxy').begin() as c:
|
||
def query(sql, params):
|
||
sql = map_tables(sql)
|
||
assert_single_table(sql)
|
||
return c.execute(text(sql), params)
|
||
n = query("UPDATE pms_proposal SET status='ACCEPTED',decided_at=:at,hard_numbers_json=:hn "
|
||
"WHERE proposal_id=:pid AND status='WAIT_USER' AND hard_numbers_json=:old "
|
||
"AND expire_at>=:at AND qty=:qty",
|
||
{'at':now,'hn':encode(hard),'pid':old['proposal_id'],
|
||
'old':old['hard_numbers_json'],'qty':int(old['qty'])}).rowcount
|
||
if not n:
|
||
return False
|
||
query('INSERT INTO pms_action_ledger (ts_code,decided_at,action,arbiter,verdict,price_at,'
|
||
'hard_numbers_json,failed_checks_json,reason,ref_id,outcome_scored) '
|
||
"VALUES (:code,:at,'OPEN','user','PASS',:px,:hn,'[]',:reason,:pid,0)",
|
||
{'code':old['ts_code'],'at':now,'px':float(hard.get('price') or 0),
|
||
'hn':encode(hard),'reason':reason[:500],'pid':old['proposal_id']})
|
||
query('INSERT INTO pms_instruction (instruction_id,origin_type,origin_id,ts_code,action,side,'
|
||
'qty,limit_price,window_tdays,status,exec_qty,progress_json,created_at,updated_at) '
|
||
"VALUES (:iid,'proposal',:pid,:code,'OPEN','buy',:qty,:price,:window,'PROPOSED',0,:pg,:at,:at)",
|
||
{'iid':instruction['id'],'pid':old['proposal_id'],'code':old['ts_code'],
|
||
'qty':int(old['qty']),'price':hard.get('price'),'window':instruction['window'],
|
||
'pg':encode(instruction['progress']),'at':now})
|
||
return True
|