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

"""Celery tasks для NSPD quarter harvest.

Архитектура (Sprint 1.1 item #3):
    harvest_quarter      — один квартал: search_by_quarter → UPSERT nspd_quarter_dumps.
    harvest_stale_quarters — beat-triggered bulk: ищет устаревшие/отсутствующие дампы,
                             фанаут harvest_quarter.apply_async() на каждый.

Sanity check при startup:
    worker_ready сигнал проверяет наличие таблицы nspd_quarter_dumps. Если отсутствует
    (migration 88_nspd_quarter_dumps.sql не применена) — пишет critical log, не падает.

UPSERT pattern:
    ON CONFLICT (quarter_cad) DO UPDATE — обновляет все колонки кроме region_code
    (иммутабельный per-row, задаётся при первом insert).
    Геометрия квартала: ST_Multi(ST_Transform(ST_SetSRID(ST_GeomFromGeoJSON(...),3857),4326))
    Bbox: ST_MakeEnvelope(xmin,ymin,xmax,ymax,3857).
"""

from __future__ import annotations

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

from sqlalchemy import text

from app.core.db import SessionLocal
from app.services.scrapers.nspd_client import (
    NSPDClient,
    NSPDFeature,
    NspdLiteWafError,
    QuarterDump,
)
from app.workers.celery_app import celery_app

logger = logging.getLogger(__name__)


# ── Geometry helpers ──────────────────────────────────────────────────────────


def _build_features_json(dump: QuarterDump) -> list[dict[str, Any]]:
    """Собрать все NSPDFeature из QuarterDump в плоский list[dict] для JSONB.

    Каждый элемент: {layer, feature_id, geometry, properties}.
    Порядок: core layers → zouit layers → risk layers (для консистентности).
    """

    def _feat_to_dict(layer: str, f: NSPDFeature) -> dict[str, Any]:
        return {
            "layer": layer,
            "feature_id": f.feature_id,
            "geometry": f.geometry,
            "properties": f.properties,
        }

    out: list[dict[str, Any]] = []

    # Core layers
    for feat in dump.parcels:
        out.append(_feat_to_dict("parcels", feat))
    for feat in dump.buildings:
        out.append(_feat_to_dict("buildings", feat))
    for feat in dump.territorial_zones:
        out.append(_feat_to_dict("territorial_zones", feat))
    for feat in dump.red_lines:
        out.append(_feat_to_dict("red_lines", feat))
    for feat in dump.engineering_structures:
        out.append(_feat_to_dict("engineering_structures", feat))

    # ЗОУИТ groups — keys: okn, engineering, natural, protected, other
    for short_name, features in dump.zouit.items():
        layer_name = f"zouit_{short_name}"
        for feat in features:
            out.append(_feat_to_dict(layer_name, feat))

    # Risk groups — keys: flooding_underground, flooding, swampification, ...
    for short_name, features in dump.risks.items():
        layer_name = f"risk_{short_name}"
        for feat in features:
            out.append(_feat_to_dict(layer_name, feat))

    return out


def _build_zouit_count(dump: QuarterDump) -> int:
    """Сумма features по всем 5 ЗОУИТ группам."""
    return sum(len(v) for v in dump.zouit.values())


def _build_risks_count(dump: QuarterDump) -> int:
    """Сумма features по всем risk группам."""
    return sum(len(v) for v in dump.risks.values())


# ── UPSERT helper ─────────────────────────────────────────────────────────────

_UPSERT_SQL = text(
    """
    INSERT INTO nspd_quarter_dumps (
        quarter_cad, quarter_geom, bbox_3857,
        parcels_count, buildings_count, territorial_zones_count,
        red_lines_count, engineering_count, zouit_count, risks_count, total_features,
        features_json, layers_fetched, fetched_at_utc, harvest_duration_ms,
        harvest_error, region_code
    ) VALUES (
        :quarter_cad,
        CASE WHEN :geom_json IS NULL THEN NULL
             ELSE ST_Multi(ST_Transform(ST_SetSRID(ST_GeomFromGeoJSON(:geom_json), 3857), 4326))
        END,
        CASE WHEN :bbox_xmin IS NULL THEN NULL
             ELSE ST_MakeEnvelope(:bbox_xmin, :bbox_ymin, :bbox_xmax, :bbox_ymax, 3857)
        END,
        :parcels_count, :buildings_count, :territorial_zones_count,
        :red_lines_count, :engineering_count, :zouit_count, :risks_count, :total_features,
        CAST(:features_json AS jsonb),
        CAST(:layers_fetched AS text[]),
        :fetched_at_utc::timestamptz,
        :harvest_duration_ms,
        :harvest_error,
        :region_code
    )
    ON CONFLICT (quarter_cad) DO UPDATE SET
        quarter_geom            = EXCLUDED.quarter_geom,
        bbox_3857               = EXCLUDED.bbox_3857,
        parcels_count           = EXCLUDED.parcels_count,
        buildings_count         = EXCLUDED.buildings_count,
        territorial_zones_count = EXCLUDED.territorial_zones_count,
        red_lines_count         = EXCLUDED.red_lines_count,
        engineering_count       = EXCLUDED.engineering_count,
        zouit_count             = EXCLUDED.zouit_count,
        risks_count             = EXCLUDED.risks_count,
        total_features          = EXCLUDED.total_features,
        features_json           = EXCLUDED.features_json,
        layers_fetched          = EXCLUDED.layers_fetched,
        fetched_at_utc          = EXCLUDED.fetched_at_utc,
        harvest_duration_ms     = EXCLUDED.harvest_duration_ms,
        harvest_error           = EXCLUDED.harvest_error
    -- region_code исключён из UPDATE: иммутабельное свойство квартала
    """
)


