[DOP-22532] change job type filter to list

Author: TiGrib

data_rentgen/db/repositories/job.py

@@ -77,7 +77,7 @@ async def paginate(
         job_ids: Collection[int],
         search_query: str | None,
         location_id: int | None,
-        job_type: str | None,
+        job_type: Collection[str],
     ) -> PaginationDTO[Job]:
         where = []
         if job_ids:
@@ -86,7 +86,7 @@ async def paginate(
         query: Select | CompoundSelect
         order_by: list[ColumnElement | SQLColumnExpression]
         if job_type:
-            where.append(Job.type == job_type)  # type: ignore[arg-type]
+            where.append(Job.type == any_(list(job_type)))  # type: ignore[arg-type]
         if location_id:
             where.append(Job.location_id == location_id)  # type: ignore[arg-type]
 

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

@@ -22,7 +22,9 @@
 router = APIRouter(
     prefix="/jobs",
     tags=["Jobs"],
-    responses=get_error_responses(include={NotAuthorizedSchema, NotAuthorizedRedirectSchema, InvalidRequestSchema}),
+    responses=get_error_responses(
+        include={NotAuthorizedSchema, NotAuthorizedRedirectSchema, InvalidRequestSchema},
+    ),
 )
 
 
@@ -62,7 +64,7 @@ async def get_jobs_lineage(
     return build_lineage_response(lineage)
 
 
-@router.get("/job_types", summary="Get distinct types of Jobs")
+@router.get("/job-types", summary="Get distinct types of Jobs")
 async def get_job_types(
     job_service: Annotated[JobService, Depends()],
     current_user: Annotated[User, Depends(get_user())],

data_rentgen/server/schemas/v1/job.py

@@ -43,13 +43,13 @@ class JobPaginateQueryV1(PaginateQueryV1):
         min_length=3,
         description="Search query",
     )
-    job_type: str | None = Field(
-        default=None,
-        description="Filter for searching",
+    job_type: list[str] = Field(
+        default_factory=list,
+        description="Specify job types",
     )
     location_id: int | None = Field(
         default=None,
-        description="",
+        description="The location id which jobs belong",
     )
 
     model_config = ConfigDict(extra="forbid")

data_rentgen/server/services/job.py

@@ -32,7 +32,7 @@ async def paginate(
         job_ids: Collection[int],
         search_query: str | None,
         location_id: int | None,
-        job_type: str | None,
+        job_type: Collection[str],
     ) -> JobServicePaginatedResult:
         pagination = await self._uow.job.paginate(
             page=page,

tests/test_server/fixtures/factories/job.py

@@ -155,7 +155,10 @@ async def jobs_search(
             await create_location(
                 async_session,
                 location_kwargs=kwargs,
-                address_kwargs={"urls": [random_string(32), random_string(32)]},  # Each location has 2 addresses
+                # Each location has 2 addresses
+                address_kwargs={
+                    "urls": [random_string(32), random_string(32)],
+                },
             )
             for kwargs in location_kwargs
         ]
@@ -170,7 +173,9 @@ async def jobs_search(
             await create_location(
                 async_session,
                 location_kwargs={"type": "random"},
-                address_kwargs={"urls": [random_string(32), random_string(32)]},
+                address_kwargs={
+                    "urls": [random_string(32), random_string(32)],
+                },
             )
             for _ in range(3)
         ]
@@ -207,9 +212,19 @@ async def jobs_search(
 
     jobs_by_name = {job.name: job for job in jobs_with_names}
     jobs_by_location = dict(
-        zip([location.name for location in locations_with_names], jobs_with_location_names, strict=False),
+        zip(
+            [location.name for location in locations_with_names],
+            jobs_with_location_names,
+            strict=False,
+        ),
+    )
+    jobs_by_address = dict(
+        zip(
+            addresses_url,
+            [job for job in jobs_with_address_urls for _ in range(2)],
+            strict=False,
+        ),
     )
-    jobs_by_address = dict(zip(addresses_url, [job for job in jobs_with_address_urls for _ in range(2)], strict=False))
 
     yield jobs_by_name, jobs_by_location, jobs_by_address
 
@@ -218,36 +233,7 @@ async def jobs_search(
 
 
 @pytest_asyncio.fixture
-async def jobs_search_with_different_types(
-    async_session_maker: Callable[[], AbstractAsyncContextManager[AsyncSession]],
-) -> AsyncGenerator[tuple[Job]]:
-    async with async_session_maker() as async_session:
-        location = await create_location(async_session)
-        airflow_dag_type = await create_job_type(async_session, {"type": "AIRFLOW_DAG"})
-        airflow_task_type = await create_job_type(async_session, {"type": "AIRFLOW_TASK"})
-
-        job_with_dag_type = await create_job(
-            async_session,
-            location_id=location.id,
-            job_type_id=airflow_dag_type.id,
-            job_kwargs={"name": "airflow-dag"},
-        )
-        job_with_task_type = await create_job(
-            async_session,
-            location_id=location.id,
-            job_type_id=airflow_task_type.id,
-            job_kwargs={"name": "airflow-task"},
-        )
-        async_session.expunge_all()
-
-    yield (job_with_dag_type, job_with_task_type)
-
-    async with async_session_maker() as async_session:
-        await clean_db(async_session)
-
-
-@pytest_asyncio.fixture
-async def jobs_search_with_different_locations(
+async def jobs_with_locations_and_types(
     async_session_maker: Callable[[], AbstractAsyncContextManager[AsyncSession]],
 ) -> AsyncGenerator[tuple[Job]]:
     async with async_session_maker() as async_session:
@@ -256,23 +242,25 @@ async def jobs_search_with_different_locations(
             async_session,
             location_kwargs={"name": "airflow-host", "type": "http"},
         )
-        job_type = await create_job_type(async_session)
+        cluster_type = await create_job_type(async_session, {"type": "SPARK_APPLICATION"})
+        airflow_dag_type = await create_job_type(async_session, {"type": "AIRFLOW_DAG"})
+        airflow_task_type = await create_job_type(async_session, {"type": "AIRFLOW_TASK"})
         cluster_job = await create_job(
             async_session,
             location_id=cluster_location.id,
-            job_type_id=job_type.id,
+            job_type_id=cluster_type.id,
             job_kwargs={"name": "my-job_cluster"},
         )
         dag_job = await create_job(
             async_session,
             location_id=airflow_location.id,
-            job_type_id=job_type.id,
+            job_type_id=airflow_dag_type.id,
             job_kwargs={"name": "my-job_dag"},
         )
         task_job = await create_job(
             async_session,
             location_id=airflow_location.id,
-            job_type_id=job_type.id,
+            job_type_id=airflow_task_type.id,
             job_kwargs={"name": "my-job_task"},
         )
 

tests/test_server/test_jobs/test_search_jobs.py

@@ -254,7 +254,7 @@ async def test_get_job_types(
 ) -> None:
     unique_job_type = {item.type for item in job_types}
     response = await test_client.get(
-        "/v1/jobs/job_types",
+        "/v1/jobs/job-types",
         headers={"Authorization": f"Bearer {mocked_user.access_token}"},
     )
 
@@ -265,47 +265,21 @@ async def test_get_job_types(
 async def test_search_jobs_by_location_id(
     test_client: AsyncClient,
     async_session: AsyncSession,
-    jobs_search_with_different_locations: tuple[Job],
+    jobs_with_locations_and_types: tuple[Job],
     mocked_user: MockedUser,
 ) -> None:
-    jobs = await enrich_jobs(jobs_search_with_different_locations, async_session)
-    response = await test_client.get(
-        "/v1/jobs",
-        headers={"Authorization": f"Bearer {mocked_user.access_token}"},
-        params={"search_query": "my-job"},
-    )
-
-    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": 3,
-        },
-        "items": [
-            {
-                "id": str(job.id),
-                "data": job_to_json(job),
-            }
-            for job in jobs
-        ],
-    }
+    jobs = await enrich_jobs(jobs_with_locations_and_types, async_session)
 
     # first job in jobs has a different location unlike two others
     [_, dag_job, task_job] = jobs
-    response_with_location_id_filter = await test_client.get(
+    response = await test_client.get(
         "/v1/jobs",
         headers={"Authorization": f"Bearer {mocked_user.access_token}"},
-        params={"search_query": "my-job", "location_id": dag_job.location_id},
+        params={"location_id": dag_job.location_id},
     )
 
-    assert response_with_location_id_filter.status_code == HTTPStatus.OK, response_with_location_id_filter.json()
-    assert response_with_location_id_filter.json() == {
+    assert response.status_code == HTTPStatus.OK, response.json()
+    assert response.json() == {
         "meta": {
             "has_next": False,
             "has_previous": False,
@@ -329,17 +303,17 @@ async def test_search_jobs_by_location_id(
 async def test_search_jobs_by_location_id_non_existen_id(
     test_client: AsyncClient,
     async_session: AsyncSession,
-    jobs_search_with_different_locations: tuple[Job],
+    jobs_with_locations_and_types: tuple[Job],
     mocked_user: MockedUser,
 ):
-    response_with_location_id_filter = await test_client.get(
+    response = await test_client.get(
         "/v1/jobs",
         headers={"Authorization": f"Bearer {mocked_user.access_token}"},
-        params={"search_query": "my-job", "location_id": -1},
+        params={"location_id": -1},
     )
 
-    assert response_with_location_id_filter.status_code == HTTPStatus.OK, response_with_location_id_filter.json()
-    assert response_with_location_id_filter.json() == {
+    assert response.status_code == HTTPStatus.OK, response.json()
+    assert response.json() == {
         "meta": {
             "has_next": False,
             "has_previous": False,
@@ -354,18 +328,18 @@ async def test_search_jobs_by_location_id_non_existen_id(
     }
 
 
-async def test_search_jobs_by_name_and_type(
+async def test_search_by_jobs_type(
     test_client: AsyncClient,
     async_session: AsyncSession,
-    jobs_search_with_different_types: tuple[Job],
+    jobs_with_locations_and_types: tuple[Job],
     mocked_user: MockedUser,
 ) -> None:
-    jobs = await enrich_jobs(jobs_search_with_different_types, async_session)
-    job_with_dag_type = jobs[0]
+    jobs = await enrich_jobs(jobs_with_locations_and_types, async_session)
+    [_, dag_job, task_job] = jobs
     response = await test_client.get(
         "/v1/jobs",
         headers={"Authorization": f"Bearer {mocked_user.access_token}"},
-        params={"search_query": "airflow"},
+        params={"job_type": ["AIRFLOW_DAG", "AIRFLOW_TASK"]},
     )
 
     assert response.status_code == HTTPStatus.OK, response.json()
@@ -385,51 +359,25 @@ async def test_search_jobs_by_name_and_type(
                 "id": str(job.id),
                 "data": job_to_json(job),
             }
-            for job in jobs
+            for job in (dag_job, task_job)
         ],
     }
 
-    response_with_type_filter = await test_client.get(
-        "/v1/jobs",
-        headers={"Authorization": f"Bearer {mocked_user.access_token}"},
-        params={"job_type": "AIRFLOW_DAG", "search_query": "airflow"},
-    )
 
-    assert response_with_type_filter.status_code == HTTPStatus.OK, response_with_type_filter.json()
-    assert response_with_type_filter.json() == {
-        "meta": {
-            "has_next": False,
-            "has_previous": False,
-            "next_page": None,
-            "page": 1,
-            "page_size": 20,
-            "pages_count": 1,
-            "previous_page": None,
-            "total_count": 1,
-        },
-        "items": [
-            {
-                "id": str(job_with_dag_type.id),
-                "data": job_to_json(job_with_dag_type),
-            },
-        ],
-    }
-
-
-async def test_search_jobs_by_name_and_type_non_existen_type(
+async def test_search_jobs_by_non_existen_type(
     test_client: AsyncClient,
     async_session: AsyncSession,
-    jobs_search_with_different_types: tuple[Job],
+    jobs_with_locations_and_types: tuple[Job],
     mocked_user: MockedUser,
 ) -> None:
-    response_with_type_filter = await test_client.get(
+    response = await test_client.get(
         "/v1/jobs",
         headers={"Authorization": f"Bearer {mocked_user.access_token}"},
-        params={"job_type": "NO_EXISTEN_TYPE", "search_query": "airflow"},
+        params={"job_type": "NO_EXISTEN_TYPE"},
     )
 
-    assert response_with_type_filter.status_code == HTTPStatus.OK, response_with_type_filter.json()
-    assert response_with_type_filter.json() == {
+    assert response.status_code == HTTPStatus.OK, response.json()
+    assert response.json() == {
         "meta": {
             "has_next": False,
             "has_previous": False,