external_demand/app/festival_demand/repository.py

"""节日需求 MySQL 仓储。"""

from __future__ import annotations

from typing import Any

try:
    import pymysql
    from pymysql.cursors import DictCursor
except ImportError:  # pragma: no cover
    pymysql = None
    DictCursor = None

from app.festival_demand.demand_id import build_festival_demand_id
from app.festival_demand.exceptions import FestivalDemandError
from app.hot_content.types import MysqlConfig

TABLE_NAME = "festival_demand_pool_di"


class FestivalDemandRepository:
    def __init__(self, config: MysqlConfig):
        if pymysql is None or DictCursor is None:
            raise FestivalDemandError("missing dependency: pip install pymysql")
        self.conn = pymysql.connect(
            host=config.host,
            port=config.port,
            user=config.user,
            password=config.password,
            database=config.database,
            charset=config.charset,
            autocommit=True,
            cursorclass=DictCursor,
        )
        self._ensure_table()

    def close(self) -> None:
        self.conn.close()

    def _ensure_table(self) -> None:
        sql = f"""
            CREATE TABLE IF NOT EXISTS {TABLE_NAME} (
                id BIGINT UNSIGNED NOT NULL AUTO_INCREMENT COMMENT '自增主键',
                strategy VARCHAR(128) NOT NULL COMMENT '策略',
                demand_id CHAR(32) NOT NULL COMMENT '需求id',
                demand_name VARCHAR(1024) NOT NULL COMMENT '需求',
                partition_dt VARCHAR(8) NOT NULL DEFAULT '' COMMENT '业务日期 yyyymmdd',
                created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP COMMENT '创建时间',
                updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP
                    ON UPDATE CURRENT_TIMESTAMP COMMENT '更新时间',
                PRIMARY KEY (id),
                UNIQUE KEY uk_festival_demand_pool (strategy, demand_id),
                KEY idx_strategy_partition (strategy, partition_dt),
                KEY idx_demand_name (demand_name(191))
            ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci
        """
        with self.conn.cursor() as cursor:
            cursor.execute(sql)
            self._ensure_partition_dt_column(cursor)

    @staticmethod
    def _ensure_partition_dt_column(cursor: Any) -> None:
        cursor.execute(
            f"""
            SELECT COUNT(*) AS cnt
            FROM information_schema.COLUMNS
            WHERE TABLE_SCHEMA = DATABASE()
              AND TABLE_NAME = '{TABLE_NAME}'
              AND COLUMN_NAME = 'partition_dt'
            """
        )
        row = cursor.fetchone() or {}
        if int(row.get("cnt") or 0) > 0:
            return
        cursor.execute(
            f"""
            ALTER TABLE {TABLE_NAME}
            ADD COLUMN partition_dt VARCHAR(8) NOT NULL DEFAULT '' COMMENT '业务日期 yyyymmdd'
            AFTER demand_name,
            ADD KEY idx_strategy_partition (strategy, partition_dt)
            """
        )

    def has_partition_data(self, *, strategy: str, partition_dt: str) -> bool:
        strategy_value = strategy.strip()
        partition_value = partition_dt.strip()
        if not strategy_value or not partition_value:
            return False
        sql = f"""
            SELECT 1
            FROM {TABLE_NAME}
            WHERE strategy = %s
              AND partition_dt = %s
            LIMIT 1
        """
        with self.conn.cursor() as cursor:
            cursor.execute(sql, (strategy_value, partition_value))
            return cursor.fetchone() is not None

    def build_demand_rows(
        self,
        *,
        demand_names: list[str],
        strategy: str,
        partition_dt: str,
    ) -> list[dict[str, str]]:
        strategy_value = strategy.strip()
        partition_value = partition_dt.strip()
        if not strategy_value:
            raise FestivalDemandError("festival demand output strategy is required")
        if not partition_value:
            raise FestivalDemandError("partition_dt is required")

        rows: list[dict[str, str]] = []
        seen_demand_ids: set[str] = set()
        for raw_name in demand_names:
            demand_name = str(raw_name).strip()
            if not demand_name:
                continue
            demand_id = build_festival_demand_id(
                strategy=strategy_value,
                demand_name=demand_name,
                partition_dt=partition_value,
            )
            if demand_id in seen_demand_ids:
                continue
            seen_demand_ids.add(demand_id)
            rows.append(
                {
                    "strategy": strategy_value,
                    "demand_id": demand_id,
                    "demand_name": demand_name,
                    "partition_dt": partition_value,
                }
            )
        return rows

    def insert_demand_rows(self, rows: list[dict[str, str]]) -> int:
        if not rows:
            return 0
        sql = f"""
            INSERT INTO {TABLE_NAME} (
                strategy,
                demand_id,
                demand_name,
                partition_dt
            )
            VALUES (%s, %s, %s, %s)
        """
        payload = [
            (
                row["strategy"],
                row["demand_id"],
                row["demand_name"],
                row["partition_dt"],
            )
            for row in rows
        ]
        with self.conn.cursor() as cursor:
            cursor.executemany(sql, payload)
        return len(rows)

    def persist_generated_demands(
        self,
        *,
        demand_names: list[str],
        strategy: str,
        partition_dt: str,
    ) -> dict[str, Any]:
        """当天已有数据则跳过；否则写入 MySQL 并返回待同步 ODPS 的行。"""
        strategy_value = strategy.strip()
        partition_value = partition_dt.strip()
        base_result = {
            "table": TABLE_NAME,
            "strategy": strategy_value,
            "partition_dt": partition_value,
            "input_count": len(demand_names),
            "skipped": False,
            "saved_count": 0,
            "rows": [],
        }

        if self.has_partition_data(strategy=strategy_value, partition_dt=partition_value):
            return {
                **base_result,
                "skipped": True,
                "skip_reason": "partition_data_exists",
            }

        rows = self.build_demand_rows(
            demand_names=demand_names,
            strategy=strategy_value,
            partition_dt=partition_value,
        )
        if not rows:
            return base_result

        saved_count = self.insert_demand_rows(rows)
        return {
            **base_result,
            "saved_count": saved_count,
            "rows": rows,
        }