[DOP-25498] Do not restrict values of Job.type

Author: dolfinus

data_rentgen/consumer/extractors/batch_extraction_result.py

@@ -10,6 +10,7 @@
     DatasetSymlinkDTO,
     InputDTO,
     JobDTO,
+    JobTypeDTO,
     LocationDTO,
     OperationDTO,
     OutputDTO,
@@ -25,6 +26,7 @@
     ColumnLineageDTO,
     DatasetSymlinkDTO,
     JobDTO,
+    JobTypeDTO,
     RunDTO,
     OperationDTO,
     InputDTO,
@@ -56,6 +58,7 @@ def __init__(self):
         self._locations: dict[tuple, LocationDTO] = {}
         self._datasets: dict[tuple, DatasetDTO] = {}
         self._dataset_symlinks: dict[tuple, DatasetSymlinkDTO] = {}
+        self._job_types: dict[tuple, JobTypeDTO] = {}
         self._jobs: dict[tuple, JobDTO] = {}
         self._runs: dict[tuple, RunDTO] = {}
         self._operations: dict[tuple, OperationDTO] = {}
@@ -71,6 +74,7 @@ def __repr__(self):
             f"locations={len(self._locations)}, "
             f"datasets={len(self._datasets)}, "
             f"dataset_symlinks={len(self._dataset_symlinks)}, "
+            f"job_types={len(self._job_types)}, "
             f"jobs={len(self._jobs)}, "
             f"runs={len(self._runs)}, "
             f"operations={len(self._operations)}, "
@@ -109,8 +113,13 @@ def add_dataset_symlink(self, dataset_symlink: DatasetSymlinkDTO):
         dataset_symlink.to_dataset = self.add_dataset(dataset_symlink.to_dataset)
         return self._add(self._dataset_symlinks, dataset_symlink)
 
+    def add_job_type(self, job_type: JobTypeDTO):
+        return self._add(self._job_types, job_type)
+
     def add_job(self, job: JobDTO):
         job.location = self.add_location(job.location)
+        if job.type:
+            job.type = self.add_job_type(job.type)
         return self._add(self._jobs, job)
 
     def add_run(self, run: RunDTO):
@@ -171,9 +180,14 @@ def get_dataset_symlink(self, dataset_symlink_key: tuple) -> DatasetSymlinkDTO:
         dataset_symlink.to_dataset = self.get_dataset(dataset_symlink.to_dataset.unique_key)
         return dataset_symlink
 
+    def get_job_type(self, job_type_key: tuple) -> JobTypeDTO:
+        return self._job_types[job_type_key]
+
     def get_job(self, job_key: tuple) -> JobDTO:
         job = self._jobs[job_key]
         job.location = self.get_location(job.location.unique_key)
+        if job.type:
+            job.type = self.get_job_type(job.type.unique_key)
         return job
 
     def get_run(self, run_key: tuple) -> RunDTO:
@@ -222,6 +236,9 @@ def datasets(self) -> list[DatasetDTO]:
     def dataset_symlinks(self) -> list[DatasetSymlinkDTO]:
         return list(map(self.get_dataset_symlink, self._dataset_symlinks))
 
+    def job_types(self) -> list[JobTypeDTO]:
+        return list(map(self.get_job_type, self._job_types))
+
     def jobs(self) -> list[JobDTO]:
         return list(map(self.get_job, self._jobs))
 

data_rentgen/consumer/extractors/batch_extractor.py

@@ -8,7 +8,6 @@
 from data_rentgen.consumer.extractors.operation import extract_operation
 from data_rentgen.consumer.extractors.output import extract_output
 from data_rentgen.consumer.extractors.run import extract_run
