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

"""Bulk NSPD geo backfill — resume-friendly Celery task.

Архитектура (см. data/sql/77_nspd_geo_jobs.sql):
    nspd_geo_jobs    — журнал заданий (один job = одна bulk-операция)
    nspd_geo_targets — список cad-номеров со status (pending/done/failed)

Resume:
    1. При worker_ready signal: ищем jobs со status='running' и stale heartbeat
       (>10 мин), переключаем их статус на 'resume_needed' + re-enqueue task.
    2. process_nspd_geo_job(job_id) сразу читает SELECT WHERE status='pending'
       — идемпотентно продолжает с того места, где остановился (UPSERT в
       cad_quarters_geom / cad_buildings гарантирует no dup'ов).

Запуск (3 пути):
    1. UI: POST /api/v1/admin/scrape/geo (создаёт job, наполняет targets, enqueue task)
    2. CLI: python -m app.workers.tasks.nspd_geo --enqueue-region 66
    3. Beat: можно добавить cron-расписание в celery_app.py
"""

from __future__ import annotations

import logging
from typing import Any

from sqlalchemy import text
from sqlalchemy.orm import Session

from app.core.db import SessionLocal
from app.services.scrapers.nspd_lite import (
    NspdLiteError,
    NspdLiteWafError,
    fetch_geoportal,
    fetch_via_rosreestr2coord,
)
from app.workers.celery_app import celery_app

logger = logging.getLogger(__name__)

# Heartbeat каждые N items (компромисс: чаще = выше DB write, реже = stale-window больше)
HEARTBEAT_EVERY = 5
# WAF backoff (seconds) после первого 403/429
WAF_BACKOFF_BASE_S = 30


# ── Helpers для job/target lifecycle ────────────────────────────────────────


def _start_job(db: Session, job_id: int) -> None:
    db.execute(
        text(
            """
            UPDATE nspd_geo_jobs
               SET status = 'running',
                   started_at = COALESCE(started_at, NOW()),
                   heartbeat_at = NOW(),
                   updated_at = NOW()
             WHERE job_id = :id
            """
        ),
        {"id": job_id},
    )
    db.commit()


def _heartbeat(db: Session, job_id: int, **counts: int) -> None:
    sets = ["heartbeat_at = NOW()", "updated_at = NOW()"]
    params: dict[str, Any] = {"id": job_id}
    for k, v in counts.items():
        sets.append(f"{k} = :{k}")
        params[k] = v
    try:
        db.execute(
            text(f"UPDATE nspd_geo_jobs SET {', '.join(sets)} WHERE job_id = :id"),
            params,
        )
        db.commit()
    except Exception as e:
        logger.warning("heartbeat failed for job=%s: %s", job_id, e)
        try:
            db.rollback()
        except Exception:
            pass


def _finish_job(db: Session, job_id: int, status: str, error: str | None = None) -> None:
    db.execute(
        text(
            """
            UPDATE nspd_geo_jobs
               SET status = :st, finished_at = NOW(), error = :err, updated_at = NOW()
             WHERE job_id = :id
            """
        ),
        {"id": job_id, "st": status, "err": error},
    )
    db.commit()


# ── Persistence для разных видов targets ────────────────────────────────────


def _save_quarter(db: Session, payload: dict, cad_num: str) -> int:
    """Парсит NSPD-response для квартала и UPSERT'ит в cad_quarters_geom."""
    feats = (payload.get("data") or {}).get("features") or []
    if not feats:
        return 0
    f = feats[0]
    geom_geojson = f.get("geometry")
    if not geom_geojson:
        return 0
    db.execute(
        text(
            """
            INSERT INTO cad_quarters_geom (cad_number, geom, raw_props, fetched_at)
            VALUES (:cad, ST_Transform(ST_SetSRID(ST_GeomFromGeoJSON(:g), 3857), 4326),
                    CAST(:props AS jsonb), NOW())
            ON CONFLICT (cad_number) DO UPDATE
               SET geom = EXCLUDED.geom,
                   raw_props = EXCLUDED.raw_props,
                   fetched_at = NOW()
            """
        ),
        {
            "cad": cad_num,
            "g": str(geom_geojson).replace("'", '"'),
            "props": str(f.get("properties", {})).replace("'", '"'),
        },
    )
    return 1


