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

from sqlalchemy import text
from sqlalchemy.orm import Session

logger = logging.getLogger(__name__)


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."""
    db.execute(
        text(
            """
            UPDATE scrape_runs
            SET heartbeat_at = NOW(), counters = CAST(:counters AS jsonb)
            WHERE id = :run_id
            """
        ),
        {"run_id": run_id, "counters": json.dumps(counters)},
    )
    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."""
    db.execute(
        text(
            """
            UPDATE scrape_runs
            SET status = 'done', finished_at = NOW(), heartbeat_at = NOW(),
                counters = CAST(:counters AS jsonb)
            WHERE id = :run_id AND status = 'running'
            """
        ),
        {"run_id": run_id, "counters": json.dumps(counters)},
    )
    db.commit()


def mark_failed(db: Session, run_id: int, error: str, counters: dict[str, int]) -> None:
    """Финализация run: status='failed', error (первые 1000 символов)."""
    db.execute(
        text(
            """
            UPDATE scrape_runs
            SET status = 'failed', finished_at = NOW(), heartbeat_at = NOW(),
                error = :error, counters = CAST(:counters AS jsonb)
            WHERE id = :run_id AND status = 'running'
            """
        ),
        {"run_id": run_id, "error": error[:1000], "counters": json.dumps(counters)},
    )
    db.commit()


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

    Отличается от 'failed': это external constraint, не наш bug.
    Cooldown 2-4 часа потом retry.
    """
    db.execute(
        text(
            """
            UPDATE scrape_runs
            SET status = 'rate_limited', finished_at = NOW(), heartbeat_at = NOW(),
                error = :error, counters = CAST(:counters AS jsonb)
            WHERE id = :run_id AND status = 'running'
            """
        ),
        {"run_id": run_id, "error": error[:1000], "counters": json.dumps(counters)},
    )
    db.commit()


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]