feat: integration of job status update flow (#4)

Author: JanssenBrm

alembic/versions/dbc9c2bd9a57_update_of_processing_jobs.py

@@ -0,0 +1,48 @@
+"""Update of processing jobs
+
+Revision ID: dbc9c2bd9a57
+Revises: b3ba8db2adef
+Create Date: 2025-08-13 18:00:40.408757
+
+"""
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision: str = 'dbc9c2bd9a57'
+down_revision: Union[str, Sequence[str], None] = 'b3ba8db2adef'
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    """Upgrade schema."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.alter_column('processing_jobs', 'platform_job_id',
+               existing_type=sa.VARCHAR(),
+               nullable=False)
+    op.alter_column('processing_jobs', 'parameters',
+               existing_type=sa.VARCHAR(),
+               nullable=False)
+    op.alter_column('processing_jobs', 'service_record',
+               existing_type=sa.VARCHAR(),
+               nullable=False)
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    """Downgrade schema."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.alter_column('processing_jobs', 'service_record',
+               existing_type=sa.VARCHAR(),
+               nullable=True)
+    op.alter_column('processing_jobs', 'parameters',
+               existing_type=sa.VARCHAR(),
+               nullable=True)
+    op.alter_column('processing_jobs', 'platform_job_id',
+               existing_type=sa.VARCHAR(),
+               nullable=True)
+    # ### end Alembic commands ###

app/database/models/processing_job.py

@@ -23,10 +23,10 @@ class ProcessingJobRecord(Base):
         Enum(ProcessingStatusEnum), index=True
     )
     user_id: Mapped[str] = mapped_column(String, index=True)
-    platform_job_id: Mapped[Optional[str]] = mapped_column(String, index=True)
-    parameters: Mapped[Optional[str]] = mapped_column(String, index=False)
+    platform_job_id: Mapped[str] = mapped_column(String, index=True)
+    parameters: Mapped[str] = mapped_column(String, index=False)
     result_link: Mapped[Optional[str]] = mapped_column(String, index=False)
-    service_record: Mapped[Optional[str]] = mapped_column(String, index=True)
+    service_record: Mapped[str] = mapped_column(String, index=True)
     created: Mapped[datetime.datetime] = mapped_column(
         DateTime, default=datetime.datetime.utcnow, index=True
     )
@@ -63,6 +63,15 @@ def get_jobs_by_user_id(database: Session, user_id: str) -> List[ProcessingJobRe
     )
 
 
+def get_job_by_id(database: Session, job_id: int) -> Optional[ProcessingJobRecord]:
+    logger.info(f"Retrieving processing job with ID {job_id}")
+    return (
+        database.query(ProcessingJobRecord)
+        .filter(ProcessingJobRecord.id == job_id)
+        .first()
+    )
+
+
 def get_job_by_user_id(
     database: Session, job_id: int, user_id: str
 ) -> Optional[ProcessingJobRecord]:
@@ -74,3 +83,19 @@ def get_job_by_user_id(
         )
         .first()
     )
+
+
+def update_job_status_by_id(
+    database: Session, job_id: int, status: ProcessingStatusEnum
+):
+    logger.info(f"Updating the status of processing job with ID {job_id} to {status}")
+    job = get_job_by_id(database, job_id)
+
+    if job:
+        job.status = status
+        database.commit()
+        database.refresh(job)
+    else:
+        logger.warning(
+            f"Could not update job status of job {job_id} as it could not be found in the database"
+        )

app/platforms/base.py

@@ -1,6 +1,6 @@
 from abc import ABC, abstractmethod
 
-from app.schemas import ServiceDetails
+from app.schemas import ProcessingStatusEnum, ServiceDetails
 
 
 class BaseProcessingPlatform(ABC):
@@ -20,3 +20,16 @@ def execute_job(self, title: str, details: ServiceDetails, parameters: dict) ->
         :return: Return the ID of the job that was created
         """
         pass
+
+    @abstractmethod
+    def get_job_status(
+        self, job_id: str, details: ServiceDetails
+    ) -> ProcessingStatusEnum:
+        """
+        Retrieve the job status of a processing job that is running on the platform.
+
+        :param job_id: The ID of the job on the platform
+        :param details: The service details containing the service ID and application.
+        :return: Return the processing status
+        """
+        pass

app/platforms/implementations/ogc_api_process.py

@@ -2,7 +2,7 @@
 
 from app.platforms.base import BaseProcessingPlatform
 from app.platforms.dispatcher import register_processing_platform
-from app.schemas import ProcessTypeEnum, ServiceDetails
+from app.schemas import ProcessTypeEnum, ProcessingStatusEnum, ServiceDetails
 
 logger = logging.getLogger(__name__)
 
@@ -14,16 +14,14 @@ class OGCAPIProcessPlatform(BaseProcessingPlatform):
     """
 
     def execute_job(self, title: str, details: ServiceDetails, parameters: dict) -> str:
