[DOP-25645] add tests

Author: TiGrib

data_rentgen/consumer/extractors/batch_extraction_result.py

@@ -136,6 +136,8 @@ def add_run(self, run: RunDTO):
 
     def add_operation(self, operation: OperationDTO):
         operation.run = self.add_run(operation.run)
+        if operation.sql_query:
+            operation.sql_query = self.add_sql_query(operation.sql_query)
         return self._add(self._operations, operation)
 
     def add_input(self, input_: InputDTO):

data_rentgen/consumer/extractors/batch_extractor.py

@@ -49,7 +49,6 @@ def extract_run(self, event: OpenLineageRunEvent) -> None:
     def extract_operation(self, event: OpenLineageRunEvent) -> None:
         operation = extract_operation(event)
         self.result.add_operation(operation)
-        # TODO here we don't have a query id, so we should store relation to query dto instead of only id
 
         for input_dataset in event.inputs:
             input_dto, symlink_dtos = extract_input(operation, input_dataset)

data_rentgen/consumer/extractors/job.py

@@ -1,12 +1,11 @@
 # SPDX-FileCopyrightText: 2024-2025 MTS PJSC
 # SPDX-License-Identifier: Apache-2.0
 
-from textwrap import dedent
 from urllib.parse import urlparse
 
 from data_rentgen.consumer.openlineage.job import OpenLineageJob
 from data_rentgen.consumer.openlineage.run_facets import OpenLineageParentJob
-from data_rentgen.dto import JobDTO, JobTypeDTO, LocationDTO, SQLQueryDTO
+from data_rentgen.dto import JobDTO, JobTypeDTO, LocationDTO
 
 
 def extract_parent_job(job: OpenLineageParentJob) -> JobDTO:
@@ -21,7 +20,6 @@ def extract_job(job: OpenLineageJob) -> JobDTO:
         name=job.name,
         location=extract_job_location(job),
         type=extract_job_type(job),
-        sql_query=extract_job_sql_query(job),
     )
 
 
@@ -44,14 +42,3 @@ def extract_job_type(job: OpenLineageJob) -> JobTypeDTO | None:
         return JobTypeDTO(type=type_.upper())
 
     return None
-
-
-def extract_job_sql_query(job: OpenLineageJob) -> SQLQueryDTO | None:
-    """
-    Sql queries are usual has format of multiline string. So we remove additional spaces and end of the rows symbols.
-    """
-    if job.facets.sql_query:
-        query = str.strip(dedent(job.facets.sql_query.query))
-        return SQLQueryDTO(query=query)
-
-    return None

data_rentgen/consumer/extractors/operation.py

@@ -1,12 +1,15 @@
 # SPDX-FileCopyrightText: 2024-2025 MTS PJSC
 # SPDX-License-Identifier: Apache-2.0
 
+from textwrap import dedent
+
 from data_rentgen.consumer.extractors.run import extract_parent_run, extract_run
+from data_rentgen.consumer.openlineage.job import OpenLineageJob
 from data_rentgen.consumer.openlineage.run_event import (
     OpenLineageRunEvent,
     OpenLineageRunEventType,
 )
-from data_rentgen.dto import OperationDTO, OperationStatusDTO, OperationTypeDTO
+from data_rentgen.dto import OperationDTO, OperationStatusDTO, OperationTypeDTO, SQLQueryDTO
 
 
 def extract_operation(event: OpenLineageRunEvent) -> OperationDTO:
@@ -32,7 +35,7 @@ def extract_operation(event: OpenLineageRunEvent) -> OperationDTO:
         name=operation_name,
         type=type_,
         status=OperationStatusDTO(run.status),
-        sql_query=run.job.sql_query,
+        sql_query=extract_sql_query(event.job),
         started_at=run.started_at,
         ended_at=run.ended_at,
     )
@@ -70,3 +73,14 @@ def enrich_operation_description(operation: OperationDTO, event: OpenLineageRunE
         operation.group = spark_job_details.jobGroup
         operation.description = spark_job_details.jobDescription
     return operation
+
+
+def extract_sql_query(job: OpenLineageJob) -> SQLQueryDTO | None:
+    """
+    Sql queries are usual has format of multiline string. So we remove additional spaces and end of the rows symbols.
+    """
+    if job.facets.sql:
+        query = str.strip(dedent(job.facets.sql.query))
+        return SQLQueryDTO(query=query)
+
+    return None

data_rentgen/consumer/openlineage/job_facets/__init__.py

@@ -28,4 +28,4 @@ class OpenLineageJobFacets(OpenLineageBase):
 
     documentation: OpenLineageDocumentationJobFacet | None = None
     jobType: OpenLineageJobTypeJobFacet | None = None
-    sql_query: OpenLineageSqlJobFacet | None = None
+    sql: OpenLineageSqlJobFacet | None = None

data_rentgen/db/repositories/operation.py

@@ -28,11 +28,11 @@ async def create_or_update_bulk(self, operations: list[OperationDTO]) -> None:
                 "started_at": func.coalesce(insert_statement.excluded.started_at, Operation.started_at),
                 "ended_at": func.coalesce(insert_statement.excluded.ended_at, Operation.ended_at),
                 "description": func.coalesce(insert_statement.excluded.description, Operation.description),
+                "sql_query_id": func.coalesce(insert_statement.excluded.sql_query_id, Operation.sql_query_id),
                 "group": func.coalesce(insert_statement.excluded.group, Operation.group),
                 "position": func.coalesce(insert_statement.excluded.position, Operation.position),
             },
         )
