"""当下供需 gap 脚本主驱 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.gap_script_demand.demand_id import build_gap_script_demand_id from app.gap_script_demand.exceptions import GapScriptDemandError from app.hot_content.types import MysqlConfig TABLE_NAME = "gap_script_demand_pool_di" class GapScriptDemandRepository: def __init__(self, config: MysqlConfig): if pymysql is None or DictCursor is None: raise GapScriptDemandError("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_gap_script_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) 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 GapScriptDemandError("gap script demand output strategy is required") if not partition_value: raise GapScriptDemandError("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_gap_script_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_matched_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, }