-from data_rentgen.consumer.openlineage.job_facets.job_type import OpenLineageJobType
 from data_rentgen.consumer.openlineage.run_event import OpenLineageRunEvent
 from data_rentgen.dto import (
     DatasetDTO,
@@ -22,13 +21,25 @@ def __init__(self) -> None:
 
     def add_events(self, events: list[OpenLineageRunEvent]) -> BatchExtractionResult:
         for event in events:
-            if event.job.facets.jobType and event.job.facets.jobType.jobType == OpenLineageJobType.JOB:
+            if self.is_operation(event):
                 self.extract_operation(event)
             else:
                 self.extract_run(event)
-
         return self.result
 
+    def is_operation(self, event: OpenLineageRunEvent) -> bool:
+        job_type_facet = event.job.facets.jobType
+        if not job_type_facet:
+            return False
+
+        if job_type_facet.integration == "SPARK":
+            return job_type_facet.jobType != "APPLICATION"
+
+        if job_type_facet.integration == "AIRFLOW":
+            return False
+
+        return False
+
     def extract_run(self, event: OpenLineageRunEvent) -> None:
         run = extract_run(event)
         self.result.add_run(run)

data_rentgen/consumer/extractors/job.py

@@ -36,8 +36,9 @@ def extract_job_location(job: OpenLineageJob | OpenLineageParentJob) -> Location
 
 def extract_job_type(job: OpenLineageJob) -> JobTypeDTO | None:
     if job.facets.jobType:
-        job_type = job.facets.jobType.jobType
         integration_type = job.facets.jobType.integration
-        return JobTypeDTO(f"{integration_type}_{job_type}")
+        job_type = job.facets.jobType.jobType
+        type_ = f"{integration_type}_{job_type}" if job_type else integration_type
+        return JobTypeDTO(type=type_.upper())
 
     return None

data_rentgen/consumer/extractors/operation.py

@@ -20,11 +20,15 @@ def extract_operation(event: OpenLineageRunEvent) -> OperationDTO:
         prefix = len(run.job.name) + 1
         operation_name = operation_name[prefix:]
 
+    type_: OperationTypeDTO = OperationTypeDTO.BATCH
+    if event.job.facets.jobType and event.job.facets.jobType.processingType:
+        type_ = OperationTypeDTO(event.job.facets.jobType.processingType)
+
     operation = OperationDTO(
         id=event.run.runId,  # type: ignore [arg-type]
         run=run,
         name=operation_name,
-        type=OperationTypeDTO(event.job.facets.jobType.processingType) if event.job.facets.jobType else None,
+        type=type_,
     )
     enrich_operation_status(operation, event)
     enrich_operation_description(operation, event)

data_rentgen/consumer/openlineage/job_facets/__init__.py

@@ -7,19 +7,15 @@
     OpenLineageDocumentationJobFacet,
 )
 from data_rentgen.consumer.openlineage.job_facets.job_type import (
-    OpenLineageJobIntegrationType,
     OpenLineageJobProcessingType,
-    OpenLineageJobType,
     OpenLineageJobTypeJobFacet,
 )
 
 __all__ = [
     "OpenLineageDocumentationJobFacet",
     "OpenLineageJobFacet",
     "OpenLineageJobFacets",
-    "OpenLineageJobIntegrationType",
     "OpenLineageJobProcessingType",
-    "OpenLineageJobType",
     "OpenLineageJobTypeJobFacet",
 ]
 

data_rentgen/consumer/openlineage/job_facets/job_type.py

@@ -3,64 +3,24 @@
 
 from enum import Enum
 
-from pydantic import field_validator
-
 from data_rentgen.consumer.openlineage.job_facets.base import OpenLineageJobFacet
 
 
-class OpenLineageJobIntegrationType(str, Enum):
-    """Integration where job is running.
-    See [JobTypeJobFacet](https://github.com/OpenLineage/OpenLineage/blob/main/spec/facets/JobTypeJobFacet.json).
-    """
-
-    SPARK = "SPARK"
-    AIRFLOW = "AIRFLOW"
-
-    def __str__(self) -> str:
-        return self.value
-
-
-class OpenLineageJobType(str, Enum):
-    """Job type.
-    See [JobTypeJobFacet](https://github.com/OpenLineage/OpenLineage/blob/main/spec/facets/JobTypeJobFacet.json).
-    """
-
-    APPLICATION = "APPLICATION"
-    JOB = "JOB"
-    DAG = "DAG"
-    TASK = "TASK"
-
-    def __str__(self) -> str:
-        return self.value
-
-    @classmethod
-    def _missing_(cls, value):
-        if value in {"SQL_JOB", "RDD_JOB"}:
-            return cls.JOB
-        return None
-
-
 class OpenLineageJobProcessingType(str, Enum):
     """Job processing type.
     See [JobTypeJobFacet](https://github.com/OpenLineage/OpenLineage/blob/main/spec/facets/JobTypeJobFacet.json).
     """
 
     BATCH = "BATCH"
     STREAMING = "STREAMING"
+    NONE = "NONE"
 
 
 class OpenLineageJobTypeJobFacet(OpenLineageJobFacet):
     """Job facet describing job type.
     See [JobTypeJobFacet](https://github.com/OpenLineage/OpenLineage/blob/main/spec/facets/JobTypeJobFacet.json).
     """
 
-    integration: OpenLineageJobIntegrationType
-    jobType: OpenLineageJobType
-    processingType: OpenLineageJobProcessingType | None = None
-
-    @field_validator("processingType", mode="before")
-    @classmethod
-    def _validate_processing_type(cls, processing_type: str):
-        if processing_type == "NONE":
-            return None
-        return processing_type
+    processingType: OpenLineageJobProcessingType
+    integration: str
+    jobType: str | None = None

data_rentgen/consumer/subscribers.py

