SupplyAgent/supply_infra/scheduler/jobs/sync_global_tree_odps_to_mysql.py

"""
定时任务：从 ODPS global_category 同步有效“实质”分类树到 MySQL。

global_tree_element 已退出该流程。分类树直接以 stable_id /
parent_stable_id 建立层级，本地表继续保留稳定的 MySQL id 供现有下游引用。
"""
from __future__ import annotations

import logging
from datetime import datetime, timedelta
from typing import Any

from supply_infra.config import get_infra_settings
from supply_infra.db.repositories.global_tree_category_repo import (
    GlobalTreeCategoryRepository,
)
from supply_infra.db.session import get_session
from supply_infra.odps.client import get_odps_client

logger = logging.getLogger(__name__)


def _to_category_source_rows(
    categories: list[dict[str, Any]],
) -> list[dict[str, Any]]:
    """规范 stable_id 字段并按 stable_id 去重，保留最后一条。"""
    by_source_id: dict[int, dict[str, Any]] = {}
    for category in categories:
        stable_id = category.get("stable_id")
        name = category.get("name")
        if stable_id is None or name is None or not str(name).strip():
            logger.warning("Skip global category with missing stable_id/name: %s", category)
            continue
        parent_stable_id = category.get("parent_stable_id")
        source_id = int(stable_id)
        by_source_id[source_id] = {
            "source_id": source_id,
            "source_parent_id": (
                int(parent_stable_id)
                if parent_stable_id is not None and int(parent_stable_id) != 0
                else 0
            ),
            "name": str(name).strip(),
            "description": (
                str(category["description"])
                if category.get("description") is not None
                else None
            ),
            "level": (
                int(category["level"])
                if category.get("level") is not None
                else None
            ),
        }
    return sorted(
        by_source_id.values(),
        key=lambda row: (row.get("level") or 0, int(row["source_id"])),
    )


def _allocate_local_ids(
    categories: list[dict[str, Any]],
    source_id_to_mysql_id: dict[int, int],
    next_id: int,
) -> tuple[dict[int, int], set[int]]:
    """为新 stable_id 分配本地 id；缺少父节点的行不会进入正式树。"""
    updated_map = dict(source_id_to_mysql_id)
    pending = [
        category
        for category in categories
        if int(category["source_id"]) not in updated_map
    ]
    failed: set[int] = set()

    while pending:
        next_pending: list[dict[str, Any]] = []
        progress = False
        for category in pending:
            source_id = int(category["source_id"])
            source_parent_id = int(category["source_parent_id"])
            if source_parent_id != 0 and source_parent_id not in updated_map:
                next_pending.append(category)
                continue
            updated_map[source_id] = next_id
            next_id += 1
            progress = True

        if not progress:
            failed.update(int(category["source_id"]) for category in next_pending)
            for category in next_pending:
                logger.warning(
                    "Parent stable_id=%s not found for stable_id=%s, skip",
                    category["source_parent_id"],
                    category["source_id"],
                )
            break
        pending = next_pending

    return updated_map, failed


def _to_mysql_rows(
    categories: list[dict[str, Any]],
    source_id_to_mysql_id: dict[int, int],
    failed_source_ids: set[int],
) -> list[dict[str, Any]]:
    rows: list[dict[str, Any]] = []
    for category in categories:
        source_id = int(category["source_id"])
        if source_id in failed_source_ids:
            continue
        source_parent_id = int(category["source_parent_id"])
        parent_id = (
            source_id_to_mysql_id.get(source_parent_id)
            if source_parent_id != 0
            else None
        )
        if source_parent_id != 0 and parent_id is None:
            failed_source_ids.add(source_id)
            continue
        rows.append(
            {
                "id": source_id_to_mysql_id[source_id],
                "name": category["name"],
                "description": category.get("description"),
                "level": category.get("level"),
                "parent_id": parent_id,
                "source_id": source_id,
                "source_parent_id": source_parent_id,
                "is_delete": 0,
            }
        )
    return rows


def sync_global_tree_odps_to_mysql(partition_date: str | None = None) -> dict:
    """
    从 global_category 拉取完整有效分类树并同步到 MySQL。

    Args:
        partition_date: 分区日期 (YYYYMMDD)，默认昨天
    """
    if partition_date is None:
        partition_date = (datetime.now() - timedelta(days=1)).strftime("%Y%m%d")

    logger.info(
        "Starting global_category ODPS → MySQL sync for partition: %s",
        partition_date,
    )
    raw_categories = get_odps_client().fetch_global_categories(partition_date)
    categories = _to_category_source_rows(raw_categories)
    if not categories:
        raise RuntimeError(
            f"global_category returned no active categories for dt={partition_date}"
        )

    with get_session() as session:
        repository = GlobalTreeCategoryRepository(session)
        existing_active_count = repository.count_active()
        minimum_retained = int(
            existing_active_count
            * get_infra_settings().global_category_min_retention_ratio
        )
        if existing_active_count and len(categories) < minimum_retained:
            raise RuntimeError(
                "global_category active row count dropped below safety threshold: "
                f"current={existing_active_count} incoming={len(categories)} "
                f"minimum={minimum_retained}"
            )
        existing_map = repository.get_source_id_map()
        source_id_to_mysql_id, failed_source_ids = _allocate_local_ids(
            categories,
            existing_map,
            repository.get_max_id() + 1,
        )
        rows = _to_mysql_rows(
            categories,
            source_id_to_mysql_id,
            failed_source_ids,
        )
        affected = repository.bulk_upsert(rows)
        active_source_ids = {int(category["source_id"]) for category in categories}
        retired = repository.mark_missing_deleted(active_source_ids)

    existing_source_ids = set(existing_map)
    inserted = sum(
        1
        for row in rows
        if int(row["source_id"]) not in existing_source_ids
    )
    result = {
        "partition_date": partition_date,
        "source_type": "实质",
        "categories_fetched": len(raw_categories),
        "categories_valid": len(categories),
        "categories_inserted": inserted,
        "categories_existing_synced": len(rows) - inserted,
        "categories_failed": len(failed_source_ids),
        "categories_retired": retired,
        "rows_affected": affected,
        "global_tree_element_used": False,
        "synced_at": datetime.now().isoformat(),
    }
    logger.info("Global category sync completed: %s", result)
    return result


if __name__ == "__main__":
    import sys

    from supply_infra.scheduler.cli_result import run_cli

    run_cli(
        lambda: sync_global_tree_odps_to_mysql(
            partition_date=sys.argv[1] if len(sys.argv) > 1 else None,
        ),
        label="sync_global_tree_odps_to_mysql",
    )