Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 6 additions & 1 deletion data_rentgen/consumer/extractors/impl/spark.py
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,12 @@ def is_operation(self, event: OpenLineageRunEvent) -> bool:
def extract_run(self, event: OpenLineageRunEvent) -> RunDTO:
run = super().extract_run(event)
self._enrich_run_identifiers(run, event)
if run.job.name == "unknown":
# Workaround for https://github.com/OpenLineage/OpenLineage/issues/3846
# Avoid updating the same temporary job multiple times
run.job.parent_job = None
run.job.tag_values.clear()
run.job_dependencies.clear()
return run

def _enrich_run_identifiers(self, run: RunDTO, event: OpenLineageRunEvent):
Expand All @@ -60,7 +66,6 @@ def extract_operation(self, event: OpenLineageRunEvent) -> OperationDTO:
# For Spark, SQL_JOB --parent-> SPARK_APPLICATION = operation -> run,
# and parent is always here.
run = self.extract_parent_run(event.run.facets.parent) # type: ignore[arg-type]
# Workaround for https://github.com/OpenLineage/OpenLineage/issues/3846
self._enrich_run_identifiers(run, event)
self._enrich_run_tags(run, event)
operation = super()._extract_operation(event, run)
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,64 @@
# SPDX-FileCopyrightText: 2024-present MTS PJSC
# SPDX-License-Identifier: Apache-2.0
"""Detach Spark unknown jobs from everything

It should have no parent job, no dependencies and no tags.

Revision ID: 96fd9a096682
Revises: 4a02d2d5c8b1
Create Date: 2026-07-06 13:01:46.138108

"""

import sqlalchemy as sa
from alembic import op

# revision identifiers, used by Alembic.
revision = "96fd9a096682"
down_revision = "4a02d2d5c8b1"
branch_labels = None
depends_on = None

SPARK_UNKNOWN_JOBS_QUERY = """
SELECT job.*
FROM job
JOIN job_type ON job.type_id = job_type.id
WHERE job_type.type = 'SPARK_APPLICATION'
AND lower(job.name) = 'unknown'
"""


def upgrade() -> None:
op.execute(
sa.text(
f"""
WITH data AS ({SPARK_UNKNOWN_JOBS_QUERY})
UPDATE job
SET parent_job_id=NULL
WHERE id IN (SELECT id FROM data);
""",
),
)
op.execute(
sa.text(
f"""
WITH data AS ({SPARK_UNKNOWN_JOBS_QUERY})
DELETE FROM job_dependency
WHERE from_job_id IN (SELECT id FROM data) OR to_job_id IN (SELECT id FROM data);
""",
),
)
op.execute(
sa.text(
f"""
WITH data AS ({SPARK_UNKNOWN_JOBS_QUERY})
DELETE FROM job_tag_value
WHERE job_id IN (SELECT id FROM data);
""",
),
)


def downgrade() -> None:
# this migration is irreversible
pass
4 changes: 4 additions & 0 deletions docs/changelog/next_release/480.improvement.rst
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
If Data.Rentgen receives RunEvent with Spark job with ``name=unknown``, it should have no parent job, no tags and no dependencies.
Otherwise multiple workers will update the same row each time, although it is temporary.

