"""日报结构化入库:DB 连接、建表、写入。""" from __future__ import annotations import json import os from dataclasses import dataclass from typing import Any import pymysql from loguru import logger from .models import ReportData from .schema import DDL_STATEMENTS @dataclass(frozen=True) class DbConfig: """MySQL 连接配置(来自环境变量 NEWS_DB_*)。""" host: str port: int user: str password: str name: str def load_db_config() -> DbConfig: """从环境变量读取 NEWS_DB_*,缺失密码时抛异常(禁止默认密码)。""" host = os.environ.get("NEWS_DB_HOST", "127.0.0.1") port = int(os.environ.get("NEWS_DB_PORT", "13306")) user = os.environ.get("NEWS_DB_USER", "myquant") password = os.environ.get("NEWS_DB_PASSWORD", "") name = os.environ.get("NEWS_DB_NAME", "myquant") if not password: logger.error("NEWS_DB_PASSWORD 未配置,请在 .env 中设置") raise ValueError("NEWS_DB_PASSWORD 未配置") return DbConfig(host=host, port=port, user=user, password=password, name=name) def connect(cfg: DbConfig | None = None) -> pymysql.Connection: """建立短连接(autocommit=False)。失败时记录日志并抛出。""" cfg = cfg or load_db_config() try: conn = pymysql.connect( host=cfg.host, port=cfg.port, user=cfg.user, password=cfg.password, database=cfg.name, charset="utf8mb4", autocommit=False, cursorclass=pymysql.cursors.DictCursor, ) except Exception: logger.exception("连接 MySQL 失败: host={} port={} user={}", cfg.host, cfg.port, cfg.user) raise logger.debug("MySQL 已连接: {}/{}", cfg.host, cfg.name) return conn def init_schema(conn: pymysql.Connection) -> None: """建表(CREATE TABLE IF NOT EXISTS ×2),幂等。""" with conn.cursor() as cur: for ddl in DDL_STATEMENTS: cur.execute(ddl) conn.commit() logger.info("news_report / news_event 建表完成") def save_report(conn: pymysql.Connection, report: ReportData) -> int: """事务内写入一份日报。 幂等策略: - 主表按 (report_date, report_type, file_name) 唯一键 upsert; - 事件表 DELETE 该 report 旧行后全量 INSERT(整份覆盖一致)。 返回 report_id。 """ with conn.cursor() as cur: stats_json = json.dumps(report.stats, ensure_ascii=False) if report.stats else None cur.execute( """ INSERT INTO news_report (report_date, report_type, file_name, generated_at, ai_summary, stats) VALUES (%s, %s, %s, %s, %s, %s) ON DUPLICATE KEY UPDATE generated_at = VALUES(generated_at), ai_summary = VALUES(ai_summary), stats = VALUES(stats) """, ( report.report_date, report.report_type, report.file_name, report.generated_at, report.ai_summary, stats_json, ), ) cur.execute( "SELECT id FROM news_report WHERE report_date=%s AND report_type=%s AND file_name=%s", (report.report_date, report.report_type, report.file_name), ) row = cur.fetchone() if row is None: # pragma: no cover - 理论不可达 raise RuntimeError("写入 news_report 后查询不到 report_id") report_id: int = row["id"] cur.execute("DELETE FROM news_event WHERE report_id=%s", (report_id,)) for ev in report.events: cur.execute( """ INSERT INTO news_event (report_id, section, rank, importance, event_type, title, summary, sentiment, source, url) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s) """, ( report_id, ev.section, ev.rank, ev.importance, ev.event_type, ev.title, ev.summary, ev.sentiment, ev.source, ev.url, ), ) conn.commit() logger.info("日报已入库: report_id={} date={} type={} events={}", report_id, report.report_date, report.report_type, len(report.events)) return report_id def fetch_report(conn: pymysql.Connection, report_id: int) -> dict[str, Any] | None: """读侧辅助(联调/测试用),返回主表行。""" with conn.cursor() as cur: cur.execute("SELECT * FROM news_report WHERE id=%s", (report_id,)) return cur.fetchone() def exists_report( conn: pymysql.Connection, report_date: Any, report_type: str, file_name: str, ) -> bool: """判断主表是否已存在该唯一键记录(历史导入幂等用)。""" with conn.cursor() as cur: cur.execute( "SELECT 1 FROM news_report WHERE report_date=%s AND report_type=%s AND file_name=%s", (report_date, report_type, file_name), ) return cur.fetchone() is not None