gendesign/tradein-mvp/backend/app/tasks/dtp_stat_refresh.py

"""ДТП (dtp-stat.ru) refresh — CLI + scheduler run-lifecycle wrapper (#3410).

Наполняет `dtp_incidents` (пустая при деплое, миграция 294) через
`app/services/dtp_stat_loader.load_dtp_incidents` (скачивание ZIP + потоковый парс +
TRUNCATE+INSERT). Scheduler source='dtp_stat_refresh' (seed 293, enabled=false —
источник заморожен, см. докстринг loader'а; первый прогон ручной).

Запуск из контейнера:
    python -m app.tasks.dtp_stat_refresh
    python -m app.tasks.dtp_stat_refresh --dry-run
    python -m app.tasks.dtp_stat_refresh --src-path /path/to/local.zip

psycopg v3: `CAST(:x AS type)`, никогда `:x::type`.
"""

from __future__ import annotations

import argparse
import logging
import time
from dataclasses import dataclass, field

from sqlalchemy.orm import Session

from app.services import scrape_runs as runs_mod
from app.services.dtp_stat_loader import load_dtp_incidents

logger = logging.getLogger(__name__)


@dataclass
class DtpStatRefreshResult:
    parsed: int = 0
    inserted: int = 0
    duration_sec: float = field(default=0.0)

    def to_counters(self) -> dict[str, int]:
        return {
            "parsed": self.parsed,
            "inserted": self.inserted,
            "duration_sec": int(self.duration_sec),
        }


def run_dtp_stat_refresh(db: Session, *, run_id: int, params: dict) -> DtpStatRefreshResult:
    """Run-lifecycle wrapper для scheduler'а (source='dtp_stat_refresh') и ручного прогона.

    params:
      dry_run — bool, по умолчанию False. True: парсит и считает, БД не трогает.
      src_path — str, локальный путь к ZIP (для отладки/тестов, минуя скачивание).

    Финализирует scrape_runs (mark_done / mark_failed) со счётчиками.
    """
    dry_run = bool(params.get("dry_run", False))
    src_path = params.get("src_path")

    counters: dict[str, int] = {"parsed": 0, "inserted": 0}
    start = time.monotonic()

    try:
        runs_mod.update_heartbeat(db, run_id, counters)

        load_counts = load_dtp_incidents(db, src_path=src_path, dry_run=dry_run)
        if not dry_run:
            db.commit()

        result = DtpStatRefreshResult(
            parsed=load_counts["parsed"],
            inserted=load_counts["inserted"],
            duration_sec=time.monotonic() - start,
        )
        counters = result.to_counters()
        runs_mod.mark_done(db, run_id, counters)
        logger.info(
            "run_dtp_stat_refresh: run_id=%d DONE dry_run=%s parsed=%d inserted=%d duration=%.1fs",
            run_id,
            dry_run,
            result.parsed,
            result.inserted,
            result.duration_sec,
        )
        return result
    except Exception as exc:
        logger.exception("run_dtp_stat_refresh: run_id=%d FAILED", run_id)
        try:
            db.rollback()
        except Exception:
            pass
        runs_mod.mark_failed(db, run_id, str(exc)[:1000], counters)
        raise


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--src-path", default=None, help="Локальный ZIP вместо скачивания SRC_URL")
    parser.add_argument("--dry-run", action="store_true", help="Парсить и считать, не писать в БД")
    return parser


def main() -> None:
    logging.basicConfig(
        level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s"
    )
    args = build_parser().parse_args()

    from app.core.db import SessionLocal

    db = SessionLocal()
    try:
        counts = load_dtp_incidents(db, src_path=args.src_path, dry_run=args.dry_run)
        if not args.dry_run:
            db.commit()
        logger.info("dtp_stat_refresh CLI DONE: dry_run=%s %s", args.dry_run, counts)
    finally:
        db.close()


if __name__ == "__main__":
    main()