gendesign/tradein-mvp/backend/app/services/scrape_runs.py

"""Utility functions для scrape_runs table — tracking long-running pipeline runs.

Таблица scrape_runs создана в 015_scrape_runs.sql.
Расширена в 051_scrape_runs_extend.sql: params/counters/error/finished_at/cancelled.
"""

from __future__ import annotations

import json
import logging
from typing import Any

import sentry_sdk
from sqlalchemy import text
from sqlalchemy.orm import Session

logger = logging.getLogger(__name__)

# Количество последовательных неудачных запусков (failed/banned), при достижении
# которого отправляется Sentry alert. Алерт срабатывает ровно при N-й подряд ошибке
# (anti-spam: не на каждой последующей).
CONSECUTIVE_FAILURE_ALERT_THRESHOLD = 3


def _column_counts(counters: dict[str, int]) -> tuple[int | None, int | None]:
    """Извлечь значения для dedicated-колонок total_seen / new_count из jsonb-counters.

    Все sweep-counters (YandexCitySweepCounters / CitySweepCounters / …) пишут число
    объявлений в выдаче как ``lots_fetched`` и число впервые вставленных как
    ``lots_inserted``. Раньше эти значения попадали ТОЛЬКО в counters jsonb, а
    выделенные колонки total_seen/new_count оставались 0 → admin/observability
    показывала total_seen=0 при реально сохранённых строках (audit #1871/#1926).

    Приоритет ключей:
      - total_seen ← 'total_seen' (если уже есть в counters) иначе 'lots_fetched'
      - new_count  ← 'new_count'  (если уже есть)            иначе 'lots_inserted'

    Возвращает (total_seen, new_count); None для ключа, которого нет в counters —
    тогда соответствующая колонка не перезаписывается (COALESCE-семантика в UPDATE).
    """

    def _pick(*keys: str) -> int | None:
        for key in keys:
            val = counters.get(key)
            if val is not None:
                try:
                    return int(val)
                except (TypeError, ValueError):
                    return None
        return None

    return _pick("total_seen", "lots_fetched"), _pick("new_count", "lots_inserted")


def _alert_if_consecutive_failures(db: Session, source: str) -> None:
    """Отправить Sentry alert если последние CONSECUTIVE_FAILURE_ALERT_THRESHOLD
    завершённых запусков для данного source имеют статус 'failed' или 'banned'.

    Anti-spam: алерт срабатывает ТОЛЬКО когда стрик РОВНО равен порогу — т.е. запрос
    возвращает ровно N последних (failed|banned) и (N+1)-й, если существует, НЕ является
    failed/banned. Это предотвращает повторный алерт на каждой ошибке сверх порога.

    Best-effort: весь блок обёрнут в try/except — сбой запроса или неинициализированный
    Sentry НЕ должен нарушать вызывающий mark_* путь.
    """
    n = CONSECUTIVE_FAILURE_ALERT_THRESHOLD
    try:
        # Берём последние N+1 завершённых (non-running) запусков по source.
        # Сортируем по finished_at DESC чтобы самые свежие шли первыми.
        rows = db.execute(
            text(
                """
                SELECT status FROM scrape_runs
                WHERE source = :source
                  AND status IN ('failed', 'banned', 'done', 'cancelled')
                ORDER BY finished_at DESC NULLS LAST
                LIMIT :limit
                """
            ),
            {"source": source, "limit": n + 1},
        ).fetchall()

        if len(rows) < n:
            # Ещё не набралось N завершённых запусков вообще — алерт не нужен.
            return

        # Первые N должны быть все failed/banned.
        first_n = rows[:n]
        if not all(r.status in ("failed", "banned") for r in first_n):
            return

        # (N+1)-й запуск, если есть, тоже должен НЕ быть failed/banned — иначе мы уже
        # должны были отправить алерт раньше и не стоит дублировать.
        if len(rows) > n and rows[n].status in ("failed", "banned"):
            return

        # Стрик ровно достиг порога — отправляем алерт.
        sentry_sdk.capture_message(
            f"Scraper source '{source}' has {n} consecutive failed/banned runs — "
            "manual intervention may be required (expired cookies / ban / broken parser).",
            level="error",
        )
        logger.error(
            "sentry alert sent: source=%s has %d consecutive failed/banned runs", source, n
        )
    except Exception:
        pass  # sentry_sdk not initialised in dev, or query failed — best-effort only