Workaround for ``OpenLineage #3846<https://github.com/OpenLineage/OpenLineage/issues/3846>`_.
6 changes: 3 additions & 3 deletions tests/test_consumer/test_extractors/fixtures/spark_raw.py
Original file line number Diff line number Diff line change
Expand Up @@ -179,7 +179,7 @@ def spark_app_run_event_start() -> OpenLineageRunEvent:
facets=OpenLineageRunFacets(
spark_applicationDetails=OpenLineageSparkApplicationDetailsRunFacet(
master="local[*]",
appName="spark_session",
appName="mysession1",
applicationId="local-1719136537510",
deployMode=OpenLineageSparkDeployMode.CLIENT,
driverHost="127.0.0.1",
Expand Down Expand Up @@ -400,7 +400,7 @@ def spark_operation_run_event_start() -> OpenLineageRunEvent:
),
spark_applicationDetails=OpenLineageSparkApplicationDetailsRunFacet(
master="local[*]",
appName="spark_session",
appName="mysession1",
applicationId="local-1719136537510",
deployMode=OpenLineageSparkDeployMode.CLIENT,
driverHost="127.0.0.1",
Expand Down Expand Up @@ -445,7 +445,7 @@ def spark_operation_run_event_running() -> OpenLineageRunEvent:
),
spark_applicationDetails=OpenLineageSparkApplicationDetailsRunFacet(
master="local[*]",
appName="spark_session",
appName="mysession1",
applicationId="local-1719136537510",
deployMode=OpenLineageSparkDeployMode.CLIENT,
driverHost="127.0.0.1",
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -167,8 +167,6 @@ def test_extractors_extract_batch_spark_openlineage_emitted_unknown_name_no_job_

extracted = BatchExtractor().add_events(events)

extracted_spark_unknown_job.parent_job = extracted_airflow_task1_job_as_parent

# the result is the same as above
assert extracted.locations() == [
extracted_airflow_location,
Expand Down
72 changes: 72 additions & 0 deletions tests/test_consumer/test_extractors/test_extractors_run_spark.py
Original file line number Diff line number Diff line change
Expand Up @@ -370,3 +370,75 @@ def test_extractors_extract_run_spark_parent_job():
persistent_log_url=None,
running_log_url="http://localhost:4040",
)


def test_extractors_extract_run_spark_parent_job_for_unknown():
now = datetime(2024, 7, 5, 9, 4, 13, 979349, tzinfo=timezone.utc)
run_id = UUID("01908223-0e9b-7c52-9856-6cecfc842610")
parent_run_id = UUID("01908224-8410-79a2-8de6-a769ad6944c9")
run = OpenLineageRunEvent(
eventType=OpenLineageRunEventType.RUNNING,
eventTime=now,
job=OpenLineageJob(
namespace="host://some.host.com",
name="unknown",
facets=OpenLineageJobFacets(
jobType=OpenLineageJobTypeJobFacet(
processingType=OpenLineageJobProcessingType.NONE,
integration="SPARK",
jobType="APPLICATION",
),
),
),
run=OpenLineageRun(
runId=run_id,
facets=OpenLineageRunFacets(
parent=OpenLineageParentRunFacet(
job=OpenLineageParentJob(
namespace="anything",
name="parentjob",
),
run=OpenLineageParentRun(
runId=parent_run_id,
),
),
processing_engine=OpenLineageProcessingEngineRunFacet(
version=Version("3.4.3"),
name="spark",
openlineageAdapterVersion=Version("1.19.0"),
),
),
),
)

assert SparkExtractor().extract_run(run) == RunDTO(
id=run_id,
job=JobDTO(
name="unknown",
parent_job=None,
location=LocationDTO(type="host", name="some.host.com", addresses={"host://some.host.com"}),
type=JobTypeDTO(type="SPARK_APPLICATION"),
tag_values=set(),
),
parent_run=RunDTO(
id=parent_run_id,
job=JobDTO(
name="parentjob",
type=None,
location=LocationDTO(
type="unknown",
name="anything",
addresses={"unknown://anything"},
),
),
status=RunStatusDTO.UNKNOWN,
),
status=RunStatusDTO.STARTED,
started_at=None,
start_reason=None,
user=None,
external_id=None,
attempt=None,
persistent_log_url=None,
running_log_url=None,
)
121 changes: 121 additions & 0 deletions tests/test_database/test_migration_detach_spark_unknown_jobs.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,121 @@
# SPDX-FileCopyrightText: 2024-present MTS PJSC
# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations

from typing import TYPE_CHECKING

import pytest
from sqlalchemy import create_engine, text
from sqlalchemy.pool import NullPool

from data_rentgen.db.models import Base
from tests.test_database.fixtures.alembic import do_run_migrations

if TYPE_CHECKING:
from alembic.config import Config as AlembicConfig

pytestmark = [pytest.mark.db]

PREV_REVISION = "4a02d2d5c8b1"
THIS_REVISION = "96fd9a096682"


def test_migration_detach_spark_unknown_jobs(empty_db_url: str, alembic_config: AlembicConfig):
do_run_migrations(alembic_config, Base.metadata, PREV_REVISION)

engine = create_engine(empty_db_url, poolclass=NullPool)
try:
with engine.begin() as conn:
conn.execute(
text(
"""
INSERT INTO location (id, type, name) VALUES
(1, 'local', 'somehost');
""",
),
)

conn.execute(
text(
"""
INSERT INTO job (id, location_id, name, type_id, parent_job_id) VALUES
(1, 1, 'airflow_task1', 3, null),
(2, 1, 'airflow_task2', 3, null),
(3, 1, 'unknown', 1, 1),
(4, 1, 'some_session', 1, 2);
""",
),
)

conn.execute(
text(
"""
INSERT INTO tag (id, name) VALUES
(1, 'some_tag'),
(2, 'another_tag');
""",
),
)

conn.execute(
text(
"""
INSERT INTO tag_value (id, tag_id, value) VALUES
(1, 1, 'some_value'),
(2, 2, 'another_value');
""",
),
)

conn.execute(
text(
"""
INSERT INTO job_tag_value (job_id, tag_value_id) VALUES
(1, 1),
(2, 1),
(3, 2),
(4, 2);
""",
),
)

conn.execute(
text(
"""
INSERT INTO job_dependency (from_job_id, to_job_id) VALUES
(1, 2),
(3, 4);
""",
),
)

do_run_migrations(alembic_config, Base.metadata, THIS_REVISION)

with engine.connect() as conn:
assert conn.execute(
text("SELECT id, parent_job_id FROM job ORDER BY id"),
).fetchall() == [
(1, None),
(2, None),
(3, None), # parent_id is null
(4, 2),
]

assert conn.execute(
text("SELECT job_id, tag_value_id FROM job_tag_value ORDER BY job_id, tag_value_id"),
).fetchall() == [
(1, 1),
(2, 1),
# no item (3, 2)
(4, 2),
]

assert conn.execute(
text("SELECT from_job_id, to_job_id FROM job_dependency ORDER BY from_job_id, to_job_id"),
).fetchall() == [
(1, 2),
# no item (3, 4)
]

finally:
engine.dispose()