def _save_building(db: Session, payload: dict, cad_num: str) -> int:
    """Парсит NSPD-response для здания и UPSERT'ит в cad_buildings."""
    feats = (payload.get("data") or {}).get("features") or []
    if not feats:
        return 0
    n = 0
    for f in feats:
        geom_geojson = f.get("geometry")
        opts = (f.get("properties") or {}).get("options") or {}
        if not geom_geojson:
            continue
        # Только этого здания (NSPD может вернуть здания всего квартала)
        if opts.get("cad_num") != cad_num:
            continue
        import json as _json

        db.execute(
            text(
                """
                INSERT INTO cad_buildings (
                    cad_num, quarter_cad_num, geom, purpose, building_name,
                    readable_address, area, floors, raw_props
                ) VALUES (
                    :cad, :qcad,
                    ST_Transform(ST_SetSRID(ST_GeomFromGeoJSON(:g), 3857), 4326),
                    :purpose, :name, :addr, :area, :floors, CAST(:props AS jsonb)
                )
                ON CONFLICT (cad_num) DO UPDATE
                   SET quarter_cad_num = EXCLUDED.quarter_cad_num,
                       geom = EXCLUDED.geom,
                       purpose = EXCLUDED.purpose,
                       building_name = EXCLUDED.building_name,
                       readable_address = EXCLUDED.readable_address,
                       area = EXCLUDED.area,
                       floors = EXCLUDED.floors,
                       raw_props = EXCLUDED.raw_props
                """
            ),
            {
                "cad": cad_num,
                "qcad": opts.get("quarter_cad_number"),
                "g": _json.dumps(geom_geojson, ensure_ascii=False),
                "purpose": opts.get("purpose"),
                "name": opts.get("building_name") or opts.get("name"),
                "addr": opts.get("readable_address"),
                "area": opts.get("area"),
                "floors": opts.get("floors"),
                "props": _json.dumps(opts, ensure_ascii=False),
            },
        )
        n += 1
    return n


def _persist_target(db: Session, thematic_id: int, payload: dict, cad_num: str) -> tuple[int, str]:
    """Returns (n_features_saved, table_name). Полиморфно по thematic_id."""
    if thematic_id == 2:
        return _save_quarter(db, payload, cad_num), "cad_quarters_geom"
    if thematic_id == 5:
        return _save_building(db, payload, cad_num), "cad_buildings"
    # Для parcel (1) пока нет нашей normalized-таблицы — только raw в targets.result_features
    return 0, ""


# ── Helper для создания job + targets ──────────────────────────────────────


def enqueue_geo_job(
    *,
    name: str | None,
    job_kind: str,
    source_kind: str,
    source_params: dict[str, Any],
    cad_nums_with_thematic: list[tuple[str, int]],
    rate_ms: int = 600,
    use_rosreestr2coord: bool = False,
    triggered_by: str = "manual",
) -> int:
    """Создаёт job + targets и возвращает job_id. Не enqueue'ит task сам.
    Caller должен сделать `process_nspd_geo_job.apply_async(args=[job_id])`.
    """
    import json as _json

    db = SessionLocal()
    try:
        job_id = db.execute(
            text(
                """
                INSERT INTO nspd_geo_jobs (
                    name, job_kind, source_kind, source_params,
                    rate_ms, use_rosreestr2coord, triggered_by, status, targets_total
                ) VALUES (:n, :k, :sk, CAST(:sp AS jsonb), :r, :ur, :tb, 'queued', :tt)
                RETURNING job_id
                """
            ),
            {
                "n": name,
                "k": job_kind,
                "sk": source_kind,
                "sp": _json.dumps(source_params, ensure_ascii=False),
                "r": rate_ms,
                "ur": use_rosreestr2coord,
                "tb": triggered_by,
                "tt": len(cad_nums_with_thematic),
            },
        ).scalar_one()

        # Bulk INSERT targets
        if cad_nums_with_thematic:
            values = ",".join(
                f"({job_id}, '{c.replace(chr(39), chr(39)*2)}', {t}, 'pending')"
                for c, t in cad_nums_with_thematic
            )
            db.execute(
                text(
                    f"INSERT INTO nspd_geo_targets (job_id, cad_num, thematic_id, status) "
                    f"VALUES {values} ON CONFLICT (job_id, cad_num, thematic_id) DO NOTHING"
                )
            )
        db.commit()
        return int(job_id)
    finally:
        db.close()


# ── Main task ───────────────────────────────────────────────────────────────