@@ -65,7 +65,7 @@ async def extract_events(
             event = OpenLineageRunEventAdapter.validate_json(message.value)
             extractor.add_events([event])
         except (ValueError, TypeError):
-            logger.error(  # noqa: TRY400
+            logger.exception(
                 "Failed to parse message: ConsumerRecord(topic=%r, partition=%d, offset=%d)",
                 message.topic,
                 message.partition,
@@ -109,6 +109,12 @@ async def save_to_db(
             dataset_symlink = await unit_of_work.dataset_symlink.create_or_update(dataset_symlink_dto)
             dataset_symlink_dto.id = dataset_symlink.id
 
+    logger.debug("Creating job types")
+    for job_type_dto in data.job_types():
+        async with unit_of_work:
+            job_type = await unit_of_work.job_type.get_or_create(job_type_dto)
+            job_type_dto.id = job_type.id
+
     logger.debug("Creating jobs")
     for job_dto in data.jobs():
         async with unit_of_work:

data_rentgen/db/migrations/versions/2025-04-25_2d2fe3f2f348_add_job_type.py

@@ -0,0 +1,101 @@
+# SPDX-FileCopyrightText: 2024-2025 MTS PJSC
+# SPDX-License-Identifier: Apache-2.0
+"""Add job_type
+
+Revision ID: 2d2fe3f2f348
+Revises: 976168ee4f16
+Create Date: 2025-04-25 15:09:17.556969
+
+"""
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "2d2fe3f2f348"
+down_revision = "976168ee4f16"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    op.create_table(
+        "job_type",
+        sa.Column("id", sa.BigInteger(), nullable=False),
+        sa.Column("type", sa.String(), nullable=False),
+        sa.PrimaryKeyConstraint("id", name=op.f("pk__job_type")),
+        sa.UniqueConstraint("type", name=op.f("uq__job_type__type")),
+    )
+    op.create_index(op.f("ix__job_type__type"), "job_type", ["type"], unique=False)
+
+    op.execute(
+        sa.text(
+            """
+        INSERT INTO
+            job_type (id, type)
+        VALUES
+            (0, 'UNKNOWN'),
+            (1, 'SPARK_APPLICATION'),
+            (2, 'AIRFLOW_DAG'),
+            (3, 'AIRFLOW_TASK');
+        """,
+        ),
+    )
+    op.execute(sa.text("ALTER SEQUENCE job_type_id_seq RESTART WITH 4;"))
+
+    op.execute(sa.text("LOCK TABLE job IN ACCESS EXCLUSIVE MODE;"))
+    op.drop_index("ix__job__type", table_name="job")
+    op.alter_column(
+        "job",
+        "type",
+        new_column_name="type_id",
+        existing_type=sa.String(length=32),
+        type_=sa.BigInteger(),
+        nullable=False,
+        postgresql_using="""
+            CASE
+                WHEN type = 'SPARK_APPLICATION'
+                    THEN 1
+                WHEN type = 'AIRFLOW_DAG'
+                    THEN 2
+                WHEN type = 'AIRFLOW_TASK'
+                    THEN 3
+                ELSE 0
+            END
+        """,
+    )
+    op.create_index(op.f("ix__job__type_id"), "job", ["type_id"], unique=False)
+    op.create_foreign_key(
+        op.f("fk__job__type_id__job_type"),
+        "job",
+        "job_type",
+        ["type_id"],
+        ["id"],
+        ondelete="CASCADE",
+    )
+
+
+def downgrade() -> None:
+    op.execute(sa.text("LOCK TABLE job IN ACCESS EXCLUSIVE MODE;"))
+    op.drop_constraint(op.f("fk__job__type_id__job_type"), "job", type_="foreignkey")
+    op.drop_index(op.f("ix__job__type_id"), table_name="job")
+    op.alter_column(
+        "job",
+        "type_id",
+        new_column_name="type",
+        existing_type=sa.BigInteger(),
+        type_=sa.String(length=32),
+        nullable=False,
+    )
+    op.execute(
+        sa.text(
+            """
+        UPDATE job
+        SET type = (SELECT job_type.type FROM job_type WHERE job_type.id = job.type);
+        """,
+        ),
+    )
+    op.create_index("ix__job__type", "job", ["type"], unique=False)
+
+    op.drop_index(op.f("ix__job_type__type"), table_name="job_type")
+    op.drop_table("job_type")

data_rentgen/db/models/__init__.py

@@ -13,7 +13,8 @@
 )
 from data_rentgen.db.models.dataset_symlink import DatasetSymlink, DatasetSymlinkType
 from data_rentgen.db.models.input import Input
-from data_rentgen.db.models.job import Job, JobType
+from data_rentgen.db.models.job import Job
+from data_rentgen.db.models.job_type import JobType
 from data_rentgen.db.models.location import Location
 from data_rentgen.db.models.operation import Operation, OperationStatus, OperationType
 from data_rentgen.db.models.output import Output, OutputType