SupplyAgent/scripts/backfill_demand_video_expansion_point_desc.py

#!/usr/bin/env python3
"""补全 demand_video_expansion 表中缺失的 point_desc。

从 multi_demand_video_point 按 (video_id, point_type, expanded_text=point_data) 匹配；
查不到则保持空值。

用法:
  .venv/bin/python scripts/backfill_demand_video_expansion_point_desc.py
  .venv/bin/python scripts/backfill_demand_video_expansion_point_desc.py --biz-dt 20260721
  .venv/bin/python scripts/backfill_demand_video_expansion_point_desc.py --biz-dt 20260721 --dry-run
"""
from __future__ import annotations

import argparse
import json
import logging
import sys
from pathlib import Path
from typing import Any

from sqlalchemy import or_, select, update

_ROOT = Path(__file__).resolve().parents[1]
if str(_ROOT) not in sys.path:
    sys.path.insert(0, str(_ROOT))

from agents.demand_video_expand_agent.tools.batch_save_demand_expansions import (
    _fill_missing_point_descs,
)
from supply_infra.db.models.demand_video_expansion import DemandVideoExpansion
from supply_infra.db.session import get_session

logger = logging.getLogger(__name__)
_BATCH_SIZE = 500


def _list_rows_missing_point_desc(biz_dt: str | None) -> list[dict[str, Any]]:
    stmt = select(DemandVideoExpansion).where(
        DemandVideoExpansion.is_delete == 0,
        or_(
            DemandVideoExpansion.point_desc.is_(None),
            DemandVideoExpansion.point_desc == "",
        ),
    )
    if biz_dt:
        stmt = stmt.where(DemandVideoExpansion.biz_dt == biz_dt)
    stmt = stmt.order_by(DemandVideoExpansion.id)

    with get_session() as session:
        rows = session.scalars(stmt).all()
        return [
            {
                "id": int(row.id),
                "biz_dt": str(row.biz_dt),
                "video_id": str(row.video_id),
                "point_type": str(row.point_type),
                "expanded_text": str(row.expanded_text),
                "point_desc": row.point_desc,
            }
            for row in rows
        ]


def backfill_missing_point_descs(
    biz_dt: str | None = None,
    *,
    dry_run: bool = False,
) -> dict[str, Any]:
    rows = _list_rows_missing_point_desc(biz_dt)
    result: dict[str, Any] = {
        "biz_dt": biz_dt,
        "dry_run": dry_run,
        "missing_total": len(rows),
        "filled": 0,
        "still_empty": 0,
        "updated": 0,
        "samples": [],
    }
    if not rows:
        return result

    with get_session() as session:
        _fill_missing_point_descs(rows, session)

    to_update: list[dict[str, Any]] = []
    for row in rows:
        if row.get("point_desc"):
            to_update.append(row)
            result["filled"] += 1
            if len(result["samples"]) < 10:
                result["samples"].append(
                    {
                        "id": row["id"],
                        "video_id": row["video_id"],
                        "point_type": row["point_type"],
                        "expanded_text": row["expanded_text"],
                        "point_desc": row["point_desc"][:80]
                        if len(str(row["point_desc"])) > 80
                        else row["point_desc"],
                    }
                )
        else:
            result["still_empty"] += 1

    if dry_run or not to_update:
        result["updated"] = 0
        return result

    with get_session() as session:
        for i in range(0, len(to_update), _BATCH_SIZE):
            batch = to_update[i : i + _BATCH_SIZE]
            for row in batch:
                session.execute(
                    update(DemandVideoExpansion)
                    .where(DemandVideoExpansion.id == int(row["id"]))
                    .values(point_desc=row["point_desc"])
                )
            result["updated"] += len(batch)

    return result


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        description="补全 demand_video_expansion 缺失的 point_desc",
    )
    parser.add_argument("--biz-dt", default=None, help="业务日期 YYYYMMDD，默认全表")
    parser.add_argument("--dry-run", action="store_true", help="仅统计，不写库")
    parser.add_argument("--json", action="store_true", help="以 JSON 输出结果")
    args = parser.parse_args(argv)

    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s %(levelname)s %(name)s: %(message)s",
    )

    result = backfill_missing_point_descs(args.biz_dt, dry_run=bool(args.dry_run))

    if args.json:
        print(json.dumps(result, ensure_ascii=False, indent=2, default=str))
    else:
        print("\n=== point_desc 补全 ===")
        print(f"biz_dt={result.get('biz_dt') or '全部'}")
        print(f"dry_run={result.get('dry_run')}")
        print(f"缺失记录={result.get('missing_total')}")
        print(f"可补全={result.get('filled')}")
        print(f"仍为空={result.get('still_empty')}")
        print(f"已更新={result.get('updated')}")
        if result.get("samples"):
            print("\n示例:")
            for item in result["samples"]:
                print(
                    f"  id={item['id']} video={item['video_id']} "
                    f"type={item['point_type']} text={item['expanded_text']!r} "
                    f"desc={item['point_desc']!r}"
                )

    return 0


if __name__ == "__main__":
    raise SystemExit(main())