gendesign/tradein-mvp/scripts/trigger-schedules.py

#!/usr/bin/env python3
"""Trigger due scrape schedules via admin API. Called by systemd timer every 60s."""

from __future__ import annotations

import fcntl
import logging
import os
import random
import sys
import time

import httpx
import psycopg

LOCK_FILE = "/var/run/tradein-trigger.lock"
DB_URL = os.environ["TRADEIN_DATABASE_URL"]
API_BASE = os.environ.get("TRADEIN_API_BASE", "http://localhost:8000")
# Username that maps to role=admin in roles.yaml — passed as X-Authenticated-User.
ADMIN_TOKEN = os.environ["TRADEIN_ADMIN_TOKEN"]

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger(__name__)

# Mapping from scrape_schedules.source → admin API path.
# Sources handled entirely inside the scheduler loop (no separate admin endpoint)
# are listed in _NO_ENDPOINT_SOURCES: they are skipped when the systemd trigger
# runs, and require SCHEDULER_ENABLE=true to fire.
_SOURCE_TO_PATH: dict[str, str] = {
    "avito_city_sweep": "/api/v1/admin/scrape/avito-city-sweep",
    "yandex_city_sweep": "/api/v1/admin/scrape/yandex-city-sweep",
    "n1_city_sweep": "/api/v1/admin/scrape/n1",
    "cian_history_backfill": "/api/v1/admin/scrape/cian-backfill-history",
    "yandex_address_backfill": "/api/v1/admin/scrape/yandex-address-backfill",
}

# Sources with no dedicated admin trigger endpoint.
_NO_ENDPOINT_SOURCES = frozenset(
    {
        "rosreestr_dkp_import",
        "listing_source_snapshot",
        "asking_to_sold_ratio_refresh",
        "refresh_search_matview",
        "deactivate_stale_avito",
        "sber_index_pull",
        "rosreestr_quarter_poll",
    }
)

_MAX_RETRIES = 3
_RETRY_BASE_S = 2.0


def _normalize_db_url(url: str) -> str:
    """Strip SQLAlchemy driver prefix so psycopg v3 accepts the connection string."""
    if url.startswith("postgresql+psycopg://"):
        return "postgresql://" + url[len("postgresql+psycopg://"):]
    return url


def get_due_schedules() -> list[dict[str, str]]:
    """Fetch scrape_schedules rows that are due for execution."""
    conn_url = _normalize_db_url(DB_URL)
    with psycopg.connect(conn_url) as conn:
        rows = conn.execute(
            """
            SELECT id, source, schedule_type
            FROM scrape_schedules
            WHERE enabled = true
              AND (next_run_at IS NULL OR next_run_at <= NOW())
            """,
        ).fetchall()
        description = conn.execute(
            "SELECT id, source, schedule_type FROM scrape_schedules LIMIT 0"
        ).description or []
    col_names = [desc.name for desc in description]
    if not col_names:
        col_names = ["id", "source", "schedule_type"]
    return [
        {col_names[i]: (str(row[i]) if row[i] is not None else "") for i in range(len(col_names))}
        for row in rows
    ]


def trigger_source(client: httpx.Client, source: str) -> None:
    """POST to the admin trigger endpoint for a given source. Retries on connection error."""
    path = _SOURCE_TO_PATH.get(source)
    if path is None:
        if source in _NO_ENDPOINT_SOURCES:
            log.warning(
                "trigger: source=%s has no admin API endpoint — "
                "ensure SCHEDULER_ENABLE=true or add a dedicated admin endpoint",
                source,
            )
        else:
            log.warning("trigger: unknown source=%s, skipping", source)
        return

    url = f"{API_BASE}{path}"
    last_exc: Exception | None = None

    for attempt in range(1, _MAX_RETRIES + 1):
        try:
            resp = client.post(
                url,
                headers={"X-Authenticated-User": ADMIN_TOKEN},
                timeout=50.0,
            )
            log.info(
                "trigger: source=%s url=%s status=%d attempt=%d",
                source,
                url,
                resp.status_code,
                attempt,
            )
            if resp.status_code >= 500:
                log.warning(
                    "trigger: source=%s got HTTP %d — body: %s",
                    source,
                    resp.status_code,
                    resp.text[:200],
                )
            return
        except httpx.ConnectError as exc:
            last_exc = exc
            wait = _RETRY_BASE_S * (2 ** (attempt - 1))
            log.warning(
                "trigger: source=%s connection error (attempt %d/%d), "
                "retrying in %.1fs: %s",
                source,
                attempt,
                _MAX_RETRIES,
                wait,
                exc,
            )
            if attempt < _MAX_RETRIES:
                time.sleep(wait)

    log.error(
        "trigger: source=%s failed after %d attempts: %s",
        source,
        _MAX_RETRIES,
        last_exc,
    )


def main() -> int:
    # Exclusive non-blocking flock: if a previous run is still active, exit immediately
    # with code 0 so systemd does not log a failure.
    lock_fd = open(LOCK_FILE, "w")
    try:
        fcntl.flock(lock_fd.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
    except BlockingIOError:
        log.info("trigger: lock busy — previous run still active, exiting")
        lock_fd.close()
        return 0

    try:
        due = get_due_schedules()
        if not due:
            log.info("trigger: no schedules due")
            return 0

        log.info("trigger: %d schedule(s) due", len(due))

        with httpx.Client() as client:
            for i, row in enumerate(due):
                source = row.get("source", "")
                trigger_source(client, source)
                if i < len(due) - 1:
                    # Stagger requests to avoid thundering herd on the API
                    time.sleep(random.uniform(1, 4))

        return 0

    except Exception:
        log.exception("trigger: unexpected error")
        return 1
    finally:
        try:
            fcntl.flock(lock_fd.fileno(), fcntl.LOCK_UN)
        except Exception:
            pass
        lock_fd.close()


if __name__ == "__main__":
    sys.exit(main())