524 lines
17 KiB
Python
524 lines
17 KiB
Python
from __future__ import annotations
|
||
|
||
from pathlib import Path
|
||
from typing import Any, Dict, List, Optional
|
||
|
||
from fastapi import Depends, FastAPI, File, Form, HTTPException, Query, UploadFile
|
||
from fastapi.responses import Response
|
||
from pydantic import BaseModel, Field
|
||
|
||
from domain.models import JobStatus
|
||
from infrastructure.auth import require_bearer
|
||
from infrastructure.storage_minio import MinioSettings, MinioStorage
|
||
from infrastructure.store import JobStore
|
||
|
||
store = JobStore()
|
||
minio = MinioStorage(MinioSettings.from_env())
|
||
|
||
app = FastAPI(title="BotDescarga Remote API", version="1.0.0")
|
||
|
||
|
||
class RegisterWorkerBody(BaseModel):
|
||
worker_id: str
|
||
hostname: str = ""
|
||
max_concurrent: int = 1
|
||
|
||
|
||
class HeartbeatBody(BaseModel):
|
||
worker_id: str
|
||
|
||
|
||
class EventBody(BaseModel):
|
||
event_type: str
|
||
message: str
|
||
radicado: Optional[str] = None
|
||
nit: Optional[str] = None
|
||
index: Optional[int] = None
|
||
total: Optional[int] = None
|
||
metadata: Dict[str, Any] = Field(default_factory=dict)
|
||
timestamp: Optional[str] = None
|
||
|
||
|
||
class PatchJobBody(BaseModel):
|
||
status: Optional[str] = None
|
||
result: Optional[Dict[str, Any]] = None
|
||
error: Optional[str] = None
|
||
|
||
|
||
@app.get("/health")
|
||
def health() -> dict:
|
||
return {"status": "ok", "service": "bot-descarga"}
|
||
|
||
|
||
@app.post("/workers/register")
|
||
def register_worker(body: RegisterWorkerBody, _: str = Depends(require_bearer)) -> dict:
|
||
worker = store.register_worker(
|
||
body.worker_id,
|
||
hostname=body.hostname,
|
||
max_concurrent=body.max_concurrent,
|
||
)
|
||
return worker.to_dict()
|
||
|
||
|
||
@app.post("/workers/heartbeat")
|
||
def heartbeat(body: HeartbeatBody, _: str = Depends(require_bearer)) -> dict:
|
||
worker = store.heartbeat(body.worker_id)
|
||
if not worker:
|
||
raise HTTPException(status_code=404, detail="Worker no registrado")
|
||
return worker.to_dict()
|
||
|
||
|
||
@app.get("/workers")
|
||
def list_workers(_: str = Depends(require_bearer)) -> List[dict]:
|
||
return [w.to_dict() for w in store.list_workers()]
|
||
|
||
|
||
@app.get("/workers/me/pending")
|
||
def claim_pending(
|
||
worker_id: str = Query(...),
|
||
_: str = Depends(require_bearer),
|
||
) -> dict:
|
||
job = store.claim_pending(worker_id)
|
||
if not job:
|
||
return {"job": None}
|
||
return {"job": job.to_dict()}
|
||
|
||
|
||
@app.post("/jobs")
|
||
async def create_job(
|
||
worker_id: str = Form(""),
|
||
worker_ids: str = Form(""),
|
||
job_id: str = Form(...),
|
||
title: str = Form(""),
|
||
file: UploadFile = File(...),
|
||
_: str = Depends(require_bearer),
|
||
) -> dict:
|
||
"""Crea una corrida de descarga.
|
||
|
||
- 1 worker → flujo clásico (un solo job claimable).
|
||
- N workers (worker_ids JSON o todos online) → padre + N shards,
|
||
mismo downloads_prefix del padre.
|
||
"""
|
||
import json as _json
|
||
|
||
from domain.excel_split import parse_radicado_rows, rows_to_xlsx_bytes, split_rows
|
||
from infrastructure.storage_minio import downloads_prefix as mk_prefix
|
||
|
||
parent_job_id = job_id.strip()
|
||
if not parent_job_id:
|
||
raise HTTPException(status_code=400, detail="job_id (Job Desktop) obligatorio")
|
||
if not store.parent_job_exists(parent_job_id):
|
||
raise HTTPException(status_code=404, detail="Trabajo (job) no encontrado")
|
||
|
||
data = await file.read()
|
||
if not data:
|
||
raise HTTPException(status_code=400, detail="Archivo vacío")
|
||
|
||
original = file.filename or "input.xlsx"
|
||
content_type = file.content_type or "application/octet-stream"
|
||
title_clean = title.strip()
|
||
|
||
# Resolver workers
|
||
selected: List[str] = []
|
||
raw_ids = (worker_ids or "").strip()
|
||
if raw_ids:
|
||
try:
|
||
parsed = _json.loads(raw_ids)
|
||
if isinstance(parsed, list):
|
||
selected = [str(x).strip() for x in parsed if str(x).strip()]
|
||
elif isinstance(parsed, str) and parsed.strip():
|
||
selected = [parsed.strip()]
|
||
except _json.JSONDecodeError:
|
||
selected = [x.strip() for x in raw_ids.split(",") if x.strip()]
|
||
if not selected and worker_id.strip():
|
||
selected = [worker_id.strip()]
|
||
|
||
online = store.list_online_worker_ids()
|
||
online_set = set(online)
|
||
if not selected:
|
||
selected = list(online)
|
||
else:
|
||
missing = [w for w in selected if w not in online_set]
|
||
if missing:
|
||
raise HTTPException(
|
||
status_code=409,
|
||
detail=f"Workers no online: {', '.join(missing)}. Enciende el agent primero.",
|
||
)
|
||
|
||
if not selected:
|
||
raise HTTPException(
|
||
status_code=409,
|
||
detail="No hay bots online. Enciende al menos un agent.",
|
||
)
|
||
|
||
# Dedup preservando orden
|
||
seen: set[str] = set()
|
||
workers_unique: List[str] = []
|
||
for w in selected:
|
||
if w not in seen:
|
||
seen.add(w)
|
||
workers_unique.append(w)
|
||
|
||
# Parse filas
|
||
try:
|
||
rows = parse_radicado_rows(data, original)
|
||
except Exception as e:
|
||
raise HTTPException(status_code=400, detail=f"No se pudo leer el Excel/CSV: {e}") from e
|
||
if not rows:
|
||
raise HTTPException(status_code=400, detail="El archivo no tiene filas de radicado/NIT")
|
||
|
||
# Un solo bot (o una sola fila) → flujo actual
|
||
n_workers = min(len(workers_unique), len(rows))
|
||
if n_workers <= 1:
|
||
only = workers_unique[0]
|
||
job = store.create_job(
|
||
worker_id=only,
|
||
excel_key="",
|
||
original_filename=original,
|
||
title=title_clean,
|
||
job_id=parent_job_id,
|
||
)
|
||
task_id = store.create_downloads_task(
|
||
parent_job_id=parent_job_id,
|
||
download_job_id=job.id,
|
||
worker_id=only,
|
||
title=title_clean or f"Bot descarga ({original})",
|
||
original_filename=original,
|
||
worker_ids=[only],
|
||
mode="single",
|
||
)
|
||
try:
|
||
key = minio.upload_excel(job.id, data, original, content_type=content_type)
|
||
except Exception as e:
|
||
store.patch_job(job.id, status=JobStatus.FAILED.value, error=f"MinIO upload failed: {e}")
|
||
raise HTTPException(status_code=502, detail=f"Error subiendo a MinIO: {e}") from e
|
||
|
||
job = store.update_excel_key(job.id, key)
|
||
assert job is not None
|
||
store.add_event(
|
||
job.id,
|
||
"run_started",
|
||
f"Job creado y asignado a {only}",
|
||
metadata={
|
||
"excel_key": key,
|
||
"worker_id": only,
|
||
"job_id": parent_job_id,
|
||
"task_id": task_id,
|
||
"mode": "single",
|
||
"total_rows": len(rows),
|
||
},
|
||
)
|
||
return job.to_dict()
|
||
|
||
# Distribuido: padre + shards
|
||
workers_unique = workers_unique[:n_workers]
|
||
chunks = split_rows(rows, n_workers)
|
||
n = len(chunks)
|
||
workers_unique = workers_unique[:n]
|
||
|
||
parent = store.create_job(
|
||
worker_id=None,
|
||
excel_key="",
|
||
original_filename=original,
|
||
title=title_clean or f"Descarga distribuida ({original})",
|
||
job_id=parent_job_id,
|
||
shard_count=n,
|
||
status=JobStatus.PENDING.value,
|
||
result={
|
||
"mode": "distributed",
|
||
"shard_ids": [],
|
||
"total_rows": len(rows),
|
||
"total": len(rows),
|
||
"processed": 0,
|
||
"exitosos": 0,
|
||
"fallidos": 0,
|
||
},
|
||
)
|
||
prefix = parent.downloads_prefix or mk_prefix(parent.id)
|
||
|
||
task_id = store.create_downloads_task(
|
||
parent_job_id=parent_job_id,
|
||
download_job_id=parent.id,
|
||
worker_id=None,
|
||
title=title_clean or f"Bot descarga ×{n} ({original})",
|
||
original_filename=original,
|
||
worker_ids=workers_unique,
|
||
mode="distributed",
|
||
)
|
||
|
||
try:
|
||
full_key = minio.upload_excel(parent.id, data, original, content_type=content_type)
|
||
except Exception as e:
|
||
store.patch_job(parent.id, status=JobStatus.FAILED.value, error=f"MinIO upload failed: {e}")
|
||
raise HTTPException(status_code=502, detail=f"Error subiendo a MinIO: {e}") from e
|
||
|
||
store.update_excel_key(parent.id, full_key)
|
||
|
||
shard_ids: List[str] = []
|
||
for i, (chunk, wid) in enumerate(zip(chunks, workers_unique)):
|
||
shard_bytes = rows_to_xlsx_bytes(chunk)
|
||
shard = store.create_job(
|
||
worker_id=wid,
|
||
excel_key="",
|
||
original_filename=f"shard_{i + 1}_of_{n}.xlsx",
|
||
title=f"{title_clean or 'Descarga'} · shard {i + 1}/{n}",
|
||
job_id=parent_job_id,
|
||
task_id=task_id,
|
||
downloads_prefix_override=prefix,
|
||
parent_download_job_id=parent.id,
|
||
shard_index=i,
|
||
shard_count=n,
|
||
result={"total": len(chunk), "processed": 0},
|
||
)
|
||
try:
|
||
shard_key = minio.upload_excel(
|
||
shard.id,
|
||
shard_bytes,
|
||
f"shard_{i + 1}_of_{n}.xlsx",
|
||
content_type=(
|
||
"application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
|
||
),
|
||
)
|
||
except Exception as e:
|
||
store.patch_job(
|
||
parent.id,
|
||
status=JobStatus.FAILED.value,
|
||
error=f"MinIO shard upload failed: {e}",
|
||
)
|
||
raise HTTPException(
|
||
status_code=502, detail=f"Error subiendo shard a MinIO: {e}"
|
||
) from e
|
||
store.update_excel_key(shard.id, shard_key)
|
||
shard_ids.append(shard.id)
|
||
store.add_event(
|
||
parent.id,
|
||
"info",
|
||
f"Shard {i + 1}/{n} → {wid} ({len(chunk)} filas)",
|
||
metadata={"shard_id": shard.id, "worker_id": wid, "rows": len(chunk)},
|
||
)
|
||
|
||
store.patch_job(
|
||
parent.id,
|
||
result={
|
||
"mode": "distributed",
|
||
"shard_ids": shard_ids,
|
||
"total_rows": len(rows),
|
||
"total": len(rows),
|
||
"processed": 0,
|
||
"exitosos": 0,
|
||
"fallidos": 0,
|
||
"downloads_prefix": prefix,
|
||
"worker_ids": workers_unique,
|
||
},
|
||
)
|
||
store.add_event(
|
||
parent.id,
|
||
"run_started",
|
||
f"Corrida distribuida: {n} bots, {len(rows)} filas",
|
||
metadata={
|
||
"excel_key": full_key,
|
||
"worker_ids": workers_unique,
|
||
"job_id": parent_job_id,
|
||
"task_id": task_id,
|
||
"mode": "distributed",
|
||
"shard_ids": shard_ids,
|
||
"total_rows": len(rows),
|
||
},
|
||
)
|
||
job = store.get_job(parent.id)
|
||
assert job is not None
|
||
return job.to_dict()
|
||
|
||
|
||
@app.post("/internal/reclaim")
|
||
def reclaim_jobs(_: str = Depends(require_bearer)) -> dict:
|
||
n = store.reclaim_stale_running_jobs()
|
||
return {"reclaimed": n}
|
||
|
||
|
||
@app.get("/jobs")
|
||
def list_jobs(
|
||
limit: int = Query(100, ge=1, le=500),
|
||
job_id: Optional[str] = Query(None, description="Filtrar por Job Desktop"),
|
||
_: str = Depends(require_bearer),
|
||
) -> List[dict]:
|
||
return [j.to_dict() for j in store.list_jobs(limit=limit, job_id=job_id)]
|
||
|
||
|
||
@app.get("/jobs/{download_job_id}")
|
||
def get_job(download_job_id: str, _: str = Depends(require_bearer)) -> dict:
|
||
job = store.get_job(download_job_id)
|
||
if not job:
|
||
raise HTTPException(status_code=404, detail="Job no encontrado")
|
||
# Refrescar consolidación si es padre distribuido
|
||
if job.is_distributed_parent:
|
||
refreshed = store.refresh_parent_from_shards(download_job_id)
|
||
if refreshed:
|
||
job = refreshed
|
||
payload = job.to_dict()
|
||
if job.is_distributed_parent:
|
||
payload["shards"] = [s.to_dict() for s in store.list_shards(job.id)]
|
||
return payload
|
||
|
||
|
||
@app.get("/jobs/{download_job_id}/files")
|
||
def list_job_files(download_job_id: str, _: str = Depends(require_bearer)) -> dict:
|
||
job = store.get_job(download_job_id)
|
||
if not job:
|
||
raise HTTPException(status_code=404, detail="Job no encontrado")
|
||
try:
|
||
return minio.list_download_files(download_job_id, prefix=job.downloads_prefix)
|
||
except Exception as e:
|
||
raise HTTPException(status_code=502, detail=f"Error listando MinIO: {e}") from e
|
||
|
||
|
||
@app.get("/jobs/{download_job_id}/object")
|
||
def get_job_object(
|
||
download_job_id: str,
|
||
key: str = Query(..., min_length=1, description="Object key bajo downloads_prefix"),
|
||
_: str = Depends(require_bearer),
|
||
):
|
||
"""Stream de un PDF del resultado MinIO (solo keys del prefijo del job)."""
|
||
job = store.get_job(download_job_id)
|
||
if not job:
|
||
raise HTTPException(status_code=404, detail="Job no encontrado")
|
||
|
||
prefix = (job.downloads_prefix or f"downloads/{download_job_id}/").lstrip("/")
|
||
if not prefix.endswith("/"):
|
||
prefix += "/"
|
||
normalized = key.lstrip("/")
|
||
if not normalized.startswith(prefix):
|
||
raise HTTPException(status_code=403, detail="Key fuera del prefijo del job")
|
||
if not normalized.lower().endswith(".pdf"):
|
||
raise HTTPException(status_code=400, detail="Solo se permiten objetos PDF")
|
||
if ".." in normalized.split("/"):
|
||
raise HTTPException(status_code=400, detail="Key no permitida")
|
||
if normalized.endswith("_manifest.json"):
|
||
raise HTTPException(status_code=400, detail="Key no permitida")
|
||
|
||
try:
|
||
data = minio.get_bytes(normalized)
|
||
except Exception as e:
|
||
raise HTTPException(status_code=502, detail=f"Error leyendo MinIO: {e}") from e
|
||
|
||
filename = Path(normalized).name
|
||
return Response(
|
||
content=data,
|
||
media_type="application/pdf",
|
||
headers={"Content-Disposition": f'inline; filename="{filename}"'},
|
||
)
|
||
|
||
|
||
@app.get("/jobs/{download_job_id}/events")
|
||
def get_events(
|
||
download_job_id: str,
|
||
since: int = Query(0, ge=0),
|
||
_: str = Depends(require_bearer),
|
||
) -> dict:
|
||
if not store.get_job(download_job_id):
|
||
raise HTTPException(status_code=404, detail="Job no encontrado")
|
||
events = store.list_events(download_job_id, since_id=since)
|
||
return {"events": [e.to_dict() for e in events]}
|
||
|
||
|
||
@app.post("/jobs/{download_job_id}/cancel")
|
||
def cancel_job(download_job_id: str, _: str = Depends(require_bearer)) -> dict:
|
||
job = store.request_cancel(download_job_id)
|
||
if not job:
|
||
raise HTTPException(status_code=404, detail="Job no encontrado")
|
||
store.add_event(download_job_id, "warning", "Cancelación solicitada por el usuario")
|
||
return job.to_dict()
|
||
|
||
|
||
@app.get("/jobs/{download_job_id}/file")
|
||
def get_job_file(download_job_id: str, _: str = Depends(require_bearer)):
|
||
job = store.get_job(download_job_id)
|
||
if not job:
|
||
raise HTTPException(status_code=404, detail="Job no encontrado")
|
||
if not job.excel_key:
|
||
raise HTTPException(status_code=404, detail="Excel no disponible")
|
||
|
||
try:
|
||
data = minio.get_bytes(job.excel_key)
|
||
except Exception as e:
|
||
raise HTTPException(status_code=502, detail=f"Error leyendo MinIO: {e}") from e
|
||
|
||
filename = Path(job.excel_key).name
|
||
return Response(
|
||
content=data,
|
||
media_type="application/octet-stream",
|
||
headers={"Content-Disposition": f'attachment; filename="{filename}"'},
|
||
)
|
||
|
||
|
||
@app.get("/jobs/{download_job_id}/file-url")
|
||
def get_job_file_url(download_job_id: str, _: str = Depends(require_bearer)) -> dict:
|
||
job = store.get_job(download_job_id)
|
||
if not job or not job.excel_key:
|
||
raise HTTPException(status_code=404, detail="Job/archivo no encontrado")
|
||
try:
|
||
url = minio.presigned_get(job.excel_key)
|
||
except Exception as e:
|
||
raise HTTPException(status_code=502, detail=str(e)) from e
|
||
return {"url": url, "excel_key": job.excel_key}
|
||
|
||
|
||
@app.post("/jobs/{download_job_id}/events")
|
||
def post_event(
|
||
download_job_id: str,
|
||
body: EventBody,
|
||
_: str = Depends(require_bearer),
|
||
) -> dict:
|
||
job = store.get_job(download_job_id)
|
||
if not job:
|
||
raise HTTPException(status_code=404, detail="Job no encontrado")
|
||
event = store.add_event(
|
||
download_job_id,
|
||
body.event_type,
|
||
body.message,
|
||
radicado=body.radicado,
|
||
nit=body.nit,
|
||
index=body.index,
|
||
total=body.total,
|
||
metadata=body.metadata,
|
||
timestamp=body.timestamp,
|
||
)
|
||
if body.index is not None and body.total is not None:
|
||
result = dict(job.result or {})
|
||
result["processed"] = body.index
|
||
result["total"] = body.total
|
||
if body.nit:
|
||
result["nit_actual"] = body.nit
|
||
store.patch_job(download_job_id, result=result)
|
||
return event.to_dict()
|
||
|
||
|
||
@app.patch("/jobs/{download_job_id}")
|
||
def patch_job(
|
||
download_job_id: str,
|
||
body: PatchJobBody,
|
||
_: str = Depends(require_bearer),
|
||
) -> dict:
|
||
if body.status and body.status not in {s.value for s in JobStatus}:
|
||
raise HTTPException(status_code=400, detail="status inválido")
|
||
job = store.patch_job(
|
||
download_job_id,
|
||
status=body.status,
|
||
result=body.result,
|
||
error=body.error,
|
||
)
|
||
if not job:
|
||
raise HTTPException(status_code=404, detail="Job no encontrado")
|
||
return job.to_dict()
|
||
|
||
|
||
@app.get("/jobs/{download_job_id}/cancel-flag")
|
||
def cancel_flag(download_job_id: str, _: str = Depends(require_bearer)) -> dict:
|
||
job = store.get_job(download_job_id)
|
||
if not job:
|
||
raise HTTPException(status_code=404, detail="Job no encontrado")
|
||
return {
|
||
"cancel_requested": store.is_cancel_requested(download_job_id),
|
||
"status": job.status,
|
||
}
|