gendesign/backend/app/workers/tasks/scrape_cadastre.py

"""Celery tasks для bulk cadastre harvest (GH issue #168, PR 3/5).

Архитектура:
  - enqueue_cadastre_harvest(job_id) — читает cadastre_jobs.raw_targets,
    запускает bulk_harvest_quarter_task для каждого квартала
  - bulk_harvest_quarter_task(quarter, job_id) — per-quarter Celery task,
    оборачивает async harvest_quarter() через asyncio.run()

Saga state в cadastre_jobs.phase_state JSONB:
  {phase, quarter, parcels_total_meta, buildings_total_meta, overflow_cats, ...}

acks_late=True + idempotent upsert (ON CONFLICT) → resume после redeploy.
WAF 403 → no retry, increment waf_blocked_count.
"""

from __future__ import annotations

import asyncio
import datetime as _dt
import json
import logging
from typing import Any

from sqlalchemy import text

from app.core.db import SessionLocal
from app.scrapers.nspd_bulk_client import NspdBulkWafError
from app.workers.celery_app import celery_app

logger = logging.getLogger(__name__)

# Очередь для cadastre задач
CADASTRE_QUEUE = "geo"


# ── Saga state helpers ───────────────────────────────────────────────────────


def _update_job(
    db: Any,
    job_id: int,
    *,
    status: str | None = None,
    started_at: _dt.datetime | None = None,
    finished_at: _dt.datetime | None = None,
    phase_state_update: dict[str, Any] | None = None,
    heartbeat: bool = False,
    targets_done_inc: int = 0,
    targets_failed_inc: int = 0,
    requests_count_inc: int = 0,
    waf_blocked_inc: int = 0,
    last_error: str | None = None,
) -> None:
    """Атомарный UPDATE одной строки cadastre_jobs.

    Все параметры опциональны — только переданные поля обновляются.
    phase_state_update мержится поверх текущего JSONB.
    """
    sets: list[str] = []
    params: dict[str, Any] = {"id": job_id}

    if heartbeat or heartbeat is True:
        sets.append("heartbeat_at = NOW()")

    if status is not None:
        sets.append("status = :status")
        params["status"] = status

    if started_at is not None:
        sets.append("started_at = COALESCE(started_at, :started_at)")
        params["started_at"] = started_at

    if finished_at is not None:
        sets.append("finished_at = :finished_at")
        params["finished_at"] = finished_at

    if last_error is not None:
        sets.append("error = :last_error")
        params["last_error"] = last_error[:2000]

    if phase_state_update is not None:
        # Merge JSONB: текущий phase_state || новый partial update
        sets.append(
            "phase_state = COALESCE(phase_state, CAST('{}' AS jsonb)) "
            "|| CAST(:phase_state_update AS jsonb)"
        )
        params["phase_state_update"] = json.dumps(phase_state_update, ensure_ascii=False)

    if targets_done_inc:
        sets.append("targets_done = targets_done + :done_inc")
        params["done_inc"] = targets_done_inc

    if targets_failed_inc:
        sets.append("targets_failed = targets_failed + :failed_inc")
        params["failed_inc"] = targets_failed_inc

    if requests_count_inc:
        sets.append("requests_count = requests_count + :req_inc")
        params["req_inc"] = requests_count_inc

    if waf_blocked_inc:
        sets.append("waf_blocked_count = waf_blocked_count + :waf_inc")
        params["waf_inc"] = waf_blocked_inc

    if not sets:
        return

    db.execute(
        text(f"UPDATE cadastre_jobs SET {', '.join(sets)} WHERE job_id = :id"),
        params,
    )


def _is_job_cancelled(db: Any, job_id: int) -> bool:
    """Проверить не отменён ли job (воркер должен stop обработку)."""
    row = (
        db.execute(
            text("SELECT status FROM cadastre_jobs WHERE job_id = :id"),
            {"id": job_id},
        )
        .mappings()
        .first()
    )
    return row is not None and row["status"] == "cancelled"


# ── Основной per-quarter task ────────────────────────────────────────────────


