[DOP-29112] add status and job_type filters for run search

Author: TiGrib

data_rentgen/db/repositories/run.py

@@ -94,6 +94,8 @@ async def paginate(
         job_id: int | None,
         parent_run_id: UUID | None,
         search_query: str | None,
+        job_types: Collection[str],
+        statuses: Collection[str],
     ) -> PaginationDTO[Run]:
         # do not use `tuple_(Run.created_at, Run.id).in_(...),
         # as this is too complex filter for Postgres to make an optimal query plan
@@ -130,6 +132,9 @@ async def paginate(
             where.append(Run.job_id == job_id)
         if parent_run_id:
             where.append(Run.parent_run_id == parent_run_id)
+        if statuses:
+            serialize_statuses: Collection[RunStatus] = [RunStatus[status] for status in statuses]
+            where.append(Run.status == any_(serialize_statuses))  # type: ignore[arg-type]
 
         query: Select | CompoundSelect
         order_by: list[ColumnElement | SQLColumnExpression]
@@ -160,6 +165,9 @@ async def paginate(
             query = select(Run).where(*where)
             order_by = [Run.created_at.desc(), Run.id.desc()]
 
+        if job_types:
+            query = query.join(Job, Run.job_id == Job.id).where(Job.type == any_(job_types))  # type: ignore[arg-type]
+
         options = [selectinload(Run.started_by_user)]
         return await self._paginate_by_query(
             query=query,

data_rentgen/server/api/v1/router/run.py

@@ -39,6 +39,8 @@ async def runs(
         job_id=query_args.job_id,
         parent_run_id=query_args.parent_run_id,
         search_query=query_args.search_query,
+        job_types=query_args.job_types,
+        statuses=query_args.statuses,
     )
     return PageResponseV1[RunDetailedResponseV1].from_pagination(pagination)
 

data_rentgen/server/schemas/v1/run.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 from datetime import datetime
-from enum import IntEnum
+from enum import IntEnum, StrEnum
 from uuid import UUID
 
 from pydantic import (
@@ -41,6 +41,14 @@ def __str__(self) -> str:
         return self.name
 
 
+class RunStatusForQueryV1(StrEnum):
+    UNKNOWN = "UNKNOWN"
+    STARTED = "STARTED"
+    SUCCEEDED = "SUCCEEDED"
+    FAILED = "FAILED"
+    KILLED = "KILLED"
+
+
 class RunResponseV1(BaseModel):
     """Run response"""
 
@@ -139,6 +147,13 @@ class RunsQueryV1(PaginateQueryV1):
         examples=["01913217-b761-7b1a-bb52-489da9c8b9c8"],
     )
 
+    job_types: list[str] = Field(
+        default_factory=list,
+        description="Filter runs by type of a Job",
+        examples=["SPARK_APPLICATION", "AIRFLOW_TASK"],
+    )
+    statuses: list[RunStatusForQueryV1] = Field(default_factory=list, description="Filter runs by status")
+
     search_query: str | None = Field(
         default=None,
         min_length=3,

data_rentgen/server/services/run.py

@@ -80,6 +80,8 @@ async def paginate(
         job_id: int | None,
         parent_run_id: UUID | None,
         search_query: str | None,
+        job_types: Collection[str],
+        statuses: Collection[str],
     ) -> RunServicePaginatedResult:
         pagination = await self._uow.run.paginate(
             page=page,
@@ -90,6 +92,8 @@ async def paginate(
             job_id=job_id,
             parent_run_id=parent_run_id,
             search_query=search_query,
+            job_types=job_types,
+            statuses=statuses,
         )
         run_ids = [item.id for item in pagination.items]
         input_stats = await self._uow.input.get_stats_by_run_ids(run_ids)

tests/test_server/fixtures/factories/run.py

@@ -135,7 +135,8 @@ async def runs_with_same_job(
                 async_session,
                 run_kwargs={
                     "job_id": job.id,
-                    "created_at": started_at + timedelta(seconds=s),  # To be sure runs has different timestamps
+                    # To be sure runs has different timestamps
+                    "created_at": started_at + timedelta(seconds=s),
                     "started_by_user_id": user.id,
                     **params,
                 },
@@ -167,7 +168,8 @@ async def runs_with_same_parent(
                 async_session,
                 run_kwargs={
                     "parent_run_id": parent_run_id,
-                    "created_at": started_at + timedelta(seconds=s),  # To be sure runs has different timestamps
+                    # To be sure runs has different timestamps
+                    "created_at": started_at + timedelta(seconds=s),
                     "started_by_user_id": user.id,
                     **params,
                 },
@@ -193,10 +195,10 @@ async def runs_search(
         {"name": "airflow_dag_name", "type": "AIRFLOW_DAG"},
     ]
     runs_kwargs = [
-        {"external_id": "application_1638922609021_0001"},
-        {"external_id": "application_1638922609021_0002"},
-        {"external_id": "extract_task_0001"},
-        {"external_id": "extract_task_0002"},
+        {"external_id": "application_1638922609021_0001", "status": 3},
+        {"external_id": "application_1638922609021_0002", "status": 1},
+        {"external_id": "extract_task_0001", "status": 0},
+        {"external_id": "extract_task_0002", "status": 2},
     ]
     started_at = datetime.now(tz=UTC)
     async with async_session_maker() as async_session:

tests/test_server/test_runs/test_get_runs.py

@@ -36,6 +36,8 @@ async def test_get_runs_missing_fields(
                     "message": "Value error, input should contain either 'since', 'run_id', 'job_id', 'parent_run_id' or 'search_query' field",
                     "context": {},
                     "input": {
+                        "statuses": [],
+                        "job_types": [],
                         "page": 1,
                         "page_size": 20,
                         "run_id": [],

tests/test_server/test_runs/test_get_runs_by_job_id.py

@@ -38,6 +38,8 @@ async def test_get_runs_by_job_id_missing_since(
                     "message": "Value error, 'job_id' can be passed only with 'since'",
                     "context": {},
                     "input": {
+                        "statuses": [],
+                        "job_types": [],
                         "page": 1,
                         "page_size": 20,
                         "job_id": str(new_run.job_id),

tests/test_server/test_runs/test_get_runs_by_parent_run_id.py

@@ -38,6 +38,8 @@ async def test_get_runs_by_job_id_missing_since(
                     "message": "Value error, 'parent_run_id' can be passed only with 'since'",
                     "context": {},
                     "input": {
+                        "statuses": [],
+                        "job_types": [],
                         "page": 1,
                         "page_size": 20,
                         "parent_run_id": str(new_run.parent_run_id),

tests/test_server/test_runs/test_search_runs.py

@@ -37,6 +37,8 @@ async def test_search_runs_missing_since(
                     "message": "Value error, 'search_query' can be passed only with 'since'",
                     "context": {},
                     "input": {
+                        "statuses": [],
+                        "job_types": [],
                         "page": 1,
                         "page_size": 20,
                         "run_id": [],
@@ -198,7 +200,71 @@ async def test_search_runs_by_job_type(
         headers={"Authorization": f"Bearer {mocked_user.access_token}"},
         params={
             "since": since.isoformat(),
-            "search_query": "SPARK",
+            "job_types": ["SPARK_APPLICATION"],
+        },
+    )
+
+    assert response.status_code == HTTPStatus.OK, response.json()
+    assert response.json() == {
+        "meta": {
+            "has_next": False,
+            "has_previous": False,
+            "next_page": None,
+            "page": 1,
+            "page_size": 20,
+            "pages_count": 1,
+            "previous_page": None,
+            "total_count": 2,
+        },
+        "items": [
+            {
+                "id": str(run.id),
+                "data": run_to_json(run),
+                "statistics": {
+                    "inputs": {
+                        "total_datasets": 0,
+                        "total_bytes": 0,
+                        "total_rows": 0,
+                        "total_files": 0,
+                    },
+                    "outputs": {
+                        "total_datasets": 0,
+                        "total_bytes": 0,
+                        "total_rows": 0,
+                        "total_files": 0,
+                    },
+                    "operations": {
+                        "total_operations": 0,
+                    },
+                },
+            }
+            for run in runs
+        ],
+    }
+
+
+async def test_search_runs_by_status(
+    test_client: AsyncClient,
+    async_session: AsyncSession,
+    runs_search: dict[str, Run],
+    mocked_user: MockedUser,
+) -> None:
+    runs = await enrich_runs(
+        [
+            # runs sorted by id in descending order
+            runs_search["extract_task_0001"],
+            runs_search["application_1638922609021_0002"],
+        ],
+        async_session,
+    )
+    since = min(run.created_at for run in runs)
+
+    response = await test_client.get(
+        "/v1/runs",
+        headers={"Authorization": f"Bearer {mocked_user.access_token}"},
+        params={
+            "since": since.isoformat(),
+            "statuses": ["SUCCEEDED", "STARTED"],
         },
     )