gendesign/backend/app/api/v1/admin_scrape.py

"""Admin trigger for ad-hoc kn-API sweeps.

POST /api/v1/admin/scrape/kn
Body: {"region_code": 66, "developers": ["6208_0"]?, "fetch_flats": true?}
Header: X-Admin-Token: <SCRAPE_ADMIN_TOKEN env var>

Disabled if SCRAPE_ADMIN_TOKEN is empty (default).
Returns the Celery task id; poll /api/v1/admin/scrape/runs/{run_id} to track DB-side progress.
"""

from __future__ import annotations

from typing import Annotated, Any

from fastapi import APIRouter, Depends, Header, HTTPException
from pydantic import BaseModel, Field
from sqlalchemy import text
from sqlalchemy.orm import Session

from app.core.config import settings
from app.core.db import get_db

router = APIRouter()


class TriggerKnRequest(BaseModel):
    region_code: int = Field(..., ge=1, le=99)
    developers: list[str] | None = Field(default=None, max_length=20)
    fetch_flats: bool = True
    extras: bool = True
    download_photos: bool = False
    # Если True — ПЕРЕД apply_async удалить existing Redis lock. Эмердженси-режим
    # для случая «зомби-worker оставил lock и блокирует новые задачи».
    force: bool = False


def _check_token(x_admin_token: str | None) -> None:
    if not settings.scrape_admin_token:
        raise HTTPException(
            status_code=503,
            detail="admin scrape disabled — set SCRAPE_ADMIN_TOKEN",
        )
    if x_admin_token != settings.scrape_admin_token:
        raise HTTPException(status_code=401, detail="invalid admin token")


@router.post("/kn")
def trigger_kn_sweep(
    payload: TriggerKnRequest,
    x_admin_token: Annotated[str | None, Header(alias="X-Admin-Token")] = None,
) -> dict[str, Any]:
    _check_token(x_admin_token)
    # Defer Celery import so the API works without a broker in dev.
    from app.workers.tasks.scrape_kn import (
        force_release_lock,
        scrape_kn_region,
    )

    lock_released = False
    if payload.force:
        lock_released = force_release_lock(payload.region_code, payload.developers)

    result = scrape_kn_region.apply_async(
        args=[payload.region_code, payload.developers],
        kwargs={
            "skip_jitter": True,
            "extras": payload.extras,
            "download_photos": payload.download_photos,
            "fetch_flats": payload.fetch_flats,
        },
    )
    return {
        "task_id": result.id,
        "region_code": payload.region_code,
        "developers": payload.developers,
        "queued_at": result.date_done.isoformat() if result.date_done else None,
        "force": payload.force,
        "lock_was_released": lock_released,
    }


