news/report_db/db.py

"""日报结构化入库：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 _ALTER_ADD_SOURCES, 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）+ 旧表 sources 列迁移，幂等。"""
    with conn.cursor() as cur:
        for ddl in DDL_STATEMENTS:
            cur.execute(ddl)
        # 旧库兼容:news_event 已存在但缺 sources 列时补列
        try:
            cur.execute(_ALTER_ADD_SOURCES)
            logger.info("news_event 迁移:新增 sources 多源列")
        except pymysql.err.OperationalError as e:
            if e.args and "Duplicate column" in str(e.args[0]):
                logger.debug("news_event.sources 列已存在,跳过迁移")
            else:
                raise
    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, sources, url)
                VALUES (%s, %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,
                    json.dumps(ev.sources, ensure_ascii=False) if ev.sources else None,
                    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