feat: add type filtering on job status endpoints

Author: JanssenBrm

app/routers/jobs_status.py

@@ -1,19 +1,22 @@
 import asyncio
 import json
+from typing import List
 
-from fastapi import APIRouter, Depends, WebSocket, WebSocketDisconnect
+from fastapi import APIRouter, Depends, Query, WebSocket, WebSocketDisconnect
 from sqlalchemy.orm import Session
 from loguru import logger
 
 from app.database.db import SessionLocal, get_db
-from app.schemas.jobs_status import JobsStatusResponse
+from app.schemas.jobs_status import JobsFilter, JobsStatusResponse
 from app.schemas.websockets import WSStatusMessage
 from app.services.processing import get_processing_jobs_by_user_id
 from app.services.upscaling import get_upscaling_tasks_by_user_id
 from app.auth import get_current_user_id, websocket_authenticate
 
 router = APIRouter()
 
+DEFAULT_FILTERS = [JobsFilter.upscaling, JobsFilter.processing]
+
 
 @router.get(
     "/jobs_status",
@@ -23,14 +26,28 @@
 async def get_jobs_status(
     db: Session = Depends(get_db),
     user: str = Depends(get_current_user_id),
+    filter: List[JobsFilter] = Query(
+        DEFAULT_FILTERS,
+        description="Filter jobs: upscaling, processing. Can be provided multiple times.",
+    ),
 ) -> JobsStatusResponse:
     """
     Return combined list of upscaling tasks and processing jobs for the authenticated user.
     """
     logger.debug(f"Fetching jobs list for user {user}")
+    upscaling_tasks = (
+        get_upscaling_tasks_by_user_id(db, user)
+        if JobsFilter.upscaling in filter
+        else []
+    )
+    processing_jobs = (
+        get_processing_jobs_by_user_id(db, user)
+        if JobsFilter.processing in filter
+        else []
+    )
     return JobsStatusResponse(
-        upscaling_tasks=get_upscaling_tasks_by_user_id(db, user),
-        processing_jobs=get_processing_jobs_by_user_id(db, user),
+        upscaling_tasks=upscaling_tasks,
+        processing_jobs=processing_jobs,
     )
 
 
@@ -40,6 +57,7 @@ async def get_jobs_status(
 async def ws_jobs_status(
     websocket: WebSocket,
     interval: int = 10,
+    filter: List[JobsFilter] = Query(DEFAULT_FILTERS),
 ):
     """
     Return combined list of upscaling tasks and processing jobs for the authenticated user.
@@ -62,7 +80,7 @@ async def ws_jobs_status(
                         message="Starting retrieval of status",
                     ).model_dump()
                 )
-                status = await get_jobs_status(db, user)
+                status = await get_jobs_status(db, user, filter=filter)
                 await websocket.send_json(
                     WSStatusMessage(
                         type="status",

app/schemas/jobs_status.py

@@ -1,3 +1,4 @@
+from enum import Enum
 from typing import List
 from pydantic import BaseModel, Field
 
@@ -12,3 +13,8 @@ class JobsStatusResponse(BaseModel):
     processing_jobs: List[ProcessingJobSummary] = Field(
         ..., description="List of processing jobs that are available for the user"
     )
+
+
+class JobsFilter(str, Enum):
+    upscaling = "upscaling"
+    processing = "processing"

openapi.json

@@ -156,6 +156,31 @@
         "summary": "Get a list of all upscaling tasks & processing jobs for the authenticated user",
         "description": "Return combined list of upscaling tasks and processing jobs for the authenticated user.",
         "operationId": "get_jobs_status_jobs_status_get",
+        "security": [
+          {
+            "OAuth2AuthorizationCodeBearer": []
+          }
+        ],
+        "parameters": [
+          {
+            "name": "filter",
+            "in": "query",
+            "required": false,
+            "schema": {
+              "type": "array",
+              "items": {
+                "$ref": "#/components/schemas/JobsFilter"
+              },
+              "description": "Filter jobs: upscaling, processing. Can be provided multiple times.",
+              "default": [
+                "upscaling",
+                "processing"
+              ],
+              "title": "Filter"
+            },
+            "description": "Filter jobs: upscaling, processing. Can be provided multiple times."
+          }
+        ],
         "responses": {
           "200": {
             "description": "Successful Response",
@@ -166,13 +191,18 @@
                 }
               }
             }
+          },
+          "422": {
+            "description": "Validation Error",
+            "content": {
+              "application/json": {
+                "schema": {
+                  "$ref": "#/components/schemas/HTTPValidationError"
+                }
+              }
+            }
           }
-        },
-        "security": [
-          {
-            "OAuth2AuthorizationCodeBearer": []
-          }
-        ]
+        }
       }
     },
     "/unit_jobs": {
@@ -972,6 +1002,14 @@
         "type": "object",
         "title": "HTTPValidationError"
       },
+      "JobsFilter": {
+        "type": "string",
+        "enum": [
+          "upscaling",
+          "processing"
+        ],
+        "title": "JobsFilter"
+      },
       "JobsStatusResponse": {
         "properties": {
           "upscaling_tasks": {

tests/routers/test_job_status.py

@@ -28,6 +28,48 @@ def test_unit_jobs_get_200(
     ).model_dump_json(indent=1)
 
 
+@patch("app.routers.jobs_status.get_processing_jobs_by_user_id")
+@patch("app.routers.jobs_status.get_upscaling_tasks_by_user_id")
+def test_unit_jobs_get_only_processing_200(
+    mock_get_upscaling_tasks,
+    mock_get_processing_jobs,
+    client,
+    fake_processing_job_summary,
+    fake_upscaling_task_summary,
+):
+
+    mock_get_processing_jobs.return_value = [fake_processing_job_summary]
+    mock_get_upscaling_tasks.return_value = [fake_upscaling_task_summary]
+
+    r = client.get("/jobs_status?filter=processing")
+    assert r.status_code == 200
+    assert json.dumps(r.json(), indent=1) == JobsStatusResponse(
+        upscaling_tasks=[],
+        processing_jobs=[fake_processing_job_summary],
+    ).model_dump_json(indent=1)
+
+
+@patch("app.routers.jobs_status.get_processing_jobs_by_user_id")
+@patch("app.routers.jobs_status.get_upscaling_tasks_by_user_id")
+def test_unit_jobs_get_only_upscaling_200(
+    mock_get_upscaling_tasks,
+    mock_get_processing_jobs,
+    client,
+    fake_processing_job_summary,
+    fake_upscaling_task_summary,
+):
+
+    mock_get_processing_jobs.return_value = [fake_processing_job_summary]
+    mock_get_upscaling_tasks.return_value = [fake_upscaling_task_summary]
+
+    r = client.get("/jobs_status?filter=upscaling")
+    assert r.status_code == 200
+    assert json.dumps(r.json(), indent=1) == JobsStatusResponse(
+        upscaling_tasks=[fake_upscaling_task_summary],
+        processing_jobs=[],
+    ).model_dump_json(indent=1)
+
+
 @pytest.mark.asyncio
 @patch("app.auth.get_current_user_id", new_callable=AsyncMock)
 @patch("app.routers.jobs_status.get_jobs_status", new_callable=AsyncMock)