def _upsert_dump(
    quarter_cad: str,
    region_code: int,
    dump: QuarterDump | None,
    features_json: list[dict[str, Any]] | None,
    duration_ms: int | None,
    harvest_error: str | None,
) -> None:
    """UPSERT строки в nspd_quarter_dumps.

    dump=None при partial-failure: вставляем строку с harvest_error + нулевыми счётчиками.
    """
    db = SessionLocal()
    try:
        if dump is not None:
            geom_json: str | None = (
                json.dumps(dump.quarter.geometry, ensure_ascii=False)
                if dump.quarter and dump.quarter.geometry
                else None
            )
            bbox = dump.bbox_3857
            params: dict[str, Any] = {
                "quarter_cad": quarter_cad,
                "geom_json": geom_json,
                "bbox_xmin": bbox[0] if bbox else None,
                "bbox_ymin": bbox[1] if bbox else None,
                "bbox_xmax": bbox[2] if bbox else None,
                "bbox_ymax": bbox[3] if bbox else None,
                "parcels_count": len(dump.parcels),
                "buildings_count": len(dump.buildings),
                "territorial_zones_count": len(dump.territorial_zones),
                "red_lines_count": len(dump.red_lines),
                "engineering_count": len(dump.engineering_structures),
                "zouit_count": _build_zouit_count(dump),
                "risks_count": _build_risks_count(dump),
                "total_features": dump.total_features,
                "features_json": json.dumps(features_json or [], ensure_ascii=False),
                "layers_fetched": list(dump.layers_fetched),
                "fetched_at_utc": dump.fetched_at_utc,
                "harvest_duration_ms": duration_ms,
                "harvest_error": harvest_error,
                "region_code": region_code,
            }
        else:
            # Error-only row: нулевые счётчики, harvest_error filled
            params = {
                "quarter_cad": quarter_cad,
                "geom_json": None,
                "bbox_xmin": None,
                "bbox_ymin": None,
                "bbox_xmax": None,
                "bbox_ymax": None,
                "parcels_count": 0,
                "buildings_count": 0,
                "territorial_zones_count": 0,
                "red_lines_count": 0,
                "engineering_count": 0,
                "zouit_count": 0,
                "risks_count": 0,
                "total_features": 0,
                "features_json": "[]",
                "layers_fetched": [],
                "fetched_at_utc": _dt.datetime.now(_dt.UTC).isoformat(),
                "harvest_duration_ms": duration_ms,
                "harvest_error": harvest_error,
                "region_code": region_code,
            }

        db.execute(_UPSERT_SQL, params)
        db.commit()
    except Exception:
        db.rollback()
        raise
    finally:
        db.close()


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


