116 lines
4.0 KiB
Python
116 lines
4.0 KiB
Python
from fastapi import APIRouter, HTTPException, status, WebSocket, WebSocketDisconnect
|
||
from pydantic import BaseModel
|
||
import asyncio
|
||
import json as _json
|
||
|
||
router = APIRouter(prefix="/api/admin", tags=["admin"])
|
||
|
||
SOURCES = {
|
||
"wikipedia_ru": "Wikipedia RU (~8 GB)",
|
||
"wikipedia_en": "Wikipedia EN (~22 GB)",
|
||
}
|
||
|
||
|
||
class JobResponse(BaseModel):
|
||
status: str
|
||
source: str | None = None
|
||
label: str | None = None
|
||
progress: int = 0
|
||
message: str = ""
|
||
started_at: str | None = None
|
||
finished_at: str | None = None
|
||
docs_done: int = 0
|
||
|
||
|
||
@router.get("/index/sources")
|
||
def list_sources():
|
||
"""Список доступных источников для индексации."""
|
||
return [{"key": k, "label": v} for k, v in SOURCES.items()]
|
||
|
||
|
||
@router.post("/index/download/{source_key}", response_model=JobResponse)
|
||
def start_download(source_key: str):
|
||
"""Скачать дамп и проиндексировать с нуля."""
|
||
if source_key not in SOURCES:
|
||
raise HTTPException(status_code=400, detail=f"Unknown source: {source_key}")
|
||
|
||
from app.tasks.indexing import get_job_status, download_and_index
|
||
current = get_job_status()
|
||
if current.get("status") == "running":
|
||
raise HTTPException(status_code=409, detail="Задание уже выполняется")
|
||
|
||
download_and_index.delay(source_key)
|
||
return JobResponse(status="started", source=source_key, label=SOURCES[source_key], message="Запуск скачивания…")
|
||
|
||
|
||
@router.post("/index/update/{source_key}", response_model=JobResponse)
|
||
def start_update(source_key: str):
|
||
"""Обновить индекс — только новые документы (дамп не скачивается заново если уже есть)."""
|
||
if source_key not in SOURCES:
|
||
raise HTTPException(status_code=400, detail=f"Unknown source: {source_key}")
|
||
|
||
from app.tasks.indexing import get_job_status, update_index
|
||
current = get_job_status()
|
||
if current.get("status") == "running":
|
||
raise HTTPException(status_code=409, detail="Задание уже выполняется")
|
||
|
||
update_index.delay(source_key)
|
||
return JobResponse(status="started", source=source_key, label=SOURCES[source_key], message="Запуск обновления…")
|
||
|
||
|
||
@router.get("/index/status", response_model=JobResponse)
|
||
def job_status():
|
||
"""Текущий статус задания индексации."""
|
||
from app.tasks.indexing import get_job_status
|
||
data = get_job_status()
|
||
return JobResponse(
|
||
status=data.get("status", "idle"),
|
||
source=data.get("source"),
|
||
label=data.get("label"),
|
||
progress=data.get("progress", 0),
|
||
message=data.get("message", ""),
|
||
started_at=data.get("started_at"),
|
||
finished_at=data.get("finished_at"),
|
||
docs_done=data.get("docs_done", 0),
|
||
)
|
||
|
||
|
||
@router.websocket("/ws/index-progress")
|
||
async def ws_index_progress(websocket: WebSocket):
|
||
"""Stream indexing job progress via Redis pub/sub."""
|
||
await websocket.accept()
|
||
|
||
import redis.asyncio as aioredis
|
||
from app.config import settings
|
||
|
||
r = aioredis.from_url(settings.REDIS_URL)
|
||
pubsub = r.pubsub()
|
||
await pubsub.subscribe("index_job_progress")
|
||
|
||
# Send current state immediately
|
||
from app.tasks.indexing import get_job_status
|
||
await websocket.send_json(get_job_status())
|
||
|
||
try:
|
||
async for message in pubsub.listen():
|
||
if message["type"] != "message":
|
||
continue
|
||
try:
|
||
data = _json.loads(message["data"])
|
||
except Exception:
|
||
continue
|
||
await websocket.send_json(data)
|
||
if data.get("status") in ("done", "error"):
|
||
break
|
||
except (WebSocketDisconnect, asyncio.CancelledError):
|
||
pass
|
||
finally:
|
||
try:
|
||
await pubsub.unsubscribe("index_job_progress")
|
||
except Exception:
|
||
pass
|
||
try:
|
||
await r.aclose()
|
||
except Exception:
|
||
pass
|