core:services:kraken: Allow consuming jobs results as stream

core/services/kraken/api/v2/routers/jobs.py

@@ -3,9 +3,10 @@
 from typing import Any, Callable, List, Tuple
 
 from fastapi import APIRouter, Body, HTTPException, status
+from fastapi.responses import StreamingResponse
 from fastapi_versioning import versioned_api_route
 from jobs import JobsManager
-from jobs.exceptions import JobNotFound
+from jobs.exceptions import JobIsRunning, JobNotFound
 from jobs.models import Job, JobMethod
 
 jobs_router_v2 = APIRouter(
@@ -21,6 +22,8 @@ def jobs_to_http_exception(endpoint: Callable[..., Any]) -> Callable[..., Any]:
     async def wrapper(*args: Tuple[Any], **kwargs: dict[str, Any]) -> Any:
         try:
             return await endpoint(*args, **kwargs)
+        except JobIsRunning as error:
+            raise HTTPException(status_code=status.HTTP_423_LOCKED, detail=str(error)) from error
         except JobNotFound as error:
             raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=str(error)) from error
         except Exception as error:
@@ -45,6 +48,12 @@ async def fetch() -> List[Job]:
     return JobsManager.get()
 
 
+@jobs_router_v2.get("/{identifier}/stream", status_code=status.HTTP_200_OK)
+@jobs_to_http_exception
+async def stream(identifier: str) -> StreamingResponse:
+    return StreamingResponse(JobsManager.stream(identifier), media_type="text/plain")
+
+
 @jobs_router_v2.get("/{identifier}", status_code=status.HTTP_200_OK)
 @jobs_to_http_exception
 async def fetch_by_identifier(identifier: str) -> Job:

core/services/kraken/jobs/exceptions.py

@@ -1,2 +1,6 @@
 class JobNotFound(Exception):
     pass
+
+
+class JobIsRunning(Exception):
+    pass

core/services/kraken/jobs/jobs.py

@@ -1,15 +1,17 @@
 import asyncio
-from typing import List, Optional
+from typing import AsyncGenerator, Dict, List, Optional
 
 import aiohttp
-from jobs.exceptions import JobNotFound
-from jobs.models import Job
+from jobs.exceptions import JobIsRunning, JobNotFound
+from jobs.models import Job, JobStatus
+from jobs.stream import JobStream
 from loguru import logger
 
 
 class JobsManager:
     _jobs: List[Job] = []
     _executing_job: Optional[Job] = None
+    _streams: Dict[str, JobStream] = {}
 
     def __init__(self) -> None:
         self.is_running = True
@@ -18,27 +20,38 @@ def __init__(self) -> None:
     async def execute_job(self, job: Job) -> None:
         job_name = f"{job.method.value} - {job.route}"
         logger.info(f"Executing job {job_name}")
+        stream = self._streams[job.id]
         async with aiohttp.ClientSession() as session:
             for i in range(job.retries):
                 try:
                     async with session.request(
                         method=job.method, url=f"{self.base_host}/{job.route}", json=job.body
                     ) as response:
                         response.raise_for_status()
-                        await response.read()
+                        async for chunk in response.content.iter_any():
+                            stream.publish(chunk)
+                        job.status = JobStatus.SUCCESS
                         return
                 except Exception:
                     logger.warning(f"Failed job {job_name} attempt {i + 1}/{job.retries}")
+                    stream.reset()
                     await asyncio.sleep(5)
+            job.status = JobStatus.ERROR
             logger.error(f"Job {job_name} failed to be executed")
 
     async def start(self) -> None:
         while self.is_running:
             await asyncio.sleep(1)
-            if self._jobs:
-                self._executing_job = self._jobs.pop(0)
-                await self.execute_job(self._executing_job)
-                self._executing_job = None
+            if JobsManager._jobs:
+                JobsManager._executing_job = JobsManager._jobs.pop(0)
+                JobsManager._executing_job.status = JobStatus.RUNNING
+                stream = JobsManager._streams[JobsManager._executing_job.id]
+                try:
+                    await self.execute_job(JobsManager._executing_job)
+                finally:
+                    stream.close()
+                    JobsManager._streams.pop(JobsManager._executing_job.id, None)
+                    JobsManager._executing_job = None
 
     async def stop(self) -> None:
         self.is_running = False
@@ -48,6 +61,7 @@ def set_base_host(self, host: str) -> None:
 
     @classmethod
     def add(cls, job: Job) -> None:
+        cls._streams[job.id] = JobStream()
         cls._jobs.append(job)
 
     @classmethod
@@ -61,7 +75,19 @@ def get_by_identifier(cls, identifier: str) -> Job:
             raise JobNotFound(f"Job with id {identifier} not found")
         return job
 
+    @classmethod
+    def stream(cls, identifier: str) -> AsyncGenerator[bytes, None]:
+        cls.get_by_identifier(identifier)
+        return cls._streams[identifier].subscribe()
+
     @classmethod
     def delete(cls, identifier: str) -> None:
-        job = cls.get_by_identifier(identifier)
-        cls._jobs.remove(job)
+        cls.get_by_identifier(identifier)
+        job = next((job for job in cls._jobs if job.id == identifier), None)
+        if job:
+            cls._jobs.remove(job)
+            stream = cls._streams.pop(identifier, None)
+            if stream:
+                stream.close()
+        elif cls._executing_job and cls._executing_job.id == identifier:
+            raise JobIsRunning(f"Cannot delete job {identifier} while it is running")

core/services/kraken/jobs/models.py

@@ -13,9 +13,17 @@ class JobMethod(str, Enum):
     DELETE = "DELETE"
 
 
+class JobStatus(str, Enum):
+    PENDING = "PENDING"
+    RUNNING = "RUNNING"
+    SUCCESS = "SUCCESS"
+    ERROR = "ERROR"
+
+
 class Job(BaseModel):
     id: str
     route: str
     method: JobMethod
     body: Any
     retries: int = 5
+    status: JobStatus = JobStatus.PENDING

core/services/kraken/jobs/stream.py

@@ -0,0 +1,38 @@
+import asyncio
+from typing import AsyncGenerator, List
+
+
+class JobStream:
+    def __init__(self) -> None:
+        self._chunks: List[bytes] = []
+        self._updated = asyncio.Event()
+        self._done = False
+        self._retries = 0
+
+    def publish(self, chunk: bytes) -> None:
+        self._chunks.append(chunk)
+        self._updated.set()
+
+    def reset(self) -> None:
+        self._chunks = []
+        self._retries += 1
+        self._updated.set()
+
+    def close(self) -> None:
+        self._done = True
+        self._updated.set()
+
+    async def subscribe(self) -> AsyncGenerator[bytes, None]:
+        index = 0
+        retries = self._retries
+        while True:
+            self._updated.clear()
+            if retries != self._retries:
+                retries = self._retries
+                index = 0
+            while index < len(self._chunks):
+                yield self._chunks[index]
+                index += 1
+            if self._done:
+                return
+            await self._updated.wait()