@celery_app.task(
    bind=True,
    name="tasks.cadastre.harvest_quarter",
    acks_late=True,
    max_retries=3,
    autoretry_for=(Exception,),
    retry_backoff=True,
    retry_backoff_max=60,
    dont_autoretry_for=(NspdBulkWafError,),
    soft_time_limit=600,  # 10 мин: 225 grid-walk запросов × ~0.1с + overhead
    time_limit=900,
)
def bulk_harvest_quarter_task(self: Any, quarter: str, job_id: int) -> dict[str, Any]:
    """Celery wrapper для harvest_quarter одного кадастрового квартала.

    Sync→async bridge через asyncio.run().
    Saga state пишется в cadastre_jobs.phase_state через update_progress callback.
    WAF 403 → не retry, waf_blocked_count++.
    """
    from app.services.cadastre.bulk_harvest import harvest_quarter

    db = SessionLocal()
    try:
        # Проверить не отменён ли job
        if _is_job_cancelled(db, job_id):
            logger.info(
                "bulk_harvest_quarter_task: job=%d cancelled, skip quarter=%s",
                job_id,
                quarter,
            )
            db.commit()
            return {"quarter": quarter, "job_id": job_id, "skipped": "cancelled"}

        # Обновить heartbeat + started_at
        _update_job(
            db,
            job_id,
            status="running",
            started_at=_dt.datetime.now(_dt.UTC),
            heartbeat=True,
        )
        db.commit()

        def progress_cb(state: dict[str, Any]) -> None:
            """Callback: пишет phase_state в DB при каждом шаге фазы."""
            try:
                _update_job(db, job_id, phase_state_update=state, heartbeat=True)
                db.commit()
            except Exception as cb_err:
                logger.warning(
                    "progress_cb failed for job=%d quarter=%s: %s",
                    job_id,
                    quarter,
                    cb_err,
                )
                try:
                    db.rollback()
                except Exception:
                    pass

        # Запустить async harvest_quarter в синхронном контексте Celery
        async def _run() -> dict[str, Any]:
            from app.scrapers.nspd_bulk_client import NSPDBulkClient

            async with NSPDBulkClient() as client:
                result = await harvest_quarter(
                    db=db,
                    client=client,
                    quarter=quarter,
                    job_id=job_id,
                    update_progress=progress_cb,
                )
            return {
                "quarter": quarter,
                "job_id": job_id,
                "parcels": result.parcels_upserted,
                "buildings": result.buildings_upserted,
                "constructions": result.constructions_upserted,
                "oncs": result.oncs_upserted,
                "enks": result.enks_upserted,
                "zouit": result.zouit_upserted,
                "snapshot_requests": result.snapshot_requests,
                "grid_walk_requests": result.grid_walk_requests,
                "total_requests": result.total_requests,
            }

        result_dict = asyncio.run(_run())

        # Успех: increment targets_done. requests_count учитывает ВСЕ NSPD HTTP calls
        # (Phase 1 search + Phase 1.5 per-cat probes + Phase 2-3 grid-walk), не только
        # grid_walk_requests — иначе snapshot-only quarters давали бы 0 запросов в metrics.
        _update_job(
            db,
            job_id,
            targets_done_inc=1,
            requests_count_inc=result_dict.get("total_requests", 0),
            heartbeat=True,
        )
        # Проверить не завершён ли job (все targets done)
        _maybe_finish_job(db, job_id)
        db.commit()

        logger.info(
            "bulk_harvest_quarter_task done: quarter=%s job=%d parcels=%d buildings=%d",
            quarter,
            job_id,
            result_dict.get("parcels", 0),
            result_dict.get("buildings", 0),
        )
        return result_dict

    except NspdBulkWafError as e:
        # WAF 403 — не retry, только счётчик
        logger.warning(
            "bulk_harvest_quarter_task WAF block: quarter=%s job=%d: %s",
            quarter,
            job_id,
            e,
        )
        _update_job(
            db,
            job_id,
            waf_blocked_inc=1,
            targets_failed_inc=1,
            last_error=f"WAF: {e}",
            heartbeat=True,
        )
        db.commit()
        # NspdBulkWafError в dont_autoretry_for → не retry
        raise

    except Exception as e:
        logger.exception(
            "bulk_harvest_quarter_task failed: quarter=%s job=%d: %s",
            quarter,
            job_id,
            e,
        )
        try:
            _update_job(
                db,
                job_id,
                last_error=f"{type(e).__name__}: {e}"[:500],
                heartbeat=True,
            )
            db.commit()
        except Exception as inner_e:
            logger.warning("failed to update job state: %s", inner_e)
            try:
                db.rollback()
            except Exception:
                pass
        raise

    finally:
        db.close()


def _maybe_finish_job(db: Any, job_id: int) -> None:
    """Отметить job как done если targets_done + targets_failed == targets_total."""
    row = (
        db.execute(
            text("""
                SELECT targets_total, targets_done, targets_failed, status
                FROM cadastre_jobs WHERE job_id = :id
            """),
            {"id": job_id},
        )
        .mappings()
        .first()
    )
    if not row:
        return
    if row["status"] in ("cancelled", "done", "failed"):
        return
    total = row["targets_total"] or 0
    done = row["targets_done"] or 0
    failed = row["targets_failed"] or 0
    if total > 0 and (done + failed) >= total:
        db.execute(
            text("""
                UPDATE cadastre_jobs
                SET status = 'done', finished_at = NOW()
                WHERE job_id = :id AND status NOT IN ('cancelled', 'done', 'failed')
            """),
            {"id": job_id},
        )
        logger.info(
            "cadastre job %d finished (done=%d failed=%d total=%d)", job_id, done, failed, total
        )


# ── Enqueue task ─────────────────────────────────────────────────────────────