@celery_app.task(
    bind=True,
    name="tasks.nspd_sync.harvest_quarter",
    max_retries=3,
    soft_time_limit=120,  # 2 мин — worst-case 22 HTTP × ~600ms = ~13s + margin
    autoretry_for=(NspdLiteWafError,),
    retry_backoff=True,
    retry_backoff_max=120,
    retry_jitter=True,
)
def harvest_quarter(
    self: Any,
    quarter_cad: str,
    region_code: int = 66,
    include_zouit: bool = True,
    include_risks: bool = False,
) -> dict[str, Any]:
    """Single-quarter harvest. NSPDClient.search_by_quarter → UPSERT nspd_quarter_dumps.

    Идемпотентен: повторный вызов обновляет строку (ON CONFLICT DO UPDATE).
    WAF 403/429 → autoretry с exponential backoff (max 3 попытки).
    Другие исключения → запись harvest_error в строку, return error dict (не raise).
    """
    t0 = time.monotonic()
    logger.info(
        "harvest_quarter start: cad=%s region=%d include_zouit=%s include_risks=%s",
        quarter_cad,
        region_code,
        include_zouit,
        include_risks,
    )

    client = NSPDClient()
    dump: QuarterDump | None = None
    features_json: list[dict[str, Any]] | None = None

    try:
        dump = client.search_by_quarter(
            quarter_cad,
            include_zouit=include_zouit,
            include_risks=include_risks,
        )
        features_json = _build_features_json(dump)
        duration_ms = int((time.monotonic() - t0) * 1000)

        _upsert_dump(quarter_cad, region_code, dump, features_json, duration_ms, None)

        logger.info(
            "harvest_quarter done: cad=%s region=%d duration=%dms total=%d",
            quarter_cad,
            region_code,
            duration_ms,
            dump.total_features,
        )
        return {
            "quarter_cad": quarter_cad,
            "region_code": region_code,
            "total_features": dump.total_features,
            "parcels_count": len(dump.parcels),
            "buildings_count": len(dump.buildings),
            "harvest_duration_ms": duration_ms,
            "layers_fetched": list(dump.layers_fetched),
            "bbox_3857": dump.bbox_3857,
        }

    except NspdLiteWafError:
        # autoretry_for перехватит и сделает retry с backoff.
        # Просто re-raise чтобы Celery знал о WAF.
        duration_ms = int((time.monotonic() - t0) * 1000)
        logger.warning(
            "harvest_quarter WAF: cad=%s attempt=%d duration=%dms",
            quarter_cad,
            self.request.retries,
            duration_ms,
        )
        raise

    except Exception as exc:
        duration_ms = int((time.monotonic() - t0) * 1000)
        err_str = f"{type(exc).__name__}: {exc}"
        logger.error(
            "harvest_quarter error: cad=%s region=%d duration=%dms error=%s",
            quarter_cad,
            region_code,
            duration_ms,
            err_str,
        )
        # Пишем строку с harvest_error — не даём silent failure.
        try:
            _upsert_dump(quarter_cad, region_code, None, None, duration_ms, err_str[:500])
        except Exception as db_exc:
            logger.error(
                "harvest_quarter: failed to write error row for cad=%s: %s",
                quarter_cad,
                db_exc,
            )
        return {
            "quarter_cad": quarter_cad,
            "region_code": region_code,
            "error": err_str,
            "harvest_duration_ms": duration_ms,
        }


# ── Bulk beat-triggered task ──────────────────────────────────────────────────


@celery_app.task(name="tasks.nspd_sync.harvest_stale_quarters")
def harvest_stale_quarters(
    region_code: int = 66,
    max_age_days: int = 90,
    batch_size: int = 50,
) -> dict[str, Any]:
    """Beat-triggered: найти устаревшие/отсутствующие дампы и сфанаутить harvest_quarter.

    Стратегия:
        Берём кадастровые кварталы из cad_quarters_geom WHERE region_code=:rc
        которых нет в nspd_quarter_dumps (или есть, но старше max_age_days суток
        или имеют harvest_error).
        Limit batch_size — чтобы не сжечь весь rate-limit за один запуск beat.
        Fanout: harvest_quarter.apply_async(args=[cad, region_code]).

    Идемпотентен: следующий запуск подхватит остатки если batch_size < stale total.
    """
    logger.info(
        "harvest_stale_quarters: region=%d max_age_days=%d batch_size=%d",
        region_code,
        max_age_days,
        batch_size,
    )

    db = SessionLocal()
    stale_cads: list[str] = []
    try:
        rows = db.execute(
            text(
                """
                SELECT cad_number
                  FROM cad_quarters_geom
                 WHERE cad_number NOT IN (
                       SELECT quarter_cad
                         FROM nspd_quarter_dumps
                        WHERE region_code = :rc
                          AND fetched_at_utc > NOW() - INTERVAL '1 day' * :days
                          AND harvest_error IS NULL
                   )
                 ORDER BY cad_number
                 LIMIT :lim
                """
            ),
            {"rc": region_code, "days": max_age_days, "lim": batch_size},
        ).all()
        stale_cads = [r[0] for r in rows]
    except Exception as e:
        logger.error("harvest_stale_quarters: DB query failed: %s", e)
        try:
            db.rollback()
        except Exception:
            pass
        return {"error": str(e), "enqueued": 0}
    finally:
        db.close()

    enqueued = 0
    for cad in stale_cads:
        try:
            harvest_quarter.apply_async(args=[cad, region_code])
            enqueued += 1
        except Exception as e:
            logger.warning("harvest_stale_quarters: enqueue failed for cad=%s: %s", cad, e)

    logger.info(
        "harvest_stale_quarters done: region=%d enqueued=%d stale_found=%d",
        region_code,
        enqueued,
        len(stale_cads),
    )
    return {
        "enqueued": enqueued,
        "stale_found": len(stale_cads),
        "region_code": region_code,
        "max_age_days": max_age_days,
    }