@router.get("/queue")
def queue_status(
    db: Annotated[Session, Depends(get_db)],
    x_admin_token: Annotated[str | None, Header(alias="X-Admin-Token")] = None,
) -> dict[str, Any]:
    """Snapshot of Celery state, optimised for UI poll latency.

    - active: from kn_scrape_runs WHERE status='running' (instant, richer info
      than celery inspect — includes objects_count/flats_count/requests_count).
    - reserved + workers: celery inspect with short timeout (500 ms), called in
      a thread so a missing worker can't block the request.
    - queue_depth: Redis LLEN on default queue (single round-trip).
    - scheduled: omitted — we don't use ETA tasks.

    Worst-case latency ≈ 600 ms even if no worker is reachable, vs ~8 s with
    the previous 4×2 s blocking inspect calls.
    """
    _check_token(x_admin_token)

    import concurrent.futures
    import time

    from app.workers.celery_app import celery_app

    # 0) Reap zombies by heartbeat. After Resume_Checkpoint refactor the worker
    # writes heartbeat_at every ~30 sec (per 10 objects). 10 min без heartbeat
    # = точно мёртв. Используем COALESCE(heartbeat_at, started_at) чтобы
    # подхватить и legacy-runs (без heartbeat) старше 10 мин.
    db.execute(
        text(
            """
            UPDATE kn_scrape_runs
               SET status = 'zombie',
                   finished_at = NOW(),
                   error = 'auto-marked: no heartbeat 10+ min'
             WHERE status = 'running'
               AND finished_at IS NULL
               AND COALESCE(heartbeat_at, started_at)
                   < NOW() - INTERVAL '10 minutes'
            """
        )
    )
    db.commit()

    # 1) Active runs straight from DB — no broker round-trip.
    active_rows = (
        db.execute(
            text(
                """
                SELECT run_id, started_at, region_codes, developer_ids,
                       objects_count, flats_count, requests_count
                  FROM kn_scrape_runs
                 WHERE status = 'running'
                 ORDER BY started_at DESC
                 LIMIT 20
                """
            )
        )
        .mappings()
        .all()
    )
    active: list[dict[str, Any]] = [
        {
            "worker": "db",
            "task_id": f"run-{r['run_id']}",
            "name": "scrape_kn_region",
            "args": [list(r["region_codes"] or []), list(r["developer_ids"] or [])],
            "kwargs": {
                "objects": r["objects_count"],
                "flats": r["flats_count"],
                "requests": r["requests_count"],
            },
            "time_start": (int(r["started_at"].timestamp()) if r["started_at"] else None),
            "eta": None,
        }
        for r in active_rows
    ]

    # 2) Run reserved + ping in parallel with a tight timeout so the endpoint
    # never hangs when the broker is unreachable or worker is gone.
    inspect = celery_app.control.inspect(timeout=0.5)

    def _flatten(per_worker: dict[str, list[dict[str, Any]]] | None) -> list[dict[str, Any]]:
        if not per_worker:
            return []
        out: list[dict[str, Any]] = []
        for worker, tasks in per_worker.items():
            for t in tasks or []:
                out.append(
                    {
                        "worker": worker,
                        "task_id": t.get("id"),
                        "name": t.get("name"),
                        "args": t.get("args"),
                        "kwargs": t.get("kwargs"),
                        "time_start": t.get("time_start"),
                        "eta": t.get("eta"),
                    }
                )
        return out

    def _safe(fn):
        try:
            return fn()
        except Exception:
            return None

    deadline = time.monotonic() + 0.8
    with concurrent.futures.ThreadPoolExecutor(max_workers=2) as ex:
        f_reserved = ex.submit(_safe, inspect.reserved)
        f_ping = ex.submit(_safe, inspect.ping)
        try:
            reserved_raw = f_reserved.result(timeout=max(0.1, deadline - time.monotonic()))
        except concurrent.futures.TimeoutError:
            reserved_raw = None
        try:
            ping_resp = f_ping.result(timeout=max(0.1, deadline - time.monotonic()))
        except concurrent.futures.TimeoutError:
            ping_resp = None

    reserved = _flatten(reserved_raw)
    workers = list((ping_resp or {}).keys())

    # 3) Pending in broker queue (not yet picked up by any worker).
    queue_depth: int | None = None
    try:
        with celery_app.connection_or_acquire() as conn:
            with conn.channel() as channel:
                queue_depth = channel.client.llen("celery")
    except Exception:
        queue_depth = None

    return {
        "workers": workers,
        "queue_depth": queue_depth,
        "active": active,
        "reserved": reserved,
        "scheduled": [],
    }


class ReleaseLockRequest(BaseModel):
    region_code: int = Field(..., ge=1, le=99)
    developers: list[str] | None = Field(default=None, max_length=20)


@router.post("/release-lock")
def release_lock(
    payload: ReleaseLockRequest,
    x_admin_token: Annotated[str | None, Header(alias="X-Admin-Token")] = None,
) -> dict[str, Any]:
    """Принудительно удалить Redis-lock для (region, devs). Использовать когда
    зомби-worker оставил lock и новые задачи скипаются с reason='lock_held'."""
    _check_token(x_admin_token)
    from app.workers.tasks.scrape_kn import force_release_lock

    released = force_release_lock(payload.region_code, payload.developers)
    return {
        "region_code": payload.region_code,
        "developers": payload.developers,
        "released": released,
    }


class RevokeRequest(BaseModel):
    task_id: str = Field(..., min_length=1, max_length=255)
    terminate: bool = False  # SIGTERM running task vs just unmark from queue


@router.post("/revoke")
def revoke_task(
    payload: RevokeRequest,
    x_admin_token: Annotated[str | None, Header(alias="X-Admin-Token")] = None,
) -> dict[str, Any]:
    """Cancel a queued or running task by id."""
    _check_token(x_admin_token)
    from app.workers.celery_app import celery_app

    celery_app.control.revoke(payload.task_id, terminate=payload.terminate, signal="SIGTERM")
    return {"task_id": payload.task_id, "revoked": True, "terminate": payload.terminate}