@celery_app.task(
    bind=True,
    name="tasks.cadastre.enqueue_harvest",
)
def enqueue_cadastre_harvest(self: Any, job_id: int) -> dict[str, Any]:
    """Прочитать cadastre_jobs.raw_targets и запустить per-quarter task'и.

    Идемпотентен: повторный вызов просто повторно enqueue'ит те же quarters.
    Это OK т.к. bulk_harvest_quarter_task idempotent (ON CONFLICT DO UPDATE +
    phase_state='done' early exit).
    """
    db = SessionLocal()
    try:
        job = (
            db.execute(
                text("SELECT * FROM cadastre_jobs WHERE job_id = :id"),
                {"id": job_id},
            )
            .mappings()
            .first()
        )
        if not job:
            logger.error("enqueue_cadastre_harvest: job_id=%d not found", job_id)
            return {"error": "job_not_found", "job_id": job_id}

        if job["status"] == "cancelled":
            return {"job_id": job_id, "skipped": "cancelled"}

        raw_targets = job["raw_targets"] or {}
        quarters: list[str] = raw_targets.get("quarters", [])

        if not quarters:
            logger.warning("enqueue_cadastre_harvest: job=%d имеет пустой список quarters", job_id)
            return {"job_id": job_id, "enqueued": 0}

        # Fix H: per-quarter freshness skip — не переобрабатывать quarters с свежим
        # cad_quarter_stats.fetched_at (последний успешный обход <FRESHNESS_HOURS).
        # raw_targets.skip_fresh_hours=N управляет порогом; 0 → отключает skip.
        # Хранится в существующем `raw_targets` JSONB (no schema change).
        # Это сэкономит ~30k NSPD calls при повторных запусках того же scope.
        skip_fresh_hours = int(raw_targets.get("skip_fresh_hours", 1) or 0)
        skipped_fresh: set[str] = set()
        if skip_fresh_hours > 0:
            threshold = _dt.datetime.now(_dt.UTC) - _dt.timedelta(hours=skip_fresh_hours)
            rows = (
                db.execute(
                    text(
                        "SELECT cad_number FROM cad_quarter_stats "
                        "WHERE cad_number = ANY(CAST(:qs AS text[])) "
                        "AND fetched_at >= :t"
                    ),
                    {"qs": quarters, "t": threshold},
                )
                .scalars()
                .all()
            )
            skipped_fresh = set(rows)
            if skipped_fresh:
                # Корректируем targets_total — skipped quarters не входят в прогресс
                db.execute(
                    text(
                        "UPDATE cadastre_jobs SET targets_total = :new_total, "
                        "targets_skipped = COALESCE(targets_skipped, 0) + :sk "
                        "WHERE job_id = :id"
                    ),
                    {
                        "new_total": len(quarters) - len(skipped_fresh),
                        "sk": len(skipped_fresh),
                        "id": job_id,
                    },
                )

        # Обновить статус на running
        _update_job(db, job_id, status="running", started_at=_dt.datetime.now(_dt.UTC))
        db.commit()

        enqueued = 0
        for quarter in quarters:
            if quarter in skipped_fresh:
                continue
            bulk_harvest_quarter_task.apply_async(
                args=[quarter, job_id],
                queue=CADASTRE_QUEUE,
            )
            enqueued += 1

        logger.info(
            "enqueue_cadastre_harvest: job=%d enqueued %d (skipped_fresh=%d, threshold=%dh)",
            job_id,
            enqueued,
            len(skipped_fresh),
            skip_fresh_hours,
        )
        return {"job_id": job_id, "enqueued": enqueued, "skipped_fresh": len(skipped_fresh)}

    finally:
        db.close()


# ── Zombie cleanup ───────────────────────────────────────────────────────────


@celery_app.task(name="tasks.cadastre.cleanup_zombies")
def cleanup_cadastre_zombies() -> dict[str, Any]:
    """Periodic: найти stale cadastre_jobs (heartbeat > 5 мин) и переставить в паузу.

    Запускается через beat (beat_schedule.py). Idempotent.
    """
    db = SessionLocal()
    stale_ids: list[int] = []
    try:
        rows = (
            db.execute(
                text("""
                    UPDATE cadastre_jobs
                    SET status = 'paused',
                        error = COALESCE(error, 'auto-paused: stale heartbeat')
                    WHERE status = 'running'
                      AND heartbeat_at < NOW() - INTERVAL '5 minutes'
                    RETURNING job_id
                """)
            )
            .mappings()
            .all()
        )
        db.commit()
        stale_ids = [int(r["job_id"]) for r in rows]
        if stale_ids:
            logger.warning("cleanup_cadastre_zombies: paused stale jobs: %s", stale_ids)
    except Exception as e:
        logger.exception("cleanup_cadastre_zombies failed: %s", e)
        try:
            db.rollback()
        except Exception:
            pass
    finally:
        db.close()
    return {"stale_paused": stale_ids}