268 lines
9.0 KiB
Python
268 lines
9.0 KiB
Python
from __future__ import annotations
|
|
|
|
import socket
|
|
import tempfile
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Any, Dict, List, Optional
|
|
|
|
import httpx
|
|
|
|
from bot_descarga.core.download_service import DownloadService
|
|
from bot_descarga.worker.minio_upload import MinioSettings, WorkerMinioUploader
|
|
|
|
JOB_COMPLETED = "completed"
|
|
JOB_FAILED = "failed"
|
|
JOB_CANCELLED = "cancelled"
|
|
|
|
|
|
class RemoteApiClient:
|
|
"""Cliente HTTP hacia la API en IPSExtractor (vía gateway). Sin puertos locales."""
|
|
|
|
def __init__(self, base_url: str, token: str):
|
|
self.base_url = base_url.rstrip("/")
|
|
self.token = token
|
|
self._client = httpx.Client(
|
|
timeout=httpx.Timeout(60.0, connect=10.0),
|
|
headers={"Authorization": f"Bearer {token}"},
|
|
)
|
|
|
|
def close(self) -> None:
|
|
self._client.close()
|
|
|
|
def _url(self, path: str) -> str:
|
|
return f"{self.base_url}{path}"
|
|
|
|
def register(self, worker_id: str, hostname: str = "", max_concurrent: int = 1) -> dict:
|
|
r = self._client.post(
|
|
self._url("/workers/register"),
|
|
json={
|
|
"worker_id": worker_id,
|
|
"hostname": hostname or socket.gethostname(),
|
|
"max_concurrent": max_concurrent,
|
|
},
|
|
)
|
|
r.raise_for_status()
|
|
return r.json()
|
|
|
|
def heartbeat(self, worker_id: str) -> dict:
|
|
r = self._client.post(
|
|
self._url("/workers/heartbeat"),
|
|
json={"worker_id": worker_id},
|
|
)
|
|
r.raise_for_status()
|
|
return r.json()
|
|
|
|
def claim_pending(self, worker_id: str) -> Optional[dict]:
|
|
r = self._client.get(
|
|
self._url("/workers/me/pending"),
|
|
params={"worker_id": worker_id},
|
|
)
|
|
r.raise_for_status()
|
|
data = r.json()
|
|
return data.get("job")
|
|
|
|
def download_excel(self, job_id: str, dest: Path) -> Path:
|
|
r = self._client.get(self._url(f"/jobs/{job_id}/file"))
|
|
r.raise_for_status()
|
|
dest.parent.mkdir(parents=True, exist_ok=True)
|
|
dest.write_bytes(r.content)
|
|
return dest
|
|
|
|
def post_event(self, job_id: str, event: Dict[str, Any]) -> None:
|
|
r = self._client.post(self._url(f"/jobs/{job_id}/events"), json=event)
|
|
r.raise_for_status()
|
|
|
|
def cancel_flag(self, job_id: str) -> bool:
|
|
r = self._client.get(self._url(f"/jobs/{job_id}/cancel-flag"))
|
|
r.raise_for_status()
|
|
return bool(r.json().get("cancel_requested"))
|
|
|
|
def patch_job(self, job_id: str, **body) -> dict:
|
|
r = self._client.patch(self._url(f"/jobs/{job_id}"), json=body)
|
|
r.raise_for_status()
|
|
return r.json()
|
|
|
|
|
|
class WorkerAgent:
|
|
"""Agent idle: heartbeat + claim solo jobs asignados por el usuario."""
|
|
|
|
def __init__(
|
|
self,
|
|
api_url: str,
|
|
worker_id: str,
|
|
token: str,
|
|
poll_seconds: float = 8.0,
|
|
staging_root: Optional[Path] = None,
|
|
log_func=None,
|
|
):
|
|
self.api = RemoteApiClient(api_url, token)
|
|
self.worker_id = worker_id
|
|
self.poll_seconds = poll_seconds
|
|
self.staging_root = staging_root or Path(tempfile.gettempdir()) / "bot_descarga"
|
|
self.minio = WorkerMinioUploader(MinioSettings.from_env())
|
|
self._running = True
|
|
self._uploaded_keys: List[str] = []
|
|
self._log = log_func or print
|
|
|
|
def stop(self) -> None:
|
|
self._running = False
|
|
|
|
def run_forever(self) -> None:
|
|
self.api.register(self.worker_id)
|
|
self._log(f"[worker {self.worker_id}] registrado — idle (esperando jobs manuales)")
|
|
try:
|
|
while self._running:
|
|
try:
|
|
self.api.heartbeat(self.worker_id)
|
|
job = self.api.claim_pending(self.worker_id)
|
|
if job:
|
|
self._run_job(job)
|
|
else:
|
|
time.sleep(self.poll_seconds)
|
|
except KeyboardInterrupt:
|
|
raise
|
|
except Exception as e:
|
|
self._log(f"[worker {self.worker_id}] error en loop: {e}")
|
|
time.sleep(self.poll_seconds)
|
|
finally:
|
|
self.api.close()
|
|
self._log(f"[worker {self.worker_id}] detenido")
|
|
|
|
def _run_job(self, job: dict) -> None:
|
|
job_id = job["id"]
|
|
downloads_prefix = (job.get("downloads_prefix") or "").strip() or f"downloads/{job_id}/"
|
|
self._uploaded_keys = []
|
|
work_dir = self.staging_root / job_id
|
|
out_dir = work_dir / "out"
|
|
out_dir.mkdir(parents=True, exist_ok=True)
|
|
excel_path = work_dir / (Path(job.get("excel_key") or "input.xlsx").name)
|
|
|
|
self._log(f"[worker {self.worker_id}] job {job_id} iniciado (prefix={downloads_prefix})")
|
|
|
|
try:
|
|
self.api.download_excel(job_id, excel_path)
|
|
|
|
def should_cancel() -> bool:
|
|
try:
|
|
return self.api.cancel_flag(job_id)
|
|
except Exception:
|
|
return False
|
|
|
|
def on_event(event) -> None:
|
|
self._log(event.display_message())
|
|
payload = {
|
|
"event_type": event.event_type.value
|
|
if hasattr(event.event_type, "value")
|
|
else str(event.event_type),
|
|
"message": event.message,
|
|
"radicado": event.radicado,
|
|
"nit": event.nit,
|
|
"index": event.index,
|
|
"total": event.total,
|
|
"metadata": event.metadata,
|
|
"timestamp": event.timestamp,
|
|
}
|
|
try:
|
|
self.api.post_event(job_id, payload)
|
|
except Exception as e:
|
|
self._log(f"[worker] no se pudo enviar evento: {e}")
|
|
|
|
def on_pdf_saved(nit: str, radicado: str, local_path: str) -> None:
|
|
key = self.minio.upload_pdf(
|
|
job_id,
|
|
nit,
|
|
radicado,
|
|
local_path,
|
|
downloads_prefix=downloads_prefix,
|
|
skip_if_exists=True,
|
|
)
|
|
self._uploaded_keys.append(key)
|
|
|
|
service = DownloadService(
|
|
ruta_excel=str(excel_path),
|
|
carpeta_descarga=str(out_dir),
|
|
event_callback=on_event,
|
|
should_cancel=should_cancel,
|
|
on_pdf_saved=on_pdf_saved,
|
|
)
|
|
summary = service.run()
|
|
|
|
try:
|
|
self.minio.write_manifest(
|
|
job_id,
|
|
self._uploaded_keys,
|
|
downloads_prefix=downloads_prefix,
|
|
)
|
|
except Exception as e:
|
|
self._log(f"[worker] manifest no escrito: {e}")
|
|
|
|
fallidos_count = len(summary.fallidos)
|
|
if summary.fatal_error and fallidos_count == 0 and summary.total:
|
|
fallidos_count = summary.total - summary.exitosos
|
|
|
|
result = {
|
|
"total": summary.total,
|
|
"exitosos": summary.exitosos,
|
|
"fallidos": fallidos_count,
|
|
"cancelled": summary.cancelled,
|
|
"downloads_prefix": downloads_prefix,
|
|
"uploaded_count": len(self._uploaded_keys),
|
|
}
|
|
|
|
if summary.cancelled:
|
|
status = JOB_CANCELLED
|
|
error_msg = None
|
|
elif summary.fatal_error:
|
|
status = JOB_FAILED
|
|
error_msg = summary.fatal_error
|
|
elif summary.fallidos and summary.exitosos == 0:
|
|
status = JOB_FAILED
|
|
error_msg = "Todos los radicados fallaron"
|
|
else:
|
|
status = JOB_COMPLETED
|
|
error_msg = None
|
|
|
|
if error_msg:
|
|
try:
|
|
self.api.post_event(
|
|
job_id,
|
|
{
|
|
"event_type": "warning",
|
|
"message": error_msg,
|
|
"metadata": {"fatal": bool(summary.fatal_error)},
|
|
},
|
|
)
|
|
except Exception:
|
|
pass
|
|
|
|
self.api.patch_job(
|
|
job_id,
|
|
status=status,
|
|
result=result,
|
|
error=error_msg,
|
|
)
|
|
self._log(f"[worker {self.worker_id}] job {job_id} → {status}")
|
|
|
|
except Exception as e:
|
|
self._log(f"[worker {self.worker_id}] job {job_id} error: {e}")
|
|
try:
|
|
self.api.post_event(
|
|
job_id,
|
|
{
|
|
"event_type": "warning",
|
|
"message": f"Error fatal: {e}",
|
|
"metadata": {"fatal": True},
|
|
},
|
|
)
|
|
except Exception:
|
|
pass
|
|
try:
|
|
self.api.patch_job(
|
|
job_id,
|
|
status=JOB_FAILED,
|
|
error=str(e),
|
|
)
|
|
except Exception:
|
|
pass
|