-
         await self._session.execute(
             statement,
             [

data_rentgen/dto/job.py

@@ -7,15 +7,13 @@
 
 from data_rentgen.dto.job_type import JobTypeDTO
 from data_rentgen.dto.location import LocationDTO
-from data_rentgen.dto.sql_query import SQLQueryDTO
 
 
 @dataclass
 class JobDTO:
     name: str
     location: LocationDTO
     type: JobTypeDTO | None = None
-    sql_query: SQLQueryDTO | None = None
     id: int | None = field(default=None, compare=False)
 
     @property
@@ -33,16 +31,9 @@ def merge(self, new: JobDTO) -> JobDTO:
         else:
             type_ = new.type or self.type
 
-        sql_query: SQLQueryDTO | None
-        if new.sql_query and self.sql_query:
-            sql_query = self.sql_query.merge(new.sql_query)
-        else:
-            sql_query = new.sql_query or self.sql_query
-
         return JobDTO(
             location=self.location.merge(new.location),
             name=self.name,
             type=type_,
-            sql_query=sql_query,
             id=new.id or self.id,
         )

tests/test_consumer/resources/events_spark.jsonl

@@ -11,6 +11,7 @@
     OpenLineageJobFacets,
     OpenLineageJobProcessingType,
     OpenLineageJobTypeJobFacet,
+    OpenLineageSqlJobFacet,
 )
 from data_rentgen.consumer.openlineage.run import OpenLineageRun
 from data_rentgen.consumer.openlineage.run_event import (
@@ -29,6 +30,7 @@
 from data_rentgen.dto.location import LocationDTO
 from data_rentgen.dto.operation import OperationTypeDTO
 from data_rentgen.dto.run import RunDTO
+from data_rentgen.dto.sql_query import SQLQueryDTO
 
 
 def test_extractors_extract_operation_spark_job_no_details():
@@ -302,3 +304,62 @@ def test_extractors_extract_operation_spark_job_finished(
         started_at=None,
         ended_at=ended_at,
     )
+
+
+def test_extractors_extract_operation_spark_job_sql_query():
+    now = datetime(2024, 7, 5, 9, 6, 29, 462000, tzinfo=timezone.utc)
+    run_id = UUID("01908224-8410-79a2-8de6-a769ad6944c9")
+    operation_id = UUID("01908225-1fd7-746b-910c-70d24f2898b1")
+
+    operation = OpenLineageRunEvent(
+        eventType=OpenLineageRunEventType.START,
+        eventTime=now,
+        job=OpenLineageJob(
+            namespace="anything",
+            name="mysession.execute_some_command",
+            facets=OpenLineageJobFacets(
+                jobType=OpenLineageJobTypeJobFacet(
+                    processingType=OpenLineageJobProcessingType.BATCH,
+                    integration="SPARK",
+                    jobType="SQL_JOB",
+                ),
+                sql=OpenLineageSqlJobFacet(query="select id, name from schema.table where id = 1"),
+            ),
+        ),
+        run=OpenLineageRun(
+            runId=operation_id,
+            facets=OpenLineageRunFacets(
+                parent=OpenLineageParentRunFacet(
+                    job=OpenLineageParentJob(
+                        namespace="anything",
+                        name="mysession",
+                    ),
+                    run=OpenLineageParentRun(
+                        runId=run_id,
+                    ),
+                ),
+            ),
+        ),
+    )
+    assert extract_operation(operation) == OperationDTO(
+        id=operation_id,
+        run=RunDTO(
+            id=run_id,
+            job=JobDTO(
+                name="mysession",
+                location=LocationDTO(
+                    type="unknown",
+                    name="anything",
+                    addresses={"unknown://anything"},
+                ),
+            ),
+        ),
+        name="execute_some_command",
+        type=OperationTypeDTO.BATCH,
+        position=None,
+        description=None,
+        status=OperationStatusDTO.STARTED,
+        sql_query=SQLQueryDTO(query="select id, name from schema.table where id = 1"),
+        started_at=now,
+        ended_at=None,
+    )

tests/test_consumer/test_extractors/test_extractors_operation_spark.py

@@ -27,6 +27,7 @@
     Run,
     RunStatus,
     Schema,
+    SQLQuery,
 )
 
 RESOURCES_PATH = Path(__file__).parent.parent.joinpath("resources").resolve()
@@ -95,6 +96,15 @@ async def test_runs_handler_spark(
     assert application_run.running_log_url == "http://127.0.0.1:4040"
     assert application_run.persistent_log_url is None
 
+    sql_query = select(SQLQuery).order_by(SQLQuery.id)
+    sql_query_scalars = await async_session.scalars(sql_query)
+    sql_queries = sql_query_scalars.all()
+    assert len(sql_queries) == 1
+
+    operation_sql_query = sql_queries[0]
+    assert operation_sql_query.fingerprint == UUID("250fe871-527a-5bf7-b730-3f89b42cceb2")
+    assert operation_sql_query.query == "select id, name from schema.table where id = 1"
+
     operation_query = select(Operation).order_by(Operation.id)
     operation_scalars = await async_session.scalars(operation_query)
     operations = operation_scalars.all()
@@ -108,6 +118,7 @@ async def test_runs_handler_spark(
     assert job_operation.type == OperationType.BATCH
     assert job_operation.status == OperationStatus.SUCCEEDED
     assert job_operation.started_at == datetime(2024, 7, 5, 9, 6, 29, 462000, tzinfo=timezone.utc)
+    assert job_operation.sql_query_id == operation_sql_query.id
     assert job_operation.ended_at == datetime(2024, 7, 5, 9, 7, 15, 642000, tzinfo=timezone.utc)
     assert job_operation.position == 3
     assert job_operation.description == "Hive -> Clickhouse"