"""Repositorio sobre PostgreSQL (tablas download_* + sync de tasks).""" from __future__ import annotations from datetime import datetime, timedelta, timezone from typing import Any, Dict, List, Optional from sqlalchemy import select from domain.models import ( Job, JobEventRecord, JobStatus, Worker, WorkerStatus, utc_now, ) from infrastructure.database import ( DownloadJobEventRow, DownloadJobRow, DownloadWorkerRow, Task, _iso, init_db, session_scope, utcnow, ) from infrastructure.storage_minio import downloads_prefix HEARTBEAT_OFFLINE_SECONDS = 45 def _parse_iso(value: str | None) -> datetime | None: if not value: return None dt = datetime.fromisoformat(value) if dt.tzinfo is None: dt = dt.replace(tzinfo=timezone.utc) return dt class JobStore: """API de persistencia compatible con el store SQLite anterior.""" def __init__(self, db_path: str | None = None): # db_path ignorado (compat); usa DATABASE_URL. init_db() def _row_worker(self, row: DownloadWorkerRow) -> Worker: worker = Worker( id=row.id, hostname=row.hostname or "", max_concurrent=int(row.max_concurrent or 1), status=row.status, last_heartbeat=_iso(row.last_heartbeat) or utc_now(), current_job_id=row.current_job_id, ) return self._refresh_worker_status(worker) def _row_job(self, row: DownloadJobRow) -> Job: return Job( id=row.id, worker_id=row.worker_id, status=row.status, title=row.title or "", excel_key=row.excel_key or "", downloads_prefix=row.downloads_prefix or "", original_filename=row.original_filename or "", cancel_requested=bool(row.cancel_requested), result=dict(row.result or {}), error=row.error, created_at=_iso(row.created_at) or utc_now(), started_at=_iso(row.started_at), finished_at=_iso(row.finished_at), job_id=row.job_id, task_id=row.task_id, ) def _refresh_worker_status(self, worker: Worker) -> Worker: try: hb = datetime.fromisoformat(worker.last_heartbeat) if hb.tzinfo is None: hb = hb.replace(tzinfo=timezone.utc) age = datetime.now(timezone.utc) - hb.astimezone(timezone.utc) if age > timedelta(seconds=HEARTBEAT_OFFLINE_SECONDS): worker.status = WorkerStatus.OFFLINE.value elif worker.current_job_id: worker.status = WorkerStatus.BUSY.value else: worker.status = WorkerStatus.ONLINE.value except Exception: worker.status = WorkerStatus.OFFLINE.value return worker def register_worker( self, worker_id: str, hostname: str = "", max_concurrent: int = 1, ) -> Worker: with session_scope() as session: row = session.get(DownloadWorkerRow, worker_id) now = utcnow() if row: row.hostname = hostname row.max_concurrent = max_concurrent row.status = WorkerStatus.ONLINE.value row.last_heartbeat = now row.current_job_id = None else: row = DownloadWorkerRow( id=worker_id, hostname=hostname, max_concurrent=max_concurrent, status=WorkerStatus.ONLINE.value, last_heartbeat=now, current_job_id=None, ) session.add(row) session.flush() return self._row_worker(row) def heartbeat(self, worker_id: str) -> Optional[Worker]: with session_scope() as session: row = session.get(DownloadWorkerRow, worker_id) if not row: return None row.last_heartbeat = utcnow() row.status = ( WorkerStatus.BUSY.value if row.current_job_id else WorkerStatus.ONLINE.value ) session.flush() return self._row_worker(row) def get_worker(self, worker_id: str) -> Optional[Worker]: with session_scope() as session: row = session.get(DownloadWorkerRow, worker_id) if not row: return None return self._row_worker(row) def list_workers(self) -> List[Worker]: with session_scope() as session: rows = session.scalars( select(DownloadWorkerRow).order_by(DownloadWorkerRow.id) ).all() return [self._row_worker(r) for r in rows] def is_worker_online(self, worker_id: str) -> bool: worker = self.get_worker(worker_id) if not worker: return False return worker.status in (WorkerStatus.ONLINE.value, WorkerStatus.BUSY.value) def update_excel_key(self, job_id: str, excel_key: str) -> Optional[Job]: with session_scope() as session: row = session.get(DownloadJobRow, job_id) if not row: return None row.excel_key = excel_key session.flush() return self._row_job(row) def create_job( self, worker_id: str, excel_key: str, original_filename: str = "", title: str = "", *, job_id: Optional[str] = None, task_id: Optional[str] = None, ) -> Job: import uuid as uuid_mod download_id = str(uuid_mod.uuid4()) prefix = downloads_prefix(download_id) with session_scope() as session: row = DownloadJobRow( id=download_id, job_id=job_id, task_id=task_id, worker_id=worker_id, status=JobStatus.PENDING.value, title=title or f"Descarga {original_filename}", excel_key=excel_key, downloads_prefix=prefix, original_filename=original_filename, cancel_requested=False, result={}, created_at=utcnow(), ) session.add(row) session.flush() return self._row_job(row) def get_job(self, job_id: str) -> Optional[Job]: with session_scope() as session: row = session.get(DownloadJobRow, job_id) if not row: return None return self._row_job(row) def list_jobs( self, limit: int = 100, *, job_id: Optional[str] = None, ) -> List[Job]: with session_scope() as session: stmt = select(DownloadJobRow).order_by(DownloadJobRow.created_at.desc()).limit(limit) if job_id: stmt = ( select(DownloadJobRow) .where(DownloadJobRow.job_id == job_id) .order_by(DownloadJobRow.created_at.desc()) .limit(limit) ) rows = session.scalars(stmt).all() return [self._row_job(r) for r in rows] def claim_pending(self, worker_id: str) -> Optional[Job]: with session_scope() as session: worker = session.get(DownloadWorkerRow, worker_id) if not worker or worker.current_job_id: return None row = session.scalars( select(DownloadJobRow) .where( DownloadJobRow.worker_id == worker_id, DownloadJobRow.status == JobStatus.PENDING.value, ) .order_by(DownloadJobRow.created_at.asc()) .limit(1) ).first() if not row: return None now = utcnow() row.status = JobStatus.RUNNING.value row.started_at = now worker.current_job_id = row.id worker.status = WorkerStatus.BUSY.value worker.last_heartbeat = now if row.task_id: task = session.get(Task, row.task_id) if task: task.status = JobStatus.RUNNING.value task.started_at = now task.updated_at = now session.flush() return self._row_job(row) def request_cancel(self, job_id: str) -> Optional[Job]: with session_scope() as session: row = session.get(DownloadJobRow, job_id) if not row: return None if row.status in ( JobStatus.COMPLETED.value, JobStatus.FAILED.value, JobStatus.CANCELLED.value, ): return self._row_job(row) now = utcnow() row.cancel_requested = True if row.status == JobStatus.PENDING.value: row.status = JobStatus.CANCELLED.value row.finished_at = now self._sync_task(session, row, status=JobStatus.CANCELLED.value) session.flush() return self._row_job(row) def is_cancel_requested(self, job_id: str) -> bool: job = self.get_job(job_id) return bool(job and job.cancel_requested) def _sync_task( self, session, row: DownloadJobRow, *, status: Optional[str] = None, result: Optional[Dict[str, Any]] = None, error: Optional[str] = None, ) -> None: if not row.task_id: return task = session.get(Task, row.task_id) if not task: return now = utcnow() if status: task.status = status if result is not None: task.result = { **dict(task.result or {}), "download_job_id": row.id, **result, } if error is not None: task.error = error task.updated_at = now if status in ( JobStatus.COMPLETED.value, JobStatus.FAILED.value, JobStatus.CANCELLED.value, ): task.finished_at = now def patch_job( self, job_id: str, *, status: Optional[str] = None, result: Optional[Dict[str, Any]] = None, error: Optional[str] = None, ) -> Optional[Job]: with session_scope() as session: row = session.get(DownloadJobRow, job_id) if not row: return None new_status = status or row.status if result is not None: row.result = result if error is not None: row.error = error row.status = new_status finished = False if new_status in ( JobStatus.COMPLETED.value, JobStatus.FAILED.value, JobStatus.CANCELLED.value, ): row.finished_at = row.finished_at or utcnow() finished = True self._sync_task( session, row, status=new_status if status else None, result=result, error=error, ) if finished: worker = session.get(DownloadWorkerRow, row.worker_id) if worker and worker.current_job_id == job_id: worker.current_job_id = None worker.status = WorkerStatus.ONLINE.value session.flush() return self._row_job(row) def add_event( self, job_id: str, event_type: str, message: str, *, radicado: Optional[str] = None, nit: Optional[str] = None, index: Optional[int] = None, total: Optional[int] = None, metadata: Optional[Dict[str, Any]] = None, timestamp: Optional[str] = None, ) -> JobEventRecord: with session_scope() as session: ts = _parse_iso(timestamp) or utcnow() ev = DownloadJobEventRow( download_job_id=job_id, event_type=event_type, message=message, timestamp=ts, radicado=radicado, nit=nit, idx=index, total=total, metadata_json=metadata or {}, ) session.add(ev) session.flush() return JobEventRecord( id=int(ev.id), job_id=job_id, event_type=event_type, message=message, timestamp=_iso(ts) or utc_now(), radicado=radicado, nit=nit, index=index, total=total, metadata=metadata or {}, ) def list_events(self, job_id: str, since_id: int = 0) -> List[JobEventRecord]: with session_scope() as session: rows = session.scalars( select(DownloadJobEventRow) .where( DownloadJobEventRow.download_job_id == job_id, DownloadJobEventRow.id > since_id, ) .order_by(DownloadJobEventRow.id.asc()) ).all() return [ JobEventRecord( id=int(r.id), job_id=r.download_job_id, event_type=r.event_type, message=r.message, timestamp=_iso(r.timestamp) or utc_now(), radicado=r.radicado, nit=r.nit, index=r.idx, total=r.total, metadata=dict(r.metadata_json or {}), ) for r in rows ] def create_downloads_task( self, *, parent_job_id: str, download_job_id: str, worker_id: str, title: str, original_filename: str, ) -> str: """Crea Task(type=downloads) y la asocia al download_job. Devuelve task_id.""" with session_scope() as session: now = utcnow() task = Task( job_id=parent_job_id, type="downloads", status=JobStatus.PENDING.value, title=title or f"Bot descarga ({original_filename})", request={ "worker_id": worker_id, "original_filename": original_filename, "download_job_id": download_job_id, }, result={"download_job_id": download_job_id}, created_at=now, updated_at=now, ) session.add(task) session.flush() row = session.get(DownloadJobRow, download_job_id) if row: row.task_id = task.id row.job_id = parent_job_id session.flush() return task.id def parent_job_exists(self, parent_job_id: str) -> bool: from infrastructure.database import Job as IpsJob with session_scope() as session: return session.get(IpsJob, parent_job_id) is not None