@router.get("/failures")
def list_failures(
    db: Annotated[Session, Depends(get_db)],
    x_admin_token: Annotated[str | None, Header(alias="X-Admin-Token")] = None,
    run_id: int | None = None,
    limit: int = 50,
) -> list[dict[str, Any]]:
    """Per-request failure log for manual browser verification."""
    _check_token(x_admin_token)
    where = ""
    params: dict[str, Any] = {"lim": min(limit, 200)}
    if run_id is not None:
        where = "WHERE run_id = :rid"
        params["rid"] = run_id
    rows = (
        db.execute(
            text(
                f"""
            SELECT failure_id, run_id, obj_id, endpoint, full_url,
                   http_status, error, body_preview, occurred_at
              FROM kn_scrape_failures
              {where}
             ORDER BY occurred_at DESC
             LIMIT :lim
            """
            ),
            params,
        )
        .mappings()
        .all()
    )
    return [
        {
            "failure_id": r["failure_id"],
            "run_id": r["run_id"],
            "obj_id": r["obj_id"],
            "endpoint": r["endpoint"],
            "full_url": r["full_url"],
            "http_status": r["http_status"],
            "error": r["error"],
            "body_preview": r["body_preview"],
            "occurred_at": r["occurred_at"].isoformat() if r["occurred_at"] else None,
        }
        for r in rows
    ]


@router.get("/logs")
def list_logs(
    db: Annotated[Session, Depends(get_db)],
    x_admin_token: Annotated[str | None, Header(alias="X-Admin-Token")] = None,
    run_id: int | None = None,
    since_id: int | None = None,
    limit: int = 200,
) -> list[dict[str, Any]]:
    """Per-run progress events. Use since_id to poll incrementally:
    pass the highest log_id seen → returns only newer rows."""
    _check_token(x_admin_token)
    where: list[str] = []
    params: dict[str, Any] = {"lim": min(limit, 1000)}
    if run_id is not None:
        where.append("run_id = :rid")
        params["rid"] = run_id
    if since_id is not None:
        where.append("log_id > :sid")
        params["sid"] = since_id
    where_sql = ("WHERE " + " AND ".join(where)) if where else ""
    rows = (
        db.execute(
            text(
                f"""
                SELECT log_id, run_id, ts, level, stage, obj_id, message
                  FROM kn_scrape_log
                  {where_sql}
                 ORDER BY log_id DESC
                 LIMIT :lim
                """
            ),
            params,
        )
        .mappings()
        .all()
    )
    return [
        {
            "log_id": r["log_id"],
            "run_id": r["run_id"],
            "ts": r["ts"].isoformat() if r["ts"] else None,
            "level": r["level"],
            "stage": r["stage"],
            "obj_id": r["obj_id"],
            "message": r["message"],
        }
        for r in rows
    ]


@router.get("/runs")
def list_runs(
    db: Annotated[Session, Depends(get_db)],
    x_admin_token: Annotated[str | None, Header(alias="X-Admin-Token")] = None,
    limit: int = 20,
) -> list[dict[str, Any]]:
    _check_token(x_admin_token)
    rows = (
        db.execute(
            text(
                """
            SELECT run_id, started_at, finished_at, region_codes, developer_ids,
                   objects_count, flats_count, requests_count, status, error, snapshot_date
              FROM kn_scrape_runs
             ORDER BY started_at DESC
             LIMIT :lim
            """
            ),
            {"lim": min(limit, 100)},
        )
        .mappings()
        .all()
    )
    return [
        {
            "run_id": r["run_id"],
            "started_at": r["started_at"].isoformat() if r["started_at"] else None,
            "finished_at": r["finished_at"].isoformat() if r["finished_at"] else None,
            "region_codes": list(r["region_codes"]) if r["region_codes"] else [],
            "developer_ids": list(r["developer_ids"]) if r["developer_ids"] else [],
            "objects_count": r["objects_count"],
            "flats_count": r["flats_count"],
            "requests_count": r["requests_count"],
            "status": r["status"],
            "error": r["error"],
            "snapshot_date": r["snapshot_date"].isoformat() if r["snapshot_date"] else None,
        }
        for r in rows
    ]