601 lines
20 KiB
Python
601 lines
20 KiB
Python
from __future__ import annotations
|
||
|
||
import os
|
||
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())
|
||
|
||
POBLACION_APTA_URL = os.getenv("POBLACION_APTA_URL", "http://poblacion_apta:8003").rstrip("/")
|
||
|
||
app = FastAPI(title="BotDigitador 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
|
||
|
||
|
||
def _media_type_for_key(key: str) -> str:
|
||
lower = key.lower()
|
||
if lower.endswith(".csv"):
|
||
return "text/csv; charset=utf-8"
|
||
if lower.endswith(".log") or lower.endswith(".txt"):
|
||
return "text/plain; charset=utf-8"
|
||
if lower.endswith(".json"):
|
||
return "application/json"
|
||
return "application/octet-stream"
|
||
|
||
|
||
@app.get("/health")
|
||
def health() -> dict:
|
||
return {"status": "ok", "service": "bot-digitador"}
|
||
|
||
|
||
@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(
|
||
job_id: str = Form(...),
|
||
worker_id: str = Form(""),
|
||
worker_ids: str = Form(""),
|
||
title: str = Form(""),
|
||
file: Optional[UploadFile] = File(None),
|
||
poblacion_apta_job_id: str = Form(""),
|
||
subset: str = Form("aptos"),
|
||
_: str = Depends(require_bearer),
|
||
) -> dict:
|
||
"""Crea una corrida de digitación.
|
||
|
||
Exactamente uno de: file (CSV) o poblacion_apta_job_id.
|
||
- 1 worker → flujo clásico.
|
||
- N workers → padre + N shards CSV, mismo reports_prefix del padre.
|
||
"""
|
||
import json as _json
|
||
|
||
import httpx
|
||
|
||
from domain.csv_split import parse_digitador_rows, rows_to_csv_bytes, split_rows
|
||
from infrastructure.storage_minio import reports_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")
|
||
|
||
pa_id = (poblacion_apta_job_id or "").strip()
|
||
data: bytes = b""
|
||
original = "input.csv"
|
||
content_type = "text/csv"
|
||
source_type = "upload"
|
||
|
||
file_bytes: bytes = b""
|
||
if file is not None:
|
||
file_bytes = await file.read()
|
||
|
||
if pa_id and file_bytes:
|
||
raise HTTPException(
|
||
status_code=400,
|
||
detail="Proporcione exactamente uno de: file o poblacion_apta_job_id",
|
||
)
|
||
|
||
if pa_id:
|
||
source_type = "poblacion_apta"
|
||
subset_norm = (subset or "aptos").strip().lower()
|
||
if subset_norm != "aptos":
|
||
raise HTTPException(
|
||
status_code=400,
|
||
detail="Solo se digitan aptos (PASO_BOT). Use subset=aptos.",
|
||
)
|
||
url = f"{POBLACION_APTA_URL}/jobs/{pa_id}/export.csv"
|
||
try:
|
||
async with httpx.AsyncClient(timeout=120.0) as client:
|
||
resp = await client.get(url, params={"subset": subset_norm})
|
||
except Exception as e:
|
||
raise HTTPException(
|
||
status_code=502,
|
||
detail=f"No se pudo contactar poblacion_apta: {e}",
|
||
) from e
|
||
if resp.status_code == 404:
|
||
raise HTTPException(status_code=404, detail="Job de población apta no encontrado")
|
||
if not resp.is_success:
|
||
raise HTTPException(
|
||
status_code=502,
|
||
detail=f"Error exportando CSV de población apta: HTTP {resp.status_code}",
|
||
)
|
||
data = resp.content
|
||
original = f"poblacion_apta_{subset_norm}_{pa_id[:8]}.csv"
|
||
content_type = "text/csv; charset=utf-8"
|
||
elif file_bytes:
|
||
data = file_bytes
|
||
original = (file.filename if file else None) or "input.csv"
|
||
content_type = (file.content_type if file else None) or "text/csv"
|
||
source_type = "upload"
|
||
else:
|
||
raise HTTPException(
|
||
status_code=400,
|
||
detail="Proporcione exactamente uno de: file o poblacion_apta_job_id",
|
||
)
|
||
|
||
if not data:
|
||
raise HTTPException(status_code=400, detail="CSV vacío")
|
||
|
||
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.",
|
||
)
|
||
|
||
seen: set[str] = set()
|
||
workers_unique: List[str] = []
|
||
for w in selected:
|
||
if w not in seen:
|
||
seen.add(w)
|
||
workers_unique.append(w)
|
||
|
||
try:
|
||
rows = parse_digitador_rows(data, original)
|
||
except ValueError as e:
|
||
raise HTTPException(status_code=400, detail=str(e)) from e
|
||
except Exception as e:
|
||
raise HTTPException(status_code=400, detail=f"No se pudo leer el CSV: {e}") from e
|
||
if not rows:
|
||
raise HTTPException(status_code=400, detail="El CSV no tiene filas para digitar")
|
||
|
||
# Normalizar a CSV canónico del bot (siempre pipe + headers snake_case)
|
||
canonical_bytes = rows_to_csv_bytes(rows)
|
||
|
||
n_workers = min(len(workers_unique), len(rows))
|
||
if n_workers <= 1:
|
||
only = workers_unique[0]
|
||
job = store.create_job(
|
||
worker_id=only,
|
||
csv_key="",
|
||
original_filename=original,
|
||
title=title_clean,
|
||
job_id=parent_job_id,
|
||
result={"source_type": source_type, "total_rows": len(rows), "total": len(rows)},
|
||
)
|
||
task_id = store.create_digitador_task(
|
||
parent_job_id=parent_job_id,
|
||
digitador_job_id=job.id,
|
||
worker_id=only,
|
||
title=title_clean or f"Bot digitador ({original})",
|
||
original_filename=original,
|
||
worker_ids=[only],
|
||
mode="single",
|
||
source_type=source_type,
|
||
)
|
||
try:
|
||
key = minio.upload_csv(job.id, canonical_bytes, 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_csv_key(job.id, key)
|
||
assert job is not None
|
||
store.add_event(
|
||
job.id,
|
||
"run_started",
|
||
f"Job creado y asignado a {only}",
|
||
metadata={
|
||
"csv_key": key,
|
||
"worker_id": only,
|
||
"job_id": parent_job_id,
|
||
"task_id": task_id,
|
||
"mode": "single",
|
||
"total_rows": len(rows),
|
||
"source_type": source_type,
|
||
},
|
||
)
|
||
return job.to_dict()
|
||
|
||
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,
|
||
csv_key="",
|
||
original_filename=original,
|
||
title=title_clean or f"Digitación 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,
|
||
"ok": 0,
|
||
"error": 0,
|
||
"timeout": 0,
|
||
"existentes": 0,
|
||
"source_type": source_type,
|
||
},
|
||
)
|
||
prefix = parent.reports_prefix or mk_prefix(parent.id)
|
||
|
||
task_id = store.create_digitador_task(
|
||
parent_job_id=parent_job_id,
|
||
digitador_job_id=parent.id,
|
||
worker_id=None,
|
||
title=title_clean or f"Bot digitador ×{n} ({original})",
|
||
original_filename=original,
|
||
worker_ids=workers_unique,
|
||
mode="distributed",
|
||
source_type=source_type,
|
||
)
|
||
|
||
try:
|
||
full_key = minio.upload_csv(parent.id, canonical_bytes, 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_csv_key(parent.id, full_key)
|
||
|
||
shard_ids: List[str] = []
|
||
for i, (chunk, wid) in enumerate(zip(chunks, workers_unique)):
|
||
shard_bytes = rows_to_csv_bytes(chunk)
|
||
shard = store.create_job(
|
||
worker_id=wid,
|
||
csv_key="",
|
||
original_filename=f"shard_{i + 1}_of_{n}.csv",
|
||
title=f"{title_clean or 'Digitador'} · shard {i + 1}/{n}",
|
||
job_id=parent_job_id,
|
||
task_id=task_id,
|
||
reports_prefix_override=prefix,
|
||
parent_digitador_job_id=parent.id,
|
||
shard_index=i,
|
||
shard_count=n,
|
||
result={"total": len(chunk), "processed": 0, "source_type": source_type},
|
||
)
|
||
try:
|
||
shard_key = minio.upload_csv(
|
||
shard.id,
|
||
shard_bytes,
|
||
f"shard_{i + 1}_of_{n}.csv",
|
||
content_type="text/csv; charset=utf-8",
|
||
)
|
||
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_csv_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,
|
||
"ok": 0,
|
||
"error": 0,
|
||
"timeout": 0,
|
||
"existentes": 0,
|
||
"reports_prefix": prefix,
|
||
"worker_ids": workers_unique,
|
||
"source_type": source_type,
|
||
},
|
||
)
|
||
store.add_event(
|
||
parent.id,
|
||
"run_started",
|
||
f"Corrida distribuida: {n} bots, {len(rows)} filas",
|
||
metadata={
|
||
"csv_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),
|
||
"source_type": source_type,
|
||
},
|
||
)
|
||
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/{digitador_job_id}")
|
||
def get_job(digitador_job_id: str, _: str = Depends(require_bearer)) -> dict:
|
||
job = store.get_job(digitador_job_id)
|
||
if not job:
|
||
raise HTTPException(status_code=404, detail="Job no encontrado")
|
||
if job.is_distributed_parent:
|
||
refreshed = store.refresh_parent_from_shards(digitador_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/{digitador_job_id}/files")
|
||
def list_job_files(digitador_job_id: str, _: str = Depends(require_bearer)) -> dict:
|
||
job = store.get_job(digitador_job_id)
|
||
if not job:
|
||
raise HTTPException(status_code=404, detail="Job no encontrado")
|
||
try:
|
||
return minio.list_report_files(digitador_job_id, prefix=job.reports_prefix)
|
||
except Exception as e:
|
||
raise HTTPException(status_code=502, detail=f"Error listando MinIO: {e}") from e
|
||
|
||
|
||
@app.get("/jobs/{digitador_job_id}/object")
|
||
def get_job_object(
|
||
digitador_job_id: str,
|
||
key: str = Query(..., min_length=1, description="Object key bajo reports_prefix"),
|
||
_: str = Depends(require_bearer),
|
||
):
|
||
"""Stream de un reporte MinIO (csv/log/json bajo el prefijo del job)."""
|
||
job = store.get_job(digitador_job_id)
|
||
if not job:
|
||
raise HTTPException(status_code=404, detail="Job no encontrado")
|
||
|
||
prefix = (job.reports_prefix or f"digitados/{digitador_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")
|
||
lower = normalized.lower()
|
||
if not (lower.endswith(".csv") or lower.endswith(".log") or lower.endswith(".json")):
|
||
raise HTTPException(status_code=400, detail="Solo se permiten objetos csv/log/json")
|
||
if ".." in normalized.split("/"):
|
||
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=_media_type_for_key(normalized),
|
||
headers={"Content-Disposition": f'inline; filename="{filename}"'},
|
||
)
|
||
|
||
|
||
@app.get("/jobs/{digitador_job_id}/events")
|
||
def get_events(
|
||
digitador_job_id: str,
|
||
since: int = Query(0, ge=0),
|
||
_: str = Depends(require_bearer),
|
||
) -> dict:
|
||
if not store.get_job(digitador_job_id):
|
||
raise HTTPException(status_code=404, detail="Job no encontrado")
|
||
events = store.list_events(digitador_job_id, since_id=since)
|
||
return {"events": [e.to_dict() for e in events]}
|
||
|
||
|
||
@app.post("/jobs/{digitador_job_id}/cancel")
|
||
def cancel_job(digitador_job_id: str, _: str = Depends(require_bearer)) -> dict:
|
||
job = store.request_cancel(digitador_job_id)
|
||
if not job:
|
||
raise HTTPException(status_code=404, detail="Job no encontrado")
|
||
store.add_event(digitador_job_id, "warning", "Cancelación solicitada por el usuario")
|
||
return job.to_dict()
|
||
|
||
|
||
@app.get("/jobs/{digitador_job_id}/file")
|
||
def get_job_file(digitador_job_id: str, _: str = Depends(require_bearer)):
|
||
job = store.get_job(digitador_job_id)
|
||
if not job:
|
||
raise HTTPException(status_code=404, detail="Job no encontrado")
|
||
if not job.csv_key:
|
||
raise HTTPException(status_code=404, detail="CSV no disponible")
|
||
|
||
try:
|
||
data = minio.get_bytes(job.csv_key)
|
||
except Exception as e:
|
||
raise HTTPException(status_code=502, detail=f"Error leyendo MinIO: {e}") from e
|
||
|
||
filename = Path(job.csv_key).name
|
||
return Response(
|
||
content=data,
|
||
media_type="text/csv; charset=utf-8",
|
||
headers={"Content-Disposition": f'attachment; filename="{filename}"'},
|
||
)
|
||
|
||
|
||
@app.get("/jobs/{digitador_job_id}/file-url")
|
||
def get_job_file_url(digitador_job_id: str, _: str = Depends(require_bearer)) -> dict:
|
||
job = store.get_job(digitador_job_id)
|
||
if not job or not job.csv_key:
|
||
raise HTTPException(status_code=404, detail="Job/archivo no encontrado")
|
||
try:
|
||
url = minio.presigned_get(job.csv_key)
|
||
except Exception as e:
|
||
raise HTTPException(status_code=502, detail=str(e)) from e
|
||
return {"url": url, "csv_key": job.csv_key}
|
||
|
||
|
||
@app.post("/jobs/{digitador_job_id}/events")
|
||
def post_event(
|
||
digitador_job_id: str,
|
||
body: EventBody,
|
||
_: str = Depends(require_bearer),
|
||
) -> dict:
|
||
job = store.get_job(digitador_job_id)
|
||
if not job:
|
||
raise HTTPException(status_code=404, detail="Job no encontrado")
|
||
event = store.add_event(
|
||
digitador_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(digitador_job_id, result=result)
|
||
return event.to_dict()
|
||
|
||
|
||
@app.patch("/jobs/{digitador_job_id}")
|
||
def patch_job(
|
||
digitador_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(
|
||
digitador_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/{digitador_job_id}/cancel-flag")
|
||
def cancel_flag(digitador_job_id: str, _: str = Depends(require_bearer)) -> dict:
|
||
job = store.get_job(digitador_job_id)
|
||
if not job:
|
||
raise HTTPException(status_code=404, detail="Job no encontrado")
|
||
return {
|
||
"cancel_requested": store.is_cancel_requested(digitador_job_id),
|
||
"status": job.status,
|
||
}
|