363 lines
12 KiB
Python
363 lines
12 KiB
Python
"""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
|
||
|
||
|
||
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 scrape_kn_region
|
||
|
||
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,
|
||
}
|
||
|
||
|
||
@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: workers killed by SIGKILL (OOM, container restart,
|
||
# deploy mid-run) leave kn_scrape_runs rows stuck in 'running' forever
|
||
# because the except/finally block never fires. Anything older than
|
||
# 60 minutes without finishing is almost certainly dead.
|
||
db.execute(
|
||
text(
|
||
"""
|
||
UPDATE kn_scrape_runs
|
||
SET status = 'zombie',
|
||
finished_at = NOW(),
|
||
error = 'auto-marked: no heartbeat 60+ min'
|
||
WHERE status = 'running'
|
||
AND finished_at IS NULL
|
||
AND started_at < NOW() - INTERVAL '60 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 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
|
||
]
|