[DOP-22363] Small LineageService updates

Author: dolfinus

data_rentgen/db/repositories/column_lineage.py

@@ -74,17 +74,17 @@ async def list_by_job_ids(
         job_ids: Sequence[int],
         since: datetime,
         until: datetime | None,
-        source_ids: Sequence[int],
-        target_ids: Sequence[int],
+        source_dataset_ids: Sequence[int],
+        target_dataset_ids: Sequence[int],
     ):
         if not job_ids:
             return []
 
         where = [
             ColumnLineage.created_at >= since,
             ColumnLineage.job_id == any_(job_ids),  # type: ignore[arg-type]
-            ColumnLineage.source_dataset_id == any_(source_ids),  # type: ignore[arg-type]
-            ColumnLineage.target_dataset_id == any_(target_ids),  # type: ignore[arg-type]
+            ColumnLineage.source_dataset_id == any_(source_dataset_ids),  # type: ignore[arg-type]
+            ColumnLineage.target_dataset_id == any_(target_dataset_ids),  # type: ignore[arg-type]
         ]
         if until:
             where.append(ColumnLineage.created_at <= until)
@@ -96,17 +96,17 @@ async def list_by_run_ids(
         run_ids: Sequence[UUID],
         since: datetime,
         until: datetime | None,
-        source_ids: Sequence[int],
-        target_ids: Sequence[int],
+        source_dataset_ids: Sequence[int],
+        target_dataset_ids: Sequence[int],
     ):
         if not run_ids:
             return []
 
         where = [
             ColumnLineage.created_at >= since,
             ColumnLineage.run_id == any_(run_ids),  # type: ignore[arg-type]
-            ColumnLineage.source_dataset_id == any_(source_ids),  # type: ignore[arg-type]
-            ColumnLineage.target_dataset_id == any_(target_ids),  # type: ignore[arg-type]
+            ColumnLineage.source_dataset_id == any_(source_dataset_ids),  # type: ignore[arg-type]
+            ColumnLineage.target_dataset_id == any_(target_dataset_ids),  # type: ignore[arg-type]
         ]
         if until:
             where.append(ColumnLineage.created_at <= until)
@@ -115,8 +115,8 @@ async def list_by_run_ids(
     async def list_by_operation_ids(
         self,
         operation_ids: Sequence[UUID],
-        source_ids: Sequence[int],
-        target_ids: Sequence[int],
+        source_dataset_ids: Sequence[int],
+        target_dataset_ids: Sequence[int],
     ):
         if not operation_ids:
             return []
@@ -130,8 +130,8 @@ async def list_by_operation_ids(
             ColumnLineage.created_at >= min_created_at,
             ColumnLineage.created_at <= max_created_at,
             ColumnLineage.operation_id == any_(operation_ids),  # type: ignore[arg-type]
-            ColumnLineage.source_dataset_id == any_(source_ids),  # type: ignore[arg-type]
-            ColumnLineage.target_dataset_id == any_(target_ids),  # type: ignore[arg-type]
+            ColumnLineage.source_dataset_id == any_(source_dataset_ids),  # type: ignore[arg-type]
+            ColumnLineage.target_dataset_id == any_(target_dataset_ids),  # type: ignore[arg-type]
         ]
         return await self._get_column_lineage_with_column_relations(where)
 

data_rentgen/server/services/lineage.py

@@ -1100,29 +1100,32 @@ async def _get_column_lineage(
         if not current_result.inputs or not current_result.outputs:
             return result
 
+        source_dataset_ids = [input_.dataset_id for input_ in current_result.inputs.values()]
+        target_dataset_ids = [output.dataset_id for output in current_result.outputs.values()]
+
         match granularity:
             case "OPERATION":
                 column_lineage_result = await self._uow.column_lineage.list_by_operation_ids(
                     operation_ids=sorted(current_result.operations.keys()),
                     # return column lineage only for datasets included into response
-                    source_ids=[input_.dataset_id for input_ in current_result.inputs.values()],
-                    target_ids=[output.dataset_id for output in current_result.outputs.values()],
+                    source_dataset_ids=source_dataset_ids,
+                    target_dataset_ids=target_dataset_ids,
                 )
             case "RUN":
                 column_lineage_result = await self._uow.column_lineage.list_by_run_ids(
                     run_ids=sorted(current_result.runs.keys()),
                     since=since,
                     until=until,
-                    source_ids=[input_.dataset_id for input_ in current_result.inputs.values()],
-                    target_ids=[output.dataset_id for output in current_result.outputs.values()],
+                    source_dataset_ids=source_dataset_ids,
+                    target_dataset_ids=target_dataset_ids,
                 )
             case "JOB":
                 column_lineage_result = await self._uow.column_lineage.list_by_job_ids(
                     job_ids=sorted(current_result.jobs.keys()),
                     since=since,
                     until=until,
-                    source_ids=[input_.dataset_id for input_ in current_result.inputs.values()],
-                    target_ids=[output.dataset_id for output in current_result.outputs.values()],
+                    source_dataset_ids=source_dataset_ids,
+                    target_dataset_ids=target_dataset_ids,
                 )
             case _:
                 msg = f"Unknown granularity for column lineage: {granularity}"

tests/test_server/fixtures/factories/dataset.py

@@ -156,9 +156,8 @@ async def datasets_with_symlinks(
         await clean_db(async_session)
 
 
-@pytest_asyncio.fixture(params=[{}])
+@pytest_asyncio.fixture
 async def datasets_search(
-    request: pytest.FixtureRequest,
     async_session_maker: Callable[[], AbstractAsyncContextManager[AsyncSession]],
 ) -> AsyncGenerator[tuple[dict[str, Dataset], dict[str, Dataset], dict[str, Dataset]], None]:
     """

tests/test_server/fixtures/factories/job.py

@@ -97,9 +97,8 @@ async def jobs(
         await clean_db(async_session)
 
 
-@pytest_asyncio.fixture(params=[{}])
+@pytest_asyncio.fixture
 async def jobs_search(
-    request: pytest.FixtureRequest,
     async_session_maker: Callable[[], AbstractAsyncContextManager[AsyncSession]],
 ) -> AsyncGenerator[tuple[dict[str, Job], dict[str, Job], dict[str, Job]], None]:
     """