-        """
-        Execute a processing job on the platform with the given service ID and parameters.
-
-        :param title: The title of the job to be executed.
-        :param details: The service details containing the service ID and application.
-        :param parameters: The parameters required for the job execution.
-        :return: Return the ID of the job that was created
-        """
-
         raise NotImplementedError("OGC API Process job execution not implemented yet.")
 
+    def get_job_status(
+        self, job_id: str, details: ServiceDetails
+    ) -> ProcessingStatusEnum:
+        raise NotImplementedError(
+            "OGC API Process job status retrieval not implemented yet."
+        )
+
 
 register_processing_platform(ProcessTypeEnum.OGC_API_PROCESS, OGCAPIProcessPlatform)

app/platforms/implementations/openeo.py

@@ -11,6 +11,7 @@
 from app.platforms.dispatcher import register_processing_platform
 from app.schemas import (
     ProcessTypeEnum,
+    ProcessingStatusEnum,
     ServiceDetails,
 )
 
@@ -104,15 +105,6 @@ def _get_process_id(self, url: str) -> str:
         return process_id
 
     def execute_job(self, title: str, details: ServiceDetails, parameters: dict) -> str:
-        """
-        Execute a processing job on the platform with the given service ID and parameters.
-
-        :param title: The title of the job to be executed.
-        :param details: The service details containing the service ID and application.
-        :param parameters: The parameters required for the job execution.
-        :return: Return the ID of the job that was created
-        """
-
         try:
             process_id = self._get_process_id(details.application)
 
@@ -130,8 +122,47 @@ def execute_job(self, title: str, details: ServiceDetails, parameters: dict) ->
 
             return job.job_id
         except Exception as e:
-            logger.exception(f"Failed to execute openEO job: {e}")
-            raise SystemError("Failed to execute openEO job")
+            logger.exception("Failed to execute openEO job")
+            raise SystemError("Failed to execute openEO job") from e
+
+    def _map_openeo_status(self, status: str) -> ProcessingStatusEnum:
+        """
+        Map the status returned by openEO to a status known within the API.
+
+        :param status: Status text returned by openEO.
+        :return: ProcessingStatusEnum corresponding to the input.
+        """
+
+        logger.debug("Mapping openEO status %r to ProcessingStatusEnum", status)
+
+        mapping = {
+            "created": ProcessingStatusEnum.CREATED,
+            "queued": ProcessingStatusEnum.QUEUED,
+            "running": ProcessingStatusEnum.RUNNING,
+            "cancelled": ProcessingStatusEnum.CANCELED,
+            "finished": ProcessingStatusEnum.FINISHED,
+            "error": ProcessingStatusEnum.FAILED,
+        }
+
+        try:
+            return mapping[status.lower()]
+        except (AttributeError, KeyError):
+            logger.warning("Mapping of unknown openEO status: %r", status)
+            return ProcessingStatusEnum.UNKNOWN
+
+    def get_job_status(
+        self, job_id: str, details: ServiceDetails
+    ) -> ProcessingStatusEnum:
+        try:
+            logger.debug(f"Fetching job status for openEO job with ID {job_id}")
+            connection = self._setup_connection(details.service)
+            job = connection.job(job_id)
+            return self._map_openeo_status(job.status())
+        except Exception as e:
+            logger.exception(f"Failed to fetch status for openEO job with ID {job_id}")
+            raise SystemError(
+                f"Failed to fetch status openEO job with ID {job_id}"
+            ) from e
 
 
 register_processing_platform(ProcessTypeEnum.OPENEO, OpenEOPlatform)

app/routers/jobs_status.py

@@ -1,6 +1,7 @@
+import asyncio
 import logging
 
-from fastapi import APIRouter, Depends
+from fastapi import APIRouter, Depends, WebSocket, WebSocketDisconnect
 from sqlalchemy.orm import Session
 
 from app.database.db import get_db
