[DOP-31903] Include all found parents into job dependencies response

Author: dolfinus

data_rentgen/server/services/job.py

@@ -64,7 +64,7 @@ class JobServicePaginatedResult(PaginationDTO[JobServiceResult]):
 
 @dataclass
 class JobDependenciesResult:
-    parents: list[tuple[int, int]] = field(default_factory=list)
+    parents: set[tuple[int, int]] = field(default_factory=set)
     dependencies: set[tuple[int, int, str | None]] = field(default_factory=set)
     jobs: list[JobServiceResult] = field(default_factory=list)
 
@@ -113,21 +113,25 @@ async def get_job_dependencies(
         direction: Literal["UPSTREAM", "DOWNSTREAM", "BOTH"],
     ) -> JobDependenciesResult:
         logger.info("Get Job dependencies with start at job with id %s and direction: %s", start_node_id, direction)
-        job_ids = {start_node_id}
 
         ancestor_relations = await self._uow.job.list_ancestor_relations([start_node_id])
         descendant_relations = await self._uow.job.list_descendant_relations([start_node_id])
-        job_ids |= {p_id for p_id, _ in ancestor_relations}
-        job_ids |= {c_id for _, c_id in descendant_relations}
+        job_ids = (
+            {start_node_id}
+            | {p.parent_job_id for p in ancestor_relations}
+            | {p.child_job_id for p in descendant_relations}
+        )
 
         dependencies = await self._uow.job_dependency.get_dependencies(job_ids=list(job_ids), direction=direction)
-        for dependency in dependencies:
-            job_ids.add(dependency.from_job_id)
-            job_ids.add(dependency.to_job_id)
-        jobs = await self._uow.job.list_by_ids(list(job_ids))
+        dependency_job_ids = {d.from_job_id for d in dependencies} | {d.to_job_id for d in dependencies}
+        job_ids |= dependency_job_ids
 
+        # return ancestors for all found jobs in the graph
+        ancestor_relations += await self._uow.job.list_ancestor_relations(list(dependency_job_ids))
+        job_ids |= {p.parent_job_id for p in ancestor_relations}
+        jobs = await self._uow.job.list_by_ids(list(job_ids))
         return JobDependenciesResult(
-            parents=ancestor_relations + descendant_relations,
+            parents={(p.parent_job_id, p.child_job_id) for p in ancestor_relations + descendant_relations},
             dependencies={(d.from_job_id, d.to_job_id, d.type) for d in dependencies},
             jobs=[JobServiceResult.from_orm(job) for job in jobs],
         )

tests/test_server/test_jobs/test_job_dependencies.py

@@ -83,13 +83,15 @@ async def test_get_job_dependencies_with_direction_both(
     mocked_user: MockedUser,
 ):
     (
-        (_, dag2, _),
+        (dag1, dag2, dag3),
         (task1, task2, task3),
         (_, spark2, _),
     ) = job_dependency_chain
     expected_nodes = await enrich_jobs(
         [
+            dag1,
             dag2,
+            dag3,
             task1,
             task2,
             task3,
@@ -107,7 +109,7 @@ async def test_get_job_dependencies_with_direction_both(
         assert response.status_code == HTTPStatus.OK, response.json()
         assert response.json() == {
             "relations": {
-                "parents": jobs_ancestors_to_json([task2, spark2]),
+                "parents": jobs_ancestors_to_json(expected_nodes),
                 "dependencies": [
                     {
                         "from": {"kind": "JOB", "id": str(from_id)},
@@ -132,9 +134,13 @@ async def test_get_job_dependencies_with_direction_upstream(
     async_session: AsyncSession,
     mocked_user: MockedUser,
 ):
-    (_, dag2, _), (task1, task2, _), (_, spark2, _) = job_dependency_chain
+    (
+        (dag1, dag2, _),
+        (task1, task2, _),
+        (_, spark2, _),
+    ) = job_dependency_chain
     expected_nodes = await enrich_jobs(
-        [task1, dag2, task2, spark2],
+        [dag1, task1, dag2, task2, spark2],
         async_session,
     )
 
@@ -147,7 +153,7 @@ async def test_get_job_dependencies_with_direction_upstream(
         assert response.status_code == HTTPStatus.OK, response.json()
         assert response.json() == {
             "relations": {
-                "parents": jobs_ancestors_to_json([task2, spark2]),
+                "parents": jobs_ancestors_to_json(expected_nodes),
                 "dependencies": [
                     {
                         "from": {"kind": "JOB", "id": str(task1.id)},
@@ -166,9 +172,13 @@ async def test_get_job_dependencies_with_direction_downstream(
     async_session: AsyncSession,
     mocked_user: MockedUser,
 ):
-    (_, dag2, _), (_, task2, task3), (_, spark2, _) = job_dependency_chain
+    (
+        (_, dag2, dag3),
+        (_, task2, task3),
+        (_, spark2, _),
+    ) = job_dependency_chain
     expected_nodes = await enrich_jobs(
-        [dag2, task2, spark2, task3],
+        [dag2, task2, spark2, dag3, task3],
         async_session,
     )
 
@@ -181,7 +191,7 @@ async def test_get_job_dependencies_with_direction_downstream(
         assert response.status_code == HTTPStatus.OK, response.json()
         assert response.json() == {
             "relations": {
-                "parents": jobs_ancestors_to_json([task2, spark2]),
+                "parents": jobs_ancestors_to_json(expected_nodes),
                 "dependencies": [
                     {
                         "from": {"kind": "JOB", "id": str(task2.id)},