"""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.""" row = 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' RETURNING id """ ), {"run_id": run_id, "counters": json.dumps(counters)}, ).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 row = 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' RETURNING id """ ), {"run_id": run_id, "error": error[:1000], "counters": json.dumps(counters)}, ).first() if row is None: logger.warning("mark_failed no-op: run_id=%d not in 'running' state", run_id) db.commit() 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 row = db.execute( text( """ UPDATE scrape_runs SET status = 'banned', finished_at = NOW(), heartbeat_at = NOW(), error = :error, counters = CAST(:counters AS jsonb) WHERE id = :run_id AND status = 'running' RETURNING id """ ), {"run_id": run_id, "error": error[:1000], "counters": json.dumps(counters)}, ).first() if row is None: logger.warning("mark_banned no-op: run_id=%d not in 'running' state", run_id) 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]