simonandClaude Opus 4.7
271a9343a5
7 Sprints 全部完成: Sprint 0: 基础设施 (DataManager + MariaDB) Sprint 1: 因子引擎 (34因子/12分类) Sprint 2: VectorBT 回测 (5策略+截面) Sprint 3: Optuna 优化 (+Walk-Forward) Sprint 4: ML 模型 (LightGBM+CatBoost) Sprint 5: Qwen 情绪因子 (三源新闻+日期对齐) Sprint 6: Agent 系统 (4Agent+日报.md/.html) 生产加固 (15项): Tushare双源fallback, SSH自动恢复, pool_pre_ping, save_daily先删后插, load_dotenv绝对路径, 日报5d/20d修复, RiskAgent改上证指数, 昨日对比+数据截止, mac_report utf8mb4, CLAUDE-*.md 9条已知Bug, demo全参数化, djapi数据源归一化, indexDatas API修正 Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
119 lines
3.4 KiB
Python
119 lines
3.4 KiB
Python
"""
|
|
MariaDB 连接管理。
|
|
|
|
使用 SQLAlchemy 连接池，通过 SSH 隧道访问远程数据库。
|
|
连接失败时自动尝试重建 SSH 隧道并重连。
|
|
"""
|
|
|
|
import os
|
|
import subprocess
|
|
import time
|
|
|
|
from sqlalchemy import create_engine, text
|
|
from sqlalchemy.orm import sessionmaker
|
|
|
|
from config.settings import MARIADB_CONFIG
|
|
|
|
_engine = None
|
|
_SessionLocal = None
|
|
|
|
|
|
def _build_url() -> str:
|
|
cfg = MARIADB_CONFIG
|
|
return (
|
|
"mysql+pymysql://{}:{}"
|
|
"@{}:{}/{}"
|
|
"?charset={}"
|
|
).format(cfg["user"], cfg["password"], cfg["host"], cfg["port"], cfg["database"], cfg["charset"])
|
|
|
|
|
|
def _reconnect_ssh() -> bool:
|
|
"""自动运行 autossh.sh 重建 SSH 隧道。每次调用都会尝试。"""
|
|
|
|
candidates = [
|
|
os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(
|
|
os.path.abspath(__file__)))), "shared", "script", "autossh.sh"),
|
|
os.path.expanduser("~/Downloads/cc-cursor/shared/script/autossh.sh"),
|
|
]
|
|
script = None
|
|
for c in candidates:
|
|
if os.path.exists(c):
|
|
script = c
|
|
break
|
|
|
|
if script is None:
|
|
print("[DB] SSH 脚本未找到")
|
|
return False
|
|
|
|
try:
|
|
print("[DB] 尝试重建 SSH 隧道: {}".format(script))
|
|
result = subprocess.run(["bash", script], capture_output=True, text=True, timeout=15)
|
|
time.sleep(3) # 等待隧道建立
|
|
if result.returncode == 0:
|
|
print("[DB] SSH 隧道重建完成")
|
|
return True
|
|
else:
|
|
print("[DB] SSH 隧道重建失败: {}".format(result.stderr[:200]))
|
|
return False
|
|
except Exception as e:
|
|
print("[DB] SSH 执行异常: {}".format(e))
|
|
return False
|
|
|
|
|
|
def _test_engine(engine) -> bool:
|
|
"""测试引擎连接是否存活。"""
|
|
try:
|
|
with engine.connect() as conn:
|
|
conn.execute(text("SELECT 1"))
|
|
return True
|
|
except Exception:
|
|
return False
|
|
|
|
|
|
def get_engine():
|
|
"""获取 SQLAlchemy Engine（懒初始化，断连时自动重连）。"""
|
|
global _engine, _SessionLocal
|
|
|
|
# 已有引擎但连接断开 → dispose 旧池 → SSH 重连 → 重建引擎
|
|
if _engine is not None and not _test_engine(_engine):
|
|
print("[DB] 连接已断开，尝试自动恢复...")
|
|
_engine.dispose() # 杀死所有旧连接
|
|
_engine = None
|
|
if _reconnect_ssh():
|
|
print("[DB] 连接已恢复")
|
|
else:
|
|
print("[DB] 自动恢复失败，将无法连接数据库")
|
|
|
|
# 创建引擎（首次或重建后）
|
|
if _engine is None:
|
|
cfg = MARIADB_CONFIG
|
|
_engine = create_engine(
|
|
_build_url(),
|
|
pool_size=cfg["pool_size"],
|
|
pool_recycle=cfg["pool_recycle"],
|
|
pool_pre_ping=True, # 每次取连接前先 SELECT 1 验证存活
|
|
echo=False,
|
|
)
|
|
# Session factory 下次获取时自动绑定新引擎
|
|
_SessionLocal = None
|
|
|
|
return _engine
|
|
|
|
|
|
def get_session():
|
|
"""获取一个新的数据库会话。"""
|
|
global _SessionLocal
|
|
if _SessionLocal is None:
|
|
_SessionLocal = sessionmaker(bind=get_engine())
|
|
return _SessionLocal()
|
|
|
|
|
|
def test_connection() -> bool:
|
|
"""测试数据库连接是否正常。"""
|
|
try:
|
|
engine = get_engine()
|
|
return _test_engine(engine)
|
|
except Exception as e:
|
|
print("[ERROR] 数据库连接失败: {}".format(e))
|
|
return False
|