gendesign/backend/app/workers/lifecycle.py

"""Worker lifecycle hooks — SQLAlchemy engine dispose on fork, zombie-resume on ready.

Handlers регистрируются через Celery signals декораторами при импорте модуля.
Импортируй этот модуль как side-effect из celery_app.py:

    from app.workers import lifecycle  # noqa: F401

worker_process_init — dispose SQLAlchemy engine в каждом prefork child-процессе,
    чтобы избежать shared TCP-сокетов к PostgreSQL после fork().

worker_ready — при рестарте воркера находит все 'running'/'paused' записи
    kn_scrape_runs и nspd_geo_jobs и re-enqueue'ит их как zombie-resume.
"""

import logging

from celery.signals import worker_process_init, worker_ready

logger = logging.getLogger(__name__)


@worker_process_init.connect
def _reset_db_connections(sender=None, **_kwargs) -> None:
    """Dispose SQLAlchemy engine in each prefork child process.

    Background: Celery prefork pool fork()'s child processes from the master
    that already imported SQLAlchemy + psycopg. The child inherits open TCP
    sockets to PostgreSQL, but psycopg connections cannot be safely shared
    across processes after fork — first use raises:
        psycopg.ProgrammingError: can't change 'autocommit' now:
        connection in transaction status INTRANS

    Fix: dispose the engine in each child. SQLAlchemy will lazily open fresh
    connections per-child. Standard Celery + SQLA recipe.
    """
    try:
        from app.core.db import engine

        engine.dispose()
        logger.info("worker_process_init: SQLAlchemy engine disposed for fork")
    except Exception as e:
        logger.warning("worker_process_init: engine dispose failed: %s", e)


@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' or 'paused' 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.

    No time threshold: by definition, on worker_ready ANY 'running' row is a
    zombie because there is no active worker. Previously we required heartbeat
    >5 min, which skipped jobs interrupted seconds before redeploy — they
    stayed in 'running' status forever and required manual cancel/resume.
    """
    logger.info("worker_ready: resume scan starting")
    from sqlalchemy import text

    # Persistent breadcrumb: write to nspd_geo_log so we can confirm via DB
    # query whether worker_ready actually fired (independent of container logs).
    # NULL job_id to avoid FK violation (FK is to nspd_geo_jobs.job_id which
    # has no 0 row).
    try:
        from app.core.db import SessionLocal

        _db = SessionLocal()
        try:
            _db.execute(
                text(
                    "INSERT INTO nspd_geo_log (job_id, level, stage, message) "
                    "VALUES (NULL, 'info', 'worker_ready', 'worker_ready signal fired')"
                )
            )
            _db.commit()
        except Exception:
            _db.rollback()
            raise
        finally:
            _db.close()
    except Exception as _e:
        logger.warning("worker_ready: breadcrumb insert failed: %s", _e)

    from app.core.db import SessionLocal

    # ── kn_scrape_runs resume ──
    db = SessionLocal()
    ids: list[int] = []
    try:
        rows = (
            db.execute(
                text(
                    """
                    SELECT run_id
                      FROM kn_scrape_runs
                     WHERE status = 'running'
                       AND objects_snapshot IS NOT NULL
                     ORDER BY started_at ASC
                     LIMIT 20
                    """
                )
            )
            .mappings()
            .all()
        )
        if rows:
            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()
        else:
            logger.info("worker_ready: нет stale kn runs для resume")
    except Exception as e:
        logger.exception("worker_ready kn resume scan failed: %s", e)
        try:
            db.rollback()
        except Exception:
            pass
    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 (старый Playwright-scraper) resume УДАЛЁН 2026-05-11. Скрапер
    # снят с эксплуатации в пользу bulk geo-fetcher (nspd_geo). История
    # старых runs сохранена в nspd_scrape_runs — никаких side-effects.

    # NSPD geo-jobs: bulk-fetcher с собственной resume-логикой через
    # nspd_geo_jobs / nspd_geo_targets. Resume любых 'running' / 'paused' jobs
    # — на worker_ready по определению нет активных воркеров, всё running ==
    # zombie. Раньше требовали heartbeat >10мин, что пропускало jobs убитых
    # за минуту до редеплоя и оставляло их вечно висеть.
    db = SessionLocal()
    geo_resume_jobs: list[int] = []
    try:
        rows = (
            db.execute(
                text(
                    """
                    UPDATE nspd_geo_jobs
                       SET status = 'queued',
                           error = COALESCE(error, 'auto-resume at worker_ready')
                     WHERE status IN ('running', 'paused')
                    RETURNING job_id
                    """
                )
            )
            .mappings()
            .all()
        )
        db.commit()
        geo_resume_jobs = [int(r["job_id"]) for r in rows]
        for jid in geo_resume_jobs:
            logger.info("worker_ready: NSPD geo job=%s — resume scheduled", jid)
    except Exception as e:
        logger.warning("worker_ready nspd_geo resume scan failed: %s", e)
        try:
            db.rollback()
        except Exception:
            pass
    finally:
        db.close()

    if geo_resume_jobs:
        from app.workers.tasks.nspd_geo import process_nspd_geo_job

        for jid in geo_resume_jobs:
            try:
                process_nspd_geo_job.apply_async(args=[jid], queue="geo")
                logger.info("worker_ready: process_nspd_geo_job enqueued job=%s", jid)
            except Exception as e:
                logger.warning("worker_ready: failed to enqueue geo resume job=%s: %s", jid, e)
    logger.info("worker_ready: resume scan finished (geo_jobs=%d)", len(geo_resume_jobs))

    # Sanity check: nspd_quarter_dumps table must exist (migration 88).
    # Logs critical error but does NOT crash the worker — table may be absent
    # in dev/staging before migration is applied.
    try:
        from app.core.db import SessionLocal as _SessionLocal

        _db = _SessionLocal()
        try:
            _db.execute(text("SELECT 1 FROM nspd_quarter_dumps LIMIT 0"))
            logger.info("worker_ready: nspd_quarter_dumps table OK")
        except Exception as _te:
            logger.critical(
                "worker_ready: nspd_quarter_dumps table missing or inaccessible — "
                "apply migration 88_nspd_quarter_dumps.sql before using harvest_quarter. "
                "Error: %s",
                _te,
            )
        finally:
            _db.close()
    except Exception as _e:
        logger.warning("worker_ready: nspd_quarter_dumps sanity check failed: %s", _e)