@celery_app.task(
    bind=True,
    name="tasks.nspd_geo.process_job",
    max_retries=0,  # resume через worker_ready, не через retry
    soft_time_limit=3600 * 6,  # 6 часов max — крупные jobs
)
def process_nspd_geo_job(self: Any, job_id: int) -> dict[str, Any]:
    """Идёт по nspd_geo_targets WHERE job_id=:id AND status='pending'.

    Heartbeat каждые HEARTBEAT_EVERY items, persistent counters в job-row.
    Idempotent: уже-сохранённые UPSERT'ятся, status target → 'done'.
    """
    import time

    db = SessionLocal()
    try:
        # Загружаем параметры job'а
        job = (
            db.execute(
                text("SELECT * FROM nspd_geo_jobs WHERE job_id = :id"),
                {"id": job_id},
            )
            .mappings()
            .first()
        )
        if not job:
            return {"error": "job_not_found", "job_id": job_id}

        _start_job(db, job_id)
        rate_ms = int(job["rate_ms"] or 600)
        use_lib = bool(job["use_rosreestr2coord"])

        # Текущие счётчики (для resume — продолжаем с тех значений)
        done = int(job["targets_done"] or 0)
        failed = int(job["targets_failed"] or 0)
        skipped = int(job["targets_skipped"] or 0)
        n_requests = int(job["requests_count"] or 0)
        n_waf = int(job["waf_blocked_count"] or 0)
        consecutive_waf = 0

        logger.info(
            "process_nspd_geo_job start: job=%s use_lib=%s rate=%dms", job_id, use_lib, rate_ms
        )

        while True:
            target = (
                db.execute(
                    text(
                        """
                    SELECT target_id, cad_num, thematic_id, attempts
                      FROM nspd_geo_targets
                     WHERE job_id = :id AND status = 'pending'
                     ORDER BY target_id
                     LIMIT 1
                    """
                    ),
                    {"id": job_id},
                )
                .mappings()
                .first()
            )
            if not target:
                break

            tid = target["thematic_id"]
            cad = target["cad_num"]

            try:
                if use_lib:
                    payload = {
                        "data": {
                            "features": [fetch_via_rosreestr2coord(cad, area_type=tid, timeout=15)]
                        }
                    }
                    payload["data"]["features"] = [f for f in payload["data"]["features"] if f]
                else:
                    payload = fetch_geoportal(cad, thematic_id=tid, rate_ms=rate_ms)
                n_requests += 1
                consecutive_waf = 0

                n_feats, saved_to = _persist_target(db, tid, payload, cad)
                db.execute(
                    text(
                        """
                        UPDATE nspd_geo_targets
                           SET status = 'done', result_features = :f,
                               saved_to_table = :s, fetched_at = NOW(),
                               attempts = attempts + 1
                         WHERE target_id = :tid
                        """
                    ),
                    {"f": n_feats, "s": saved_to or None, "tid": target["target_id"]},
                )
                db.commit()
                done += 1

            except NspdLiteWafError as e:
                consecutive_waf += 1
                n_waf += 1
                wait = WAF_BACKOFF_BASE_S * min(consecutive_waf, 8)
                logger.warning(
                    "WAF on %s (consecutive=%d), backoff %ds: %s",
                    cad,
                    consecutive_waf,
                    wait,
                    e,
                )
                # НЕ помечаем failed — остаётся pending для следующей итерации
                _heartbeat(db, job_id, requests_count=n_requests, waf_blocked_count=n_waf)
                time.sleep(wait)
                if consecutive_waf >= 8:
                    # 8 подряд WAF — пауза job'а до следующего worker_ready
                    _finish_job(
                        db,
                        job_id,
                        "paused",
                        error=f"WAF blocked {consecutive_waf} consecutive",
                    )
                    return {"job_id": job_id, "paused": True, "reason": "waf"}
                continue

            except (NspdLiteError, Exception) as e:
                logger.warning("fetch failed for %s: %s", cad, e)
                db.execute(
                    text(
                        """
                        UPDATE nspd_geo_targets
                           SET status = CASE WHEN attempts >= 2 THEN 'failed' ELSE 'pending' END,
                               attempts = attempts + 1,
                               error_msg = :em
                         WHERE target_id = :tid
                        """
                    ),
                    {"em": str(e)[:500], "tid": target["target_id"]},
                )
                db.commit()
                if int(target["attempts"]) >= 2:
                    failed += 1

            if (done + failed + skipped) % HEARTBEAT_EVERY == 0:
                _heartbeat(
                    db,
                    job_id,
                    targets_done=done,
                    targets_failed=failed,
                    targets_skipped=skipped,
                    requests_count=n_requests,
                    waf_blocked_count=n_waf,
                )

        _heartbeat(
            db,
            job_id,
            targets_done=done,
            targets_failed=failed,
            targets_skipped=skipped,
            requests_count=n_requests,
            waf_blocked_count=n_waf,
        )
        # done даже при failed targets — UI покажет failed_count отдельно
        _finish_job(db, job_id, "done")
        logger.info(
            "process_nspd_geo_job done: job=%s done=%d failed=%d req=%d waf=%d",
            job_id,
            done,
            failed,
            n_requests,
            n_waf,
        )
        return {
            "job_id": job_id,
            "done": done,
            "failed": failed,
            "skipped": skipped,
            "requests": n_requests,
            "waf_blocked": n_waf,
        }
    except Exception as e:
        logger.exception("process_nspd_geo_job crashed: job=%s: %s", job_id, e)
        try:
            _finish_job(db, job_id, "failed", error=f"{type(e).__name__}: {e}")
        except Exception:
            pass
        raise
    finally:
        db.close()