[DOP-24570] Add tests for lineage excluding symlinks with no IO

Author: dolfinus

Makefile

@@ -64,7 +64,7 @@ db-views: ##@DB Create views
 broker: broker-start ##@Broker Prepare broker (in docker)
 
 broker-start: ##Broker Start broker
-	docker compose -f docker-compose.test.yml up -d --wait broker $(DOCKER_COMPOSE_ARGS)
+	docker compose -f docker-compose.test.yml --profile consumer up -d --wait $(DOCKER_COMPOSE_ARGS)
 
 
 test: test-db test-broker ##@Test Run tests
@@ -84,7 +84,7 @@ test-db-start: ##@TestDB Start database
 test-broker: test-broker-start ##@TestBroker Prepare broker (in docker)
 
 test-broker-start: ##@TestBroker Start broker
-	docker compose -f docker-compose.test.yml up -d --wait broker $(DOCKER_COMPOSE_ARGS)
+	docker compose -f docker-compose.test.yml --profile consumer up -d --wait $(DOCKER_COMPOSE_ARGS)
 
 test-ci: test-db test-broker ##@Test Run CI tests
 	${POETRY} run coverage run -m pytest
@@ -93,7 +93,7 @@ test-check-fixtures: ##@Test Check declared fixtures
 	${POETRY} run pytest --dead-fixtures $(PYTEST_ARGS)
 
 test-cleanup: ##@Test Cleanup tests dependencies
-	docker compose -f docker-compose.test.yml down $(ARGS)
+	docker compose -f docker-compose.test.yml --profile all down --remove-orphans $(ARGS)
 
 
 
@@ -110,7 +110,7 @@ prod: ##@Application Run production containers
 	docker compose up -d
 
 prod-cleanup: ##@Application Stop production containers
-	docker compose down $(ARGS)
+	docker compose down --remove-orphans $(ARGS)
 
 
 .PHONY: docs

docker-compose.test.yml

@@ -27,6 +27,9 @@ services:
       interval: 10s
       timeout: 5s
       retries: 5
+    profiles:
+      - consumer
+      - all
 
   keycloak:
     image: quay.io/keycloak/keycloak:latest
@@ -39,6 +42,9 @@ services:
       - 8080:8080
     volumes:
       - keycloak_data:/opt/keycloak/data
+    profiles:
+      - keycloak
+      - all
 
 volumes:
   postgres_test_data:

tests/test_server/fixtures/factories/address.py