def _alert_on_run_id(db: Session, run_id: int) -> None:
    """Вспомогательная обёртка: извлекает source по run_id и вызывает
    _alert_if_consecutive_failures. Best-effort — не бросает исключений.
    """
    try:
        row = db.execute(
            text("SELECT source FROM scrape_runs WHERE id = :run_id"),
            {"run_id": run_id},
        ).fetchone()
        if row is None:
            return
        _alert_if_consecutive_failures(db, str(row.source))
    except Exception:
        pass


def create_run(db: Session, *, source: str, params: dict[str, Any]) -> int:
    """INSERT scrape_runs(source, status='running', params, started_at=NOW()).

    run_type DEFAULT 'city_sweep' (из 051 миграции).
    Returns run_id (bigint).
    """
    row = db.execute(
        text(
            """
            INSERT INTO scrape_runs (source, status, params, started_at, heartbeat_at)
            VALUES (:source, 'running', CAST(:params AS jsonb), NOW(), NOW())
            RETURNING id
            """
        ),
        {"source": source, "params": json.dumps(params, ensure_ascii=False)},
    ).fetchone()
    db.commit()
    assert row is not None, "scrape_runs INSERT returned no id"
    return int(row.id)


def update_heartbeat(db: Session, run_id: int, counters: dict[str, int]) -> None:
    """UPDATE heartbeat_at=NOW(), counters=:counters + total_seen/new_count колонки.

    total_seen/new_count извлекаются из counters (lots_fetched/lots_inserted) и
    пишутся в выделенные колонки, чтобы observability не показывала 0 (audit #1926).
    COALESCE: если ключа нет в counters — старое значение колонки сохраняется.
    """
    total_seen, new_count = _column_counts(counters)
    db.execute(
        text(
            """
            UPDATE scrape_runs
            SET heartbeat_at = NOW(),
                counters = CAST(:counters AS jsonb),
                total_seen = COALESCE(CAST(:total_seen AS int), total_seen),
                new_count = COALESCE(CAST(:new_count AS int), new_count)
            WHERE id = :run_id
            """
        ),
        {
            "run_id": run_id,
            "counters": json.dumps(counters),
            "total_seen": total_seen,
            "new_count": new_count,
        },
    )
    db.commit()


def is_cancelled(db: Session, run_id: int) -> bool:
    """Проверить status='cancelled' (cooperative cancel в long-running pipeline)."""
    row = db.execute(
        text("SELECT status FROM scrape_runs WHERE id = :id"),
        {"id": run_id},
    ).fetchone()
    return row is not None and row.status == "cancelled"


def mark_done(db: Session, run_id: int, counters: dict[str, int]) -> None:
    """Финализация run: status='done', finished_at=NOW(), counters + total_seen/new_count.

    total_seen/new_count извлекаются из counters (lots_fetched/lots_inserted) и пишутся
    в выделенные колонки — иначе admin/observability показывает 0 (audit #1926).
    """
    total_seen, new_count = _column_counts(counters)
    row = db.execute(
        text(
            """
            UPDATE scrape_runs
            SET status = 'done', finished_at = NOW(), heartbeat_at = NOW(),
                counters = CAST(:counters AS jsonb),
                total_seen = COALESCE(CAST(:total_seen AS int), total_seen),
                new_count = COALESCE(CAST(:new_count AS int), new_count)
            WHERE id = :run_id AND status = 'running'
            RETURNING id
            """
        ),
        {
            "run_id": run_id,
            "counters": json.dumps(counters),
            "total_seen": total_seen,
            "new_count": new_count,
        },
    ).first()
    if row is None:
        logger.warning("mark_done no-op: run_id=%d not in 'running' state", run_id)
    db.commit()


def mark_failed(db: Session, run_id: int, error: str, counters: dict[str, int]) -> None:
    """Финализация run: status='failed', error (первые 1000 символов).

    Defensive rollback: если до этого вызова в той же транзакции был ошибочный UPDATE,
    он мог оставить сессию в error state — rollback сбрасывает состояние.
    """
    try:
        db.rollback()  # defensive: cascade-safe if prior UPDATE left txn in error state
    except Exception:
        pass
    total_seen, new_count = _column_counts(counters)
    row = db.execute(
        text(
            """
            UPDATE scrape_runs
            SET status = 'failed', finished_at = NOW(), heartbeat_at = NOW(),
                error = :error, counters = CAST(:counters AS jsonb),
                total_seen = COALESCE(CAST(:total_seen AS int), total_seen),
                new_count = COALESCE(CAST(:new_count AS int), new_count)
            WHERE id = :run_id AND status = 'running'
            RETURNING id
            """
        ),
        {
            "run_id": run_id,
            "error": error[:1000],
            "counters": json.dumps(counters),
            "total_seen": total_seen,
            "new_count": new_count,
        },
    ).first()
    if row is None:
        logger.warning("mark_failed no-op: run_id=%d not in 'running' state", run_id)
    db.commit()
    # Алерт после коммита — статус уже персистирован в БД; best-effort (не бросает).
    _alert_on_run_id(db, run_id)


def mark_banned(db: Session, run_id: int, error: str, counters: dict[str, int]) -> None:
    """Финализация run: status='banned' (IP заблокирован Avito — 403/captcha).

    Per migration 015 — 'banned' задокументирован как 'Avito вернул 403/captcha'.
    Отличается от 'failed': это external constraint, не наш bug. Cooldown 2-4 часа.

    Defensive rollback: если до этого вызова в той же транзакции был ошибочный UPDATE,
    он мог оставить сессию в error state — rollback сбрасывает состояние.
    """
    try:
        db.rollback()  # defensive: cascade-safe if prior UPDATE left txn in error state
    except Exception:
        pass
    total_seen, new_count = _column_counts(counters)
    row = db.execute(
        text(
            """
            UPDATE scrape_runs
            SET status = 'banned', finished_at = NOW(), heartbeat_at = NOW(),
                error = :error, counters = CAST(:counters AS jsonb),
                total_seen = COALESCE(CAST(:total_seen AS int), total_seen),
                new_count = COALESCE(CAST(:new_count AS int), new_count)
            WHERE id = :run_id AND status = 'running'
            RETURNING id
            """
        ),
        {
            "run_id": run_id,
            "error": error[:1000],
            "counters": json.dumps(counters),
            "total_seen": total_seen,
            "new_count": new_count,
        },
    ).first()
    if row is None:
        logger.warning("mark_banned no-op: run_id=%d not in 'running' state", run_id)
    db.commit()
    # Алерт после коммита — статус уже персистирован в БД; best-effort (не бросает).
    _alert_on_run_id(db, run_id)


def mark_cancelled(db: Session, run_id: int) -> bool:
    """Set status='cancelled' если currently 'running'. Returns True если cancelled."""
    result = db.execute(
        text(
            """
            UPDATE scrape_runs
            SET status = 'cancelled', finished_at = NOW()
            WHERE id = :run_id AND status = 'running'
            RETURNING id
            """
        ),
        {"run_id": run_id},
    ).fetchone()
    db.commit()
    return result is not None


def list_recent(db: Session, *, source: str, limit: int = 10) -> list[dict[str, Any]]:
    """Список последних N runs по source, в порядке убывания started_at."""
    rows = (
        db.execute(
            text(
                """
            SELECT id AS run_id, source, status, params, counters, error,
                   started_at, finished_at, heartbeat_at
            FROM scrape_runs
            WHERE source = :source
            ORDER BY started_at DESC NULLS LAST
            LIMIT :limit
            """
            ),
            {"source": source, "limit": limit},
        )
        .mappings()
        .all()
    )
    return [dict(r) for r in rows]


def list_all(
    db: Session,
    *,
    source: str | None = None,
    status: str | None = None,
    limit: int = 50,
    offset: int = 0,
) -> tuple[int, list[dict[str, Any]]]:
    """Unified-список runs по всем source'ам с опц. фильтрами source/status.

    Возвращает (total, rows): total — кол-во строк под фильтром (без limit/offset),
    rows — текущая страница (ORDER BY started_at DESC, LIMIT/OFFSET).

    Используется единой scrapers-страницей вместо 3 per-source /runs.
    """
    where = ["1=1"]
    params: dict[str, Any] = {"limit": limit, "offset": offset}
    if source is not None:
        where.append("source = :source")
        params["source"] = source
    if status is not None:
        where.append("status = :status")
        params["status"] = status
    where_sql = " AND ".join(where)

    total_row = db.execute(
        text(f"SELECT COUNT(*) AS n FROM scrape_runs WHERE {where_sql}"),
        params,
    ).fetchone()
    total = int(total_row.n) if total_row is not None else 0

    rows = (
        db.execute(
            text(
                f"""
            SELECT id AS run_id, source, run_type, status, params, counters,
                   total_seen, new_count, started_at, finished_at,
                   heartbeat_at, error AS error_text
            FROM scrape_runs
            WHERE {where_sql}
            ORDER BY started_at DESC NULLS LAST
            LIMIT CAST(:limit AS int) OFFSET CAST(:offset AS int)
            """
            ),
            params,
        )
        .mappings()
        .all()
    )
    return total, [dict(r) for r in rows]