gendesign/backend/app/workers/celery_app.py

"""Celery app + beat schedule."""

import logging

from celery import Celery
from celery.schedules import crontab
from celery.signals import worker_ready

from app.core.config import settings

logger = logging.getLogger(__name__)


def _parse_cron(spec: str) -> crontab:
    """Parse 'M H DoM Mon DoW' crontab string into Celery crontab. Empty fields default to '*'."""
    parts = spec.strip().split()
    if len(parts) != 5:
        raise ValueError(f"crontab spec must have 5 fields, got: {spec!r}")
    minute, hour, dom, month, dow = parts
    return crontab(
        minute=minute,
        hour=hour,
        day_of_month=dom,
        month_of_year=month,
        day_of_week=dow,
    )


def _default_regions() -> list[int]:
    return [int(x.strip()) for x in settings.scrape_kn_default_regions.split(",") if x.strip()]


def _nspd_default_regions() -> list[int]:
    return [int(x.strip()) for x in settings.scrape_nspd_default_regions.split(",") if x.strip()]


celery_app = Celery(
    "gendesign",
    broker=settings.redis_url,
    backend=settings.redis_url,
    include=[
        "app.workers.tasks.scrape_kn",
        "app.workers.tasks.scrape_nspd",
        "app.workers.tasks.refresh_analytics",
    ],
)
celery_app.conf.timezone = "Europe/Moscow"

# Расписание задаётся через SCRAPE_KN_CRON env var (см. app.core.config).
# Каждый региональный sweep оборачивается случайной задержкой 0..SCRAPE_KN_JITTER_SECONDS
# внутри самой задачи (см. tasks.scrape_kn) — чтобы не бить ровно в одну минуту.
celery_app.conf.beat_schedule = {
    f"kn-region-{rc}": {
        "task": "tasks.scrape_kn.scrape_kn_region",
        "schedule": _parse_cron(settings.scrape_kn_cron),
        "args": [rc, None],
    }
    for rc in _default_regions()
}

# NSPD-скрейп — квартально, после публикации rosreestr_deals.
# Pending = (новые quarter_cad_number из rosreestr_deals) − (cad_quarters_geom).
# При limit=None прогон захватит все накопленные cad-кварталы за квартал.
celery_app.conf.beat_schedule.update(
    {
        f"nspd-region-{rc}": {
            "task": "tasks.scrape_nspd.scrape_nspd_region",
            "schedule": _parse_cron(settings.scrape_nspd_cron),
            "kwargs": {
                "region_code": rc,
                "triggered_by": "beat",
                "rate_ms": settings.scrape_nspd_rate_ms,
            },
        }
        for rc in _nspd_default_regions()
    }
)

# Refresh ekb_districts медианы — ежемесячно 5-го числа в 04:00 МСК
# (после публикации новых rosreestr-кварталов и NSPD beat-cycle).
# Лёгкая задача (1-2с на 8 районов), без locks.
celery_app.conf.beat_schedule["refresh-ekb-districts-medians"] = {
    "task": "tasks.refresh_analytics.refresh_ekb_districts_medians",
    "schedule": _parse_cron("0 4 5 * *"),
    "kwargs": {"window_months": 24, "min_deals": 50},
}


@worker_ready.connect
def _resume_zombie_runs(sender=None, **_kwargs) -> None:
    """When a worker finishes booting (after redeploy/restart), find any sweep
    that was 'running' with a stale heartbeat (>5 min) and re-enqueue a resume
    task. Each resume creates a fresh run_id linked via resumed_from_run_id;
    the original row is marked 'zombie' so the audit trail is preserved.
    """
    from sqlalchemy import text

    from app.core.db import SessionLocal

    db = SessionLocal()
    try:
        rows = (
            db.execute(
                text(
                    """
                    SELECT run_id
                      FROM kn_scrape_runs
                     WHERE status = 'running'
                       AND objects_snapshot IS NOT NULL
                       AND COALESCE(heartbeat_at, started_at)
                           < NOW() - INTERVAL '5 minutes'
                     ORDER BY started_at ASC
                     LIMIT 20
                    """
                )
            )
            .mappings()
            .all()
        )
        if not rows:
            logger.info("worker_ready: нет stale runs для resume")
            return

        ids = [int(r["run_id"]) for r in rows]
        # Помечаем найденные как 'zombie' одним апдейтом — resume создаст новые
        # run_id со ссылкой resumed_from_run_id.
        db.execute(
            text(
                """
                UPDATE kn_scrape_runs
                   SET status = 'zombie',
                       finished_at = NOW(),
                       error = COALESCE(error,
                           'auto-zombie at worker_ready, resume scheduled')
                 WHERE run_id = ANY(:ids)
                """
            ),
            {"ids": ids},
        )
        db.commit()
    except Exception as e:
        logger.exception("worker_ready resume scan failed: %s", e)
        try:
            db.rollback()
        except Exception:
            pass
        return
    finally:
        db.close()

    # Enqueue resume tasks. Late import to avoid circular at module load.
    from app.workers.tasks.scrape_kn import resume_kn_run

    for rid in ids:
        try:
            resume_kn_run.apply_async(args=[rid])
            logger.info("worker_ready: resume_kn_run enqueued for run=%s", rid)
        except Exception as e:
            logger.warning("worker_ready: failed to enqueue resume for run=%s: %s", rid, e)

    # NSPD-runs: помечаем зомби (>15 мин без heartbeat — NSPD сам по себе
    # медленнее kn, дельта больше). Resume = просто заново ставим
    # scrape_nspd_region с тем же region_code и triggered_by='resume':
    # сервис идемпотентен — pending = (rosreestr ДДУ cads) − (cad_quarters_geom done),
    # уже скрейпнутые квартала пропустятся автоматом.
    db = SessionLocal()
    nspd_resume_targets: list[int] = []
    try:
        zombie_rows = (
            db.execute(
                text(
                    """
                    UPDATE nspd_scrape_runs
                       SET status = 'zombie',
                           finished_at = NOW(),
                           error = COALESCE(error,
                               'auto-zombie at worker_ready, resume scheduled')
                     WHERE status = 'running'
                       AND COALESCE(heartbeat_at, started_at)
                           < NOW() - INTERVAL '15 minutes'
                    RETURNING run_id, region_code
                    """
                )
            )
            .mappings()
            .all()
        )
        db.commit()
        # Уникальные регионы для resume (разные ран-id для одного региона —
        # повторно не enqueueum, lock и idempotency пропустят дубль).
        seen: set[int] = set()
        for r in zombie_rows:
            rc = int(r["region_code"])
            if rc not in seen:
                seen.add(rc)
                nspd_resume_targets.append(rc)
                logger.info(
                    "worker_ready: NSPD zombie run=%s region=%s — resume scheduled",
                    r["run_id"],
                    rc,
                )
    except Exception as e:
        logger.warning("worker_ready nspd zombie scan failed: %s", e)
        try:
            db.rollback()
        except Exception:
            pass
    finally:
        db.close()

    if nspd_resume_targets:
        from app.workers.tasks.scrape_nspd import scrape_nspd_region

        for rc in nspd_resume_targets:
            try:
                scrape_nspd_region.apply_async(
                    kwargs={"region_code": rc, "triggered_by": "resume"},
                )
                logger.info("worker_ready: scrape_nspd_region enqueued region=%s", rc)
            except Exception as e:
                logger.warning(
                    "worker_ready: failed to enqueue NSPD resume for region=%s: %s",
                    rc,
                    e,
                )