@@ -30,6 +30,7 @@ async def create_address(
     else:
         address_kwargs = {"location_id": location_id}
     address = address_factory(**address_kwargs)
+    del address.id
     async_session.add(address)
     await async_session.commit()
     await async_session.refresh(address)

tests/test_server/fixtures/factories/lineage.py

@@ -691,6 +691,35 @@ async def lineage_with_symlinks(
         await clean_db(async_session)
 
 
+@pytest_asyncio.fixture()
+async def lineage_with_unconnected_symlinks(
+    lineage_with_depth: LineageResult,
+    async_session_maker: Callable[[], AbstractAsyncContextManager[AsyncSession]],
+) -> AsyncGenerator[LineageResult, None]:
+    # Same as lineage_with_depth, but each dataset has also a symlink,
+    # not connected to any input or output.
+
+    lineage = lineage_with_depth
+
+    async with async_session_maker() as async_session:
+        existing_datasets = lineage.datasets.copy()
+        for dataset in existing_datasets:
+            another_location = await create_location(async_session)
+            another_dataset = await create_dataset(async_session, location_id=another_location.id)
+            lineage.datasets.append(another_dataset)
+
+            metastore = [await make_symlink(async_session, another_dataset, dataset, DatasetSymlinkType.METASTORE)]
+            lineage.dataset_symlinks.extend(metastore)
+
+            warehouse = [await make_symlink(async_session, dataset, another_dataset, DatasetSymlinkType.WAREHOUSE)]
+            lineage.dataset_symlinks.extend(warehouse)
+
+    yield lineage
+
+    async with async_session_maker() as async_session:
+        await clean_db(async_session)
+
+
 @pytest_asyncio.fixture
 async def duplicated_lineage_with_column_lineage(
     async_session_maker: Callable[[], AbstractAsyncContextManager[AsyncSession]],

tests/test_server/test_lineage/test_dataset_lineage.py

@@ -950,6 +950,73 @@ async def test_get_dataset_lineage_with_symlink(
     }
 
 
+async def test_get_dataset_lineage_with_symlink_without_input_output(
+    test_client: AsyncClient,
+    async_session: AsyncSession,
+    lineage_with_unconnected_symlinks: LineageResult,
+    mocked_user: MockedUser,
+):
+    lineage = lineage_with_unconnected_symlinks
+    # Start from any dataset between J0 and J1, as it has both inputs and outputs
+    dataset = lineage.datasets[1]
+
+    inputs = [input for input in lineage.inputs if input.dataset_id == dataset.id]
+    assert inputs
+
+    outputs = [output for output in lineage.outputs if output.dataset_id == dataset.id]
+    assert outputs
+
+    operation_ids = {input.operation_id for input in inputs} | {output.operation_id for output in outputs}
+    operations = [operation for operation in lineage.operations if operation.id in operation_ids]
+    assert operations
+
+    run_ids = {operation.run_id for operation in operations}
+    runs = [run for run in lineage.runs if run.id in run_ids]
+    assert runs
+
+    job_ids = {run.job_id for run in runs}
+    jobs = [job for job in lineage.jobs if job.id in job_ids]
+    assert jobs
+
+    [dataset] = await enrich_datasets([dataset], async_session)
+    jobs = await enrich_jobs(jobs, async_session)
+    runs = await enrich_runs(runs, async_session)
+    since = min(run.created_at for run in runs)
+
+    response = await test_client.get(
+        "v1/datasets/lineage",
+        headers={"Authorization": f"Bearer {mocked_user.access_token}"},
+        params={
+            "since": since.isoformat(),
+            "start_node_id": dataset.id,
+        },
+    )
+
+    assert response.status_code == HTTPStatus.OK, response.json()
+    assert response.json() == {
+        "relations": {
+            "parents": run_parents_to_json(runs),
+            "symlinks": [],  # symlinks without inputs/outputs are excluded
+            "inputs": [
+                *inputs_to_json(merge_io_by_jobs(inputs), granularity="JOB"),
+                *inputs_to_json(merge_io_by_runs(inputs), granularity="RUN"),
+            ],
+            "outputs": [
+                *outputs_to_json(merge_io_by_jobs(outputs), granularity="JOB"),
+                *outputs_to_json(merge_io_by_runs(outputs), granularity="RUN"),
+            ],
+            "direct_column_lineage": [],
+            "indirect_column_lineage": [],
+        },
+        "nodes": {
+            "datasets": datasets_to_json([dataset]),
+            "jobs": jobs_to_json(jobs),
+            "runs": runs_to_json(runs),
+            "operations": {},
+        },
+    }
+
+
 @pytest.mark.parametrize("dataset_index", [0, 1], ids=["output", "input"])
 async def test_get_dataset_lineage_unmergeable_schema_and_output_type(
     test_client: AsyncClient,

tests/test_server/test_lineage/test_job_lineage.py

@@ -783,6 +783,59 @@ async def test_get_job_lineage_with_symlinks(
     }
 
 
+async def test_get_job_lineage_with_symlink_without_input_output(
+    test_client: AsyncClient,
+    async_session: AsyncSession,
+    lineage_with_unconnected_symlinks: LineageResult,
+    mocked_user: MockedUser,
+):
+    lineage = lineage_with_unconnected_symlinks
+
+    job = lineage.jobs[0]
+    inputs = [input for input in lineage.inputs if input.job_id == job.id]
+    assert inputs
+
+    outputs = [output for output in lineage.outputs if output.job_id == job.id]
+    assert outputs
+
+    dataset_ids = {input.dataset_id for input in inputs} | {output.dataset_id for output in outputs}
+    assert dataset_ids
+
+    datasets = [dataset for dataset in lineage.datasets if dataset.id in dataset_ids]
+    assert datasets
+
+    [job] = await enrich_jobs([job], async_session)
+    datasets = await enrich_datasets(datasets, async_session)
+    since = min(run.created_at for run in lineage.runs if run.job_id == job.id)
+
+    response = await test_client.get(
+        "v1/jobs/lineage",
+        headers={"Authorization": f"Bearer {mocked_user.access_token}"},
+        params={
+            "since": since.isoformat(),
+            "start_node_id": job.id,
+        },
+    )
+
+    assert response.status_code == HTTPStatus.OK, response.json()
+    assert response.json() == {
+        "relations": {
+            "parents": [],
+            "symlinks": [],  # symlinks without inputs/outputs are excluded
+            "inputs": inputs_to_json(merge_io_by_jobs(inputs), granularity="JOB"),
+            "outputs": outputs_to_json(merge_io_by_jobs(outputs), granularity="JOB"),
+            "direct_column_lineage": [],
+            "indirect_column_lineage": [],
+        },
+        "nodes": {
+            "datasets": datasets_to_json(datasets),
+            "jobs": jobs_to_json([job]),
+            "runs": {},
+            "operations": {},
+        },
+    }
+
+
 async def test_get_job_lineage_unmergeable_inputs_and_outputs(
     test_client: AsyncClient,
     async_session: AsyncSession,

tests/test_server/test_lineage/test_operation_lineage.py

@@ -640,6 +640,70 @@ async def test_get_operation_lineage_with_symlinks(
     }
 
 
+async def test_get_operation_lineage_with_symlink_without_input_output(
+    test_client: AsyncClient,
+    async_session: AsyncSession,
+    lineage_with_unconnected_symlinks: LineageResult,
+    mocked_user: MockedUser,
+):
+    lineage = lineage_with_unconnected_symlinks
+    operation = lineage.operations[0]
+
+    inputs = [input for input in lineage.inputs if input.operation_id == operation.id]
+    assert inputs
+
+    outputs = [output for output in lineage.outputs if output.operation_id == operation.id]
+    assert outputs
+
+    dataset_ids = {input.dataset_id for input in inputs} | {output.dataset_id for output in outputs}
+    assert dataset_ids
+
+    datasets = [dataset for dataset in lineage.datasets if dataset.id in dataset_ids]
+    assert datasets
+
+    run = next(run for run in lineage.runs if run.id == operation.run_id)
+    job = next(job for job in lineage.jobs if job.id == run.job_id)
+
+    [job] = await enrich_jobs([job], async_session)
+    [run] = await enrich_runs([run], async_session)
+    datasets = await enrich_datasets(datasets, async_session)
+
+    response = await test_client.get(
+        "v1/operations/lineage",
+        headers={"Authorization": f"Bearer {mocked_user.access_token}"},
+        params={
+            "since": run.created_at.isoformat(),
+            "start_node_id": str(operation.id),
+        },
+    )
+
+    assert response.status_code == HTTPStatus.OK, response.json()
+    assert response.json() == {
+        "relations": {
+            "parents": run_parents_to_json([run]) + operation_parents_to_json([operation]),
+            "symlinks": [],  # symlinks without inputs/outputs are excluded
+            "inputs": [
+                *inputs_to_json(merge_io_by_jobs(inputs), granularity="JOB"),
+                *inputs_to_json(inputs, granularity="OPERATION"),
+                *inputs_to_json(merge_io_by_runs(inputs), granularity="RUN"),
+            ],
+            "outputs": [
+                *outputs_to_json(merge_io_by_jobs(outputs), granularity="JOB"),
+                *outputs_to_json(outputs, granularity="OPERATION"),
+                *outputs_to_json(merge_io_by_runs(outputs), granularity="RUN"),
+            ],
+            "direct_column_lineage": [],
+            "indirect_column_lineage": [],
+        },
+        "nodes": {
+            "datasets": datasets_to_json(datasets),
+            "jobs": jobs_to_json([job]),
+            "runs": runs_to_json([run]),
+            "operations": operations_to_json([operation]),
+        },
+    }
+
+
 async def test_get_operation_lineage_with_empty_io_stats_and_schema(
     test_client: AsyncClient,
     async_session: AsyncSession,

tests/test_server/test_lineage/test_run_lineage.py

@@ -822,6 +822,67 @@ async def test_get_run_lineage_with_symlinks(
     }
 
 
+async def test_get_run_lineage_with_symlink_without_input_output(
+    test_client: AsyncClient,
+    async_session: AsyncSession,
+    lineage_with_unconnected_symlinks: LineageResult,
+    mocked_user: MockedUser,
+):
+    lineage = lineage_with_unconnected_symlinks
+
+    run = lineage.runs[0]
+    job = next(job for job in lineage.jobs if job.id == run.job_id)
+
+    inputs = [input for input in lineage.inputs if input.run_id == run.id]
+    assert inputs
+
+    outputs = [output for output in lineage.outputs if output.run_id == run.id]
+    assert outputs
+
+    dataset_ids = {input.dataset_id for input in inputs} | {output.dataset_id for output in outputs}
+    assert dataset_ids
+
+    datasets = [dataset for dataset in lineage.datasets if dataset.id in dataset_ids]
+    assert datasets
+
+    [job] = await enrich_jobs([job], async_session)
+    [run] = await enrich_runs([run], async_session)
+    datasets = await enrich_datasets(datasets, async_session)
+
+    response = await test_client.get(
+        "v1/runs/lineage",
+        headers={"Authorization": f"Bearer {mocked_user.access_token}"},
+        params={
+            "since": run.created_at.isoformat(),
+            "start_node_id": str(run.id),
+        },
+    )
+
+    assert response.status_code == HTTPStatus.OK, response.json()
+    assert response.json() == {
+        "relations": {
+            "parents": run_parents_to_json([run]),
+            "symlinks": [],  # symlinks without inputs/outputs are excluded
+            "inputs": [
+                *inputs_to_json(merge_io_by_jobs(inputs), granularity="JOB"),
+                *inputs_to_json(merge_io_by_runs(inputs), granularity="RUN"),
+            ],
+            "outputs": [
+                *outputs_to_json(merge_io_by_jobs(outputs), granularity="JOB"),
+                *outputs_to_json(merge_io_by_runs(outputs), granularity="RUN"),
+            ],
+            "direct_column_lineage": [],
+            "indirect_column_lineage": [],
+        },
+        "nodes": {
+            "datasets": datasets_to_json(datasets),
+            "jobs": jobs_to_json([job]),
+            "runs": runs_to_json([run]),
+            "operations": {},
+        },
+    }
+
+
 async def test_get_run_lineage_unmergeable_inputs_and_outputs(
     test_client: AsyncClient,
     async_session: AsyncSession,