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

"""Celery task wrapper for naш.дом.рф kn-API sweeps."""

from __future__ import annotations

import asyncio
import logging
import random
import time
from contextlib import contextmanager
from pathlib import Path
from typing import Any

import redis

from app.core.config import settings
from app.services.scrapers.domrf_kn import run_region_sweep
from app.workers.celery_app import celery_app

logger = logging.getLogger(__name__)

# Lock TTL — full Свердл sweep с extras занимает ~25-30 мин. 30 мин — потолок.
# Если worker умер с lock'ом — он истечёт за полчаса, а не за 2 часа (раньше).
# Heartbeat-based autoclean в /queue endpoint всё равно подхватит зомби-runs.
_LOCK_TTL_SECONDS = 30 * 60


def _lock_key(region_code: int, developers: list[str] | None) -> str:
    devs_key = ",".join(developers) if developers else "*"
    return f"scrape:kn:lock:{region_code}:{devs_key}"


def _log_task_received(region_code: int, devs: list[str] | None, task_id: str | None) -> None:
    """Записать в kn_scrape_log что worker получил задачу. Без run_id —
    он появится позже в run_region_sweep. Best-effort: SQL ошибки гасим."""
    from sqlalchemy import text

    from app.core.db import SessionLocal

    db = SessionLocal()
    try:
        db.execute(
            text(
                """
                INSERT INTO kn_scrape_log (run_id, level, stage, message)
                VALUES (NULL, 'info', 'task_received', :msg)
                """
            ),
            {
                "msg": (
                    f"Worker подхватил task_id={task_id or '?'}, region={region_code},"
                    f" devs={devs or '*'}"
                )[:500],
            },
        )
        db.commit()
    except Exception as e:
        logger.warning("log_task_received failed: %s", e)
    finally:
        db.close()


def _log_task_skipped(
    region_code: int, devs: list[str] | None, task_id: str | None, reason: str
) -> None:
    from sqlalchemy import text

    from app.core.db import SessionLocal

    db = SessionLocal()
    try:
        db.execute(
            text(
                """
                INSERT INTO kn_scrape_log (run_id, level, stage, message)
                VALUES (NULL, 'warn', 'task_skipped', :msg)
                """
            ),
            {
                "msg": (
                    f"task_id={task_id or '?'} skipped: {reason}, region={region_code},"
                    f" devs={devs or '*'}"
                )[:500],
            },
        )
        db.commit()
    except Exception as e:
        logger.warning("log_task_skipped failed: %s", e)
    finally:
        db.close()


def force_release_lock(region_code: int, developers: list[str] | None) -> bool:
    """Принудительно удалить Redis-lock (для emergency manual trigger,
    когда зомби-worker оставил key). Возвращает True если ключ был."""
    key = _lock_key(region_code, developers)
    r = redis.Redis.from_url(settings.redis_url)
    try:
        return bool(r.delete(key))
    except Exception as e:
        logger.warning("force_release_lock %s failed: %s", key, e)
        return False


@contextmanager
def _region_lock(region_code: int, developers: list[str] | None):
    """Redis singleton lock so beat double-fire or two workers can't run the
    same (region, devs) sweep concurrently. SET NX EX is atomic."""
    key = _lock_key(region_code, developers)
    r = redis.Redis.from_url(settings.redis_url)
    acquired = r.set(key, "1", nx=True, ex=_LOCK_TTL_SECONDS)
    try:
        yield bool(acquired)
    finally:
        if acquired:
            try:
                r.delete(key)
            except Exception:
                pass


@celery_app.task(bind=True, name="tasks.scrape_kn.scrape_kn_region", max_retries=2)
def scrape_kn_region(
    self: Any,
    region_code: int,
    developers: list[str] | None = None,
    skip_jitter: bool = False,
    extras: bool = True,
    download_photos: bool = False,
    fetch_flats: bool = True,
) -> dict[str, Any]:
    """Run a full kn-API sweep for one region. Returns summary dict.

    For scheduled runs (skip_jitter=False), sleeps a random 0..SCRAPE_KN_JITTER_SECONDS
    interval before starting so multiple workers / repeated weekly cycles don't hit
    DOM.РФ at exactly the same minute. Manual triggers pass skip_jitter=True for
    immediate execution.
    """
    # Сразу пишем «task received» в kn_scrape_log БЕЗ run_id (run_id ещё не
    # создан — он появится в run_region_sweep). Это позволяет UI видеть
    # «worker подхватил задачу» даже если потом сработает lock_held или
    # упадём ещё до Phase A.
    _log_task_received(region_code, developers, self.request.id)

    if not skip_jitter and settings.scrape_kn_jitter_seconds > 0:
        delay = random.uniform(0, settings.scrape_kn_jitter_seconds)
        logger.info("jitter delay %.0fs before scrape", delay)
        time.sleep(delay)

    state_path: str | None = settings.scrape_kn_state_path or None
    if state_path and not Path(state_path).exists():
        logger.warning("storage_state %s not found — bootstrap will run cold", state_path)
        state_path = None

    with _region_lock(region_code, developers) as got_lock:
        if not got_lock:
            logger.warning(
                "scrape_kn_region SKIPPED region=%s devs=%s — another sweep is"
                " already running (Redis lock held)",
                region_code,
                developers,
            )
            _log_task_skipped(region_code, developers, self.request.id, "lock_held")
            return {
                "skipped": True,
                "reason": "lock_held",
                "region_code": region_code,
                "developers": developers,
                "lock_key": _lock_key(region_code, developers),
            }

        logger.info(
            "scrape_kn_region start region=%s devs=%s extras=%s photos=%s state=%s",
            region_code,
            developers,
            extras,
            download_photos,
            state_path,
        )
        return asyncio.run(
            run_region_sweep(
                region_code=region_code,
                developers=developers,
                load_state=state_path,
                fetch_flats=fetch_flats,
                extras=extras,
                download_photos_binary=download_photos,
            )
        )


@celery_app.task(bind=True, name="tasks.scrape_kn.resume_kn_run", max_retries=2)
def resume_kn_run(self: Any, run_id: int) -> dict[str, Any]:
    """Resume a previously interrupted sweep using objects_snapshot +
    progress_obj_index from kn_scrape_runs. Triggered by worker_ready hook
    when a stale 'running' row is found (heartbeat >5 min)."""
    from sqlalchemy import text

    from app.core.db import SessionLocal

    state_path: str | None = settings.scrape_kn_state_path or None
    if state_path and not Path(state_path).exists():
        state_path = None

    # Reconstruct region/devs from the stale row to scope the lock correctly.
    db = SessionLocal()
    try:
        row = (
            db.execute(
                text(
                    """
                    SELECT region_codes, developer_ids
                      FROM kn_scrape_runs
                     WHERE run_id = :rid
                    """
                ),
                {"rid": run_id},
            )
            .mappings()
            .first()
        )
    finally:
        db.close()
    if not row:
        return {"skipped": True, "reason": "run_not_found", "run_id": run_id}

    region_codes = list(row["region_codes"] or [])
    developers = list(row["developer_ids"]) if row["developer_ids"] else None
    if not region_codes:
        return {"skipped": True, "reason": "no_region", "run_id": run_id}
    region_code = region_codes[0]

    _log_task_received(region_code, developers, self.request.id)

    with _region_lock(region_code, developers) as got_lock:
        if not got_lock:
            logger.warning(
                "resume_kn_run SKIPPED run=%s region=%s — lock held by another sweep",
                run_id,
                region_code,
            )
            _log_task_skipped(region_code, developers, self.request.id, "lock_held")
            return {"skipped": True, "reason": "lock_held", "run_id": run_id}

        logger.info("resume_kn_run start run=%s region=%s devs=%s", run_id, region_code, developers)
        return asyncio.run(
            run_region_sweep(
                region_code=region_code,
                developers=developers,
                load_state=state_path,
                resume_from_run_id=run_id,
            )
        )