|
|
@@ -0,0 +1,211 @@
|
|
|
+"""
|
|
|
+定时任务:从 ODPS 同步策略需求天级表到 MySQL,并对新词做归属分类。
|
|
|
+
|
|
|
+流程:
|
|
|
+1. 比对当天 ODPS / MySQL 行数,相同则跳过写入
|
|
|
+2. 有差异时拉取 ODPS,按 (strategy, demand_id) 只插入缺失、删除多余
|
|
|
+3. 查询当天 demand_name,按空格分词并去重
|
|
|
+4. 过滤 demand_belong_category 中已存在的词
|
|
|
+5. 剩余词按 100 个一批,调用 demand_belong_category_agent
|
|
|
+ (测试阶段仅调用 1 批)
|
|
|
+"""
|
|
|
+from __future__ import annotations
|
|
|
+
|
|
|
+import logging
|
|
|
+from datetime import datetime
|
|
|
+from typing import Any
|
|
|
+
|
|
|
+from agents.demand_belong_category_agent.run import main as classify_demand_words
|
|
|
+from supply_infra.db.repositories.demand_belong_category_repo import (
|
|
|
+ DemandBelongCategoryRepository,
|
|
|
+)
|
|
|
+from supply_infra.db.repositories.multi_demand_pool_di_repo import MultiDemandPoolDiRepository
|
|
|
+from supply_infra.db.session import get_session
|
|
|
+from supply_infra.odps.client import get_odps_client
|
|
|
+
|
|
|
+logger = logging.getLogger(__name__)
|
|
|
+
|
|
|
+_WORD_BATCH_SIZE = 100
|
|
|
+# 测试阶段只跑 1 批;上线后改为 None 表示跑完全部
|
|
|
+_MAX_CLASSIFY_BATCHES = 1
|
|
|
+
|
|
|
+RowKey = tuple[str, str]
|
|
|
+
|
|
|
+
|
|
|
+def _row_key(row: dict[str, Any]) -> RowKey:
|
|
|
+ return (row["strategy"], row["demand_id"])
|
|
|
+
|
|
|
+
|
|
|
+def _to_mysql_rows(raw_rows: list[dict[str, Any]], biz_dt: str) -> list[dict[str, Any]]:
|
|
|
+ """转换为 MySQL 行,按 (strategy, demand_id) 去重(保留最后一条)。"""
|
|
|
+ by_key: dict[RowKey, dict[str, Any]] = {}
|
|
|
+ for row in raw_rows:
|
|
|
+ strategy = row.get("strategy")
|
|
|
+ demand_id = row.get("demand_id")
|
|
|
+ demand_name = row.get("demand_name")
|
|
|
+ if strategy is None or demand_id is None or demand_name is None:
|
|
|
+ logger.warning("Skip row with missing required fields: %s", row)
|
|
|
+ continue
|
|
|
+
|
|
|
+ mapped = {
|
|
|
+ "strategy": str(strategy),
|
|
|
+ "demand_id": str(demand_id),
|
|
|
+ "demand_name": str(demand_name),
|
|
|
+ "weight": row.get("weight"),
|
|
|
+ "type": str(row["type"]) if row.get("type") is not None else None,
|
|
|
+ "video_count": row.get("video_count"),
|
|
|
+ "video_list": None,
|
|
|
+ "extend": str(row["extend"]) if row.get("extend") is not None else None,
|
|
|
+ "biz_dt": biz_dt,
|
|
|
+ }
|
|
|
+ by_key[_row_key(mapped)] = mapped
|
|
|
+ return list(by_key.values())
|
|
|
+
|
|
|
+
|
|
|
+def _tokenize_demand_names(demand_names: list[str]) -> list[str]:
|
|
|
+ """按空格分词,去空、去重(保持首次出现顺序)。"""
|
|
|
+ seen: set[str] = set()
|
|
|
+ tokens: list[str] = []
|
|
|
+ for name in demand_names:
|
|
|
+ for token in str(name).split():
|
|
|
+ word = token.strip()
|
|
|
+ if not word or word in seen:
|
|
|
+ continue
|
|
|
+ seen.add(word)
|
|
|
+ tokens.append(word)
|
|
|
+ return tokens
|
|
|
+
|
|
|
+
|
|
|
+def _chunked(items: list[str], size: int) -> list[list[str]]:
|
|
|
+ return [items[i : i + size] for i in range(0, len(items), size)]
|
|
|
+
|
|
|
+
|
|
|
+def _classify_new_words(biz_dt: str, *, max_batches: int | None = _MAX_CLASSIFY_BATCHES) -> dict:
|
|
|
+ """分词 → 过滤已存在词 → 分批调用归属分类 agent。"""
|
|
|
+ with get_session() as session:
|
|
|
+ pool_repo = MultiDemandPoolDiRepository(session)
|
|
|
+ category_repo = DemandBelongCategoryRepository(session)
|
|
|
+
|
|
|
+ demand_names = pool_repo.list_demand_names_by_biz_dt(biz_dt)
|
|
|
+ tokens = _tokenize_demand_names(demand_names)
|
|
|
+ existing = category_repo.get_existing_names(tokens)
|
|
|
+ new_words = [w for w in tokens if w not in existing]
|
|
|
+
|
|
|
+ batches = _chunked(new_words, _WORD_BATCH_SIZE)
|
|
|
+ if max_batches is not None:
|
|
|
+ batches = batches[:max_batches]
|
|
|
+
|
|
|
+ logger.info(
|
|
|
+ "Classify prepare: demand_names=%d tokens=%d existing=%d new=%d batches=%d (max=%s)",
|
|
|
+ len(demand_names),
|
|
|
+ len(tokens),
|
|
|
+ len(existing),
|
|
|
+ len(new_words),
|
|
|
+ len(batches),
|
|
|
+ max_batches,
|
|
|
+ )
|
|
|
+
|
|
|
+ classified_batches = 0
|
|
|
+ for idx, batch in enumerate(batches, start=1):
|
|
|
+ logger.info("Classifying batch %d/%d (%d words)", idx, len(batches), len(batch))
|
|
|
+ classify_demand_words(batch)
|
|
|
+ classified_batches += 1
|
|
|
+
|
|
|
+ return {
|
|
|
+ "demand_names": len(demand_names),
|
|
|
+ "tokens": len(tokens),
|
|
|
+ "existing_filtered": len(existing),
|
|
|
+ "new_words": len(new_words),
|
|
|
+ "batches_total": (len(new_words) + _WORD_BATCH_SIZE - 1) // _WORD_BATCH_SIZE
|
|
|
+ if new_words
|
|
|
+ else 0,
|
|
|
+ "batches_ran": classified_batches,
|
|
|
+ }
|
|
|
+
|
|
|
+
|
|
|
+def _sync_diff(partition_date: str) -> dict[str, Any]:
|
|
|
+ """行数不同时拉取 ODPS,只同步 (strategy, demand_id) 差异。"""
|
|
|
+ odps = get_odps_client()
|
|
|
+ raw_rows = odps.fetch_multi_demand_pool(partition_date)
|
|
|
+ mysql_rows = _to_mysql_rows(raw_rows, partition_date)
|
|
|
+ odps_keys = {_row_key(r) for r in mysql_rows}
|
|
|
+ odps_by_key = {_row_key(r): r for r in mysql_rows}
|
|
|
+
|
|
|
+ with get_session() as session:
|
|
|
+ repo = MultiDemandPoolDiRepository(session)
|
|
|
+ mysql_keys = repo.list_keys_by_biz_dt(partition_date)
|
|
|
+
|
|
|
+ to_insert_keys = odps_keys - mysql_keys
|
|
|
+ to_delete_keys = mysql_keys - odps_keys
|
|
|
+
|
|
|
+ insert_rows = [odps_by_key[k] for k in to_insert_keys]
|
|
|
+ deleted = repo.delete_by_keys(partition_date, list(to_delete_keys))
|
|
|
+ inserted = repo.bulk_insert(insert_rows)
|
|
|
+
|
|
|
+ logger.info(
|
|
|
+ "Diff sync: odps=%d mysql_before=%d insert=%d delete=%d",
|
|
|
+ len(odps_keys),
|
|
|
+ len(mysql_keys),
|
|
|
+ inserted,
|
|
|
+ deleted,
|
|
|
+ )
|
|
|
+ return {
|
|
|
+ "fetched": len(raw_rows),
|
|
|
+ "odps_unique": len(odps_keys),
|
|
|
+ "mysql_before": len(mysql_keys),
|
|
|
+ "inserted": inserted,
|
|
|
+ "deleted": deleted,
|
|
|
+ "skipped_invalid": len(raw_rows) - len(mysql_rows),
|
|
|
+ }
|
|
|
+
|
|
|
+
|
|
|
+def sync_multi_demand_pool_odps_to_mysql(partition_date: str | None = None) -> dict:
|
|
|
+ """
|
|
|
+ 从 ODPS 增量同步策略需求天级数据到 MySQL,并对新词做归属分类。
|
|
|
+
|
|
|
+ Args:
|
|
|
+ partition_date: 分区日期 (YYYYMMDD),默认当天
|
|
|
+ """
|
|
|
+ if partition_date is None:
|
|
|
+ partition_date = datetime.now().strftime("%Y%m%d")
|
|
|
+
|
|
|
+ logger.info(
|
|
|
+ "Starting multi demand pool ODPS → MySQL sync for partition: %s",
|
|
|
+ partition_date,
|
|
|
+ )
|
|
|
+
|
|
|
+ odps = get_odps_client()
|
|
|
+ odps_count = odps.count_multi_demand_pool(partition_date)
|
|
|
+ with get_session() as session:
|
|
|
+ mysql_count = MultiDemandPoolDiRepository(session).count_by_biz_dt(partition_date)
|
|
|
+
|
|
|
+ logger.info("Count check: odps=%d mysql=%d", odps_count, mysql_count)
|
|
|
+
|
|
|
+ if odps_count == mysql_count:
|
|
|
+ sync_stats: dict[str, Any] = {
|
|
|
+ "skipped_same_count": True,
|
|
|
+ "odps_count": odps_count,
|
|
|
+ "mysql_count": mysql_count,
|
|
|
+ "fetched": 0,
|
|
|
+ "inserted": 0,
|
|
|
+ "deleted": 0,
|
|
|
+ }
|
|
|
+ logger.info("Same count, skip ODPS data sync")
|
|
|
+ else:
|
|
|
+ sync_stats = {
|
|
|
+ "skipped_same_count": False,
|
|
|
+ "odps_count": odps_count,
|
|
|
+ "mysql_count": mysql_count,
|
|
|
+ **_sync_diff(partition_date),
|
|
|
+ }
|
|
|
+
|
|
|
+ classify_stats = _classify_new_words(partition_date)
|
|
|
+
|
|
|
+ result = {
|
|
|
+ "partition_date": partition_date,
|
|
|
+ **sync_stats,
|
|
|
+ "classify": classify_stats,
|
|
|
+ "synced_at": datetime.now().isoformat(),
|
|
|
+ }
|
|
|
+ logger.info("Multi demand pool sync completed: %s", result)
|
|
|
+ return result
|