@@ -16,7 +17,7 @@
     tags=["Upscale Tasks", "Unit Jobs"],
     summary="Get a list of all upscaling tasks & processing jobs for the authenticated user",
 )
-async def jobs_status(
+async def get_jobs_status(
     db: Session = Depends(get_db),
     user: str = "foobar",
 ) -> JobsStatusResponse:
@@ -29,3 +30,33 @@ async def jobs_status(
     return JobsStatusResponse(
         upscaling_tasks=[], processing_jobs=get_processing_jobs_by_user_id(db, user)
     )
+
+
+@router.websocket(
+    "/ws/jobs_status",
+)
+async def ws_jobs_status(
+    websocket: WebSocket, user: str = "foobar", interval: int = 10
+):
+    """
+    Return combined list of upscaling tasks and processing jobs for the authenticated user.
+    """
+
+    await websocket.accept()
+    logger.debug(f"WebSocket connected for user {user}")
+
+    db = next(get_db())
+    try:
+        while True:
+            status = await get_jobs_status(db, user)
+            await websocket.send_json(status.model_dump())
+
+            await asyncio.sleep(interval)
+
+    except WebSocketDisconnect:
+        logger.info(f"WebSocket disconnected for user {user}")
+    except Exception as e:
+        logger.exception(f"Error in jobs_status_ws: {e}")
+        await websocket.close(code=1011, reason="Error in job status websocket: {e}")
+    finally:
+        db.close()

app/schemas.py

@@ -11,9 +11,12 @@ class ProcessTypeEnum(str, Enum):
 
 class ProcessingStatusEnum(str, Enum):
     CREATED = "created"
+    QUEUED = "queued"
     RUNNING = "running"
     FINISHED = "finished"
+    CANCELED = "canceled"
     FAILED = "failed"
+    UNKNOWN = "unknown"
 
 
 # class TileRequest(BaseModel):

app/services/processing.py

@@ -6,6 +6,7 @@
     get_job_by_user_id,
     get_jobs_by_user_id,
     save_job_to_db,
+    update_job_status_by_id,
 )
 from app.platforms.dispatcher import get_processing_platform
 from app.schemas import (
@@ -40,28 +41,50 @@ def create_processing_job(
         platform_job_id=job_id,
         parameters=json.dumps(summary.parameters),
         result_link=None,
-        service_record=json.dumps(
-            summary.service.model_dump_json()
-        ),  # Assuming service is a dict
+        service_record=summary.service.model_dump_json(),
     )
     record = save_job_to_db(database, record)
     return ProcessingJobSummary(
         id=record.id, title=record.title, label=summary.label, status=record.status
     )
 
 
+def get_job_status(job: ProcessingJobRecord) -> ProcessingStatusEnum:
+    logger.info(f"Retrieving job status for job: {job.platform_job_id}")
+    platform = get_processing_platform(job.label)
+    details = ServiceDetails.model_validate_json(job.service_record)
+    return platform.get_job_status(job.platform_job_id, details)
+
+
 def get_processing_jobs_by_user_id(
     database: Session, user_id: str
 ) -> List[ProcessingJobSummary]:
     logger.info(f"Retrieving processing jobs for user {user_id}")
-    return list(
-        map(
-            lambda x: ProcessingJobSummary(
-                id=x.id, title=x.title, label=x.label, status=x.status
-            ),
-            get_jobs_by_user_id(database, user_id),
+
+    jobs: List[ProcessingJobSummary] = []
+    records = get_jobs_by_user_id(database, user_id)
+
+    inactive_statuses = {
+        ProcessingStatusEnum.CANCELED,
+        ProcessingStatusEnum.FAILED,
+        ProcessingStatusEnum.FINISHED,
+    }
+
+    for record in records:
+        if record.status not in inactive_statuses:
+            latest_status = get_job_status(record)
+            if latest_status != record.status:
+                update_job_status_by_id(database, record.id, latest_status)
+            status = latest_status
+        else:
+            status = record.status
+
+        jobs.append(
+            ProcessingJobSummary(
+                id=record.id, title=record.title, label=record.label, status=status
+            )
         )
-    )
+    return jobs
 
 
 def get_processing_job_by_user_id(

requirements.txt

@@ -10,6 +10,7 @@ pydantic
 pydantic-settings
 pydantic_core
 pytest
+pytest-asyncio
 pytest-cov
 python-dotenv
 requests