This is an automated email from the ASF dual-hosted git repository.
shahar1 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 0f0873630af Emit per-statement OpenLineage events for BigQuery script
jobs (#69234)
0f0873630af is described below
commit 0f0873630af036a698bb7aec31d54e5944bd9d76
Author: Aaron Chen <[email protected]>
AuthorDate: Thu Jul 23 01:24:00 2026 +0800
Emit per-statement OpenLineage events for BigQuery script jobs (#69234)
---
.../providers/google/cloud/openlineage/mixins.py | 116 +++++-
.../unit/google/cloud/openlineage/test_mixins.py | 390 ++++++++++++++++++++-
2 files changed, 497 insertions(+), 9 deletions(-)
diff --git
a/providers/google/src/airflow/providers/google/cloud/openlineage/mixins.py
b/providers/google/src/airflow/providers/google/cloud/openlineage/mixins.py
index 539c24f6533..3e5520721ed 100644
--- a/providers/google/src/airflow/providers/google/cloud/openlineage/mixins.py
+++ b/providers/google/src/airflow/providers/google/cloud/openlineage/mixins.py
@@ -21,6 +21,7 @@ import copy
import json
import traceback
from collections.abc import Iterable
+from datetime import datetime, timezone
from typing import TYPE_CHECKING, cast
from airflow.providers.common.compat.openlineage.facet import (
@@ -53,7 +54,7 @@ if TYPE_CHECKING:
class _BigQueryInsertJobOperatorOpenLineageMixin:
"""Mixin for BigQueryInsertJobOperator to extract OpenLineage metadata."""
- def get_openlineage_facets_on_complete(self, _):
+ def get_openlineage_facets_on_complete(self, task_instance):
"""
Retrieve OpenLineage data for a completed BigQuery job.
@@ -109,14 +110,45 @@ class _BigQueryInsertJobOperatorOpenLineageMixin:
run_facets["bigQueryJob"] =
self._get_bigquery_job_run_facet(job_properties)
if get_from_nullable_chain(job_properties, ["statistics",
"numChildJobs"]):
- self.log.debug("Found SCRIPT job. Extracting lineage from
child jobs instead.")
- # SCRIPT job type has no input / output information but spawns
child jobs that have one
+ self.log.debug("Found SCRIPT job. Extracting lineage from
child jobs.")
+ # SCRIPT job has no input/output of its own but spawns child
jobs that do. The parent
+ # task keeps the aggregated coarse-grained lineage (backward
compatible) while each
+ # child query is additionally emitted below as its own event
for per-statement detail.
#
https://cloud.google.com/bigquery/docs/information-schema-jobs#multi-statement_query_job
- for child_job_id in
self._client.list_jobs(parent_job=self.job_id):
- child_job_properties =
self._client.get_job(job_id=child_job_id)._properties
- child_inputs, child_outputs =
self._get_inputs_and_outputs(child_job_properties)
+ child_jobs_properties = []
+ for child_job in
self._client.list_jobs(parent_job=self.job_id):
+ try:
+ child_jobs_properties.append(
+
self._client.get_job(job_id=child_job.job_id)._properties
+ )
+ except Exception as child_exception:
+ self.log.warning(
+ "Cannot retrieve BigQuery child job `%s`. %s",
+ child_job.job_id,
+ child_exception,
+ exc_info=True,
+ )
+ child_jobs_properties.sort(key=self._get_child_job_sort_key)
+ for child_index, child_job_properties in
enumerate(child_jobs_properties, start=1):
+ try:
+ child_inputs, child_outputs =
self._get_inputs_and_outputs(child_job_properties)
+ except Exception as child_exception:
+ self.log.warning(
+ "Cannot extract lineage for BigQuery child job
`%s`. %s",
+ get_from_nullable_chain(child_job_properties,
["jobReference", "jobId"]),
+ child_exception,
+ exc_info=True,
+ )
+ continue
inputs.extend(child_inputs)
outputs.extend(child_outputs)
+ self._emit_child_query_lineage(
+ task_instance=task_instance,
+ child_index=child_index,
+ child_job_properties=child_job_properties,
+ inputs=child_inputs,
+ outputs=child_outputs,
+ )
else:
inputs, outputs = self._get_inputs_and_outputs(job_properties)
@@ -139,6 +171,78 @@ class _BigQueryInsertJobOperatorOpenLineageMixin:
job_facets={"sql":
SQLJobFacet(query=SQLParser.normalize_sql(self.sql))} if self.sql else {},
)
+ def _emit_child_query_lineage(
+ self,
+ *,
+ task_instance,
+ child_index: int,
+ child_job_properties: dict,
+ inputs: list[InputDataset | Dataset],
+ outputs: list[OutputDataset | Dataset],
+ ) -> None:
+ if task_instance is None:
+ self.log.debug("No task instance available. Skipping BigQuery
child job OpenLineage event.") # type: ignore[attr-defined]
+ return
+
+ try:
+ from airflow.providers.openlineage.api.sql import
emit_query_lineage
+ except ImportError:
+ self.log.debug( # type: ignore[attr-defined]
+ "The emit_query_lineage API requires
apache-airflow-providers-openlineage>=2.16.0. "
+ "Skipping BigQuery child job OpenLineage event."
+ )
+ return
+
+ from airflow.providers.openlineage.sqlparser import SQLParser
+
+ child_query = get_from_nullable_chain(child_job_properties,
["configuration", "query", "query"])
+ job_facets = {"sql":
SQLJobFacet(query=SQLParser.normalize_sql(child_query))} if child_query else
None
+ error_result = get_from_nullable_chain(child_job_properties,
["status", "errorResult"])
+ start_time = self._get_bigquery_job_datetime(child_job_properties,
"startTime")
+ end_time = self._get_bigquery_job_datetime(child_job_properties,
"endTime")
+ if start_time is None or end_time is None:
+ start_time = end_time = None
+ emit_query_lineage(
+ query_id=get_from_nullable_chain(child_job_properties,
["jobReference", "jobId"]),
+ query_source_namespace=BIGQUERY_NAMESPACE,
+ # Intentionally not passed as query_text: BigQuery job metadata
already provides
+ # authoritative lineage, and parser-derived datasets could
duplicate or conflict
+ # with it. The SQL is still attached via the sql facet in
additional_job_facets.
+ query_text=None,
+ inputs=inputs,
+ outputs=outputs,
+ start_time=start_time,
+ end_time=end_time,
+ is_successful=error_result is None,
+ error_message=error_result.get("message") if error_result else
None,
+ task_instance=task_instance,
+
job_name=f"{task_instance.dag_id}.{task_instance.task_id}.query.{child_index}",
+ additional_run_facets={"bigQueryJob":
self._get_bigquery_job_run_facet(child_job_properties)},
+ additional_job_facets=job_facets, # type: ignore[arg-type]
+ )
+
+ @staticmethod
+ def _get_bigquery_job_datetime(properties: dict, field_name: str) ->
datetime | None:
+ value = get_from_nullable_chain(properties, ["statistics", field_name])
+ if value is None:
+ return None
+ try:
+ return datetime.fromtimestamp(float(value) / 1000, tz=timezone.utc)
+ except (TypeError, ValueError, OverflowError):
+ return None
+
+ @classmethod
+ def _get_child_job_sort_key(cls, properties: dict) -> tuple[datetime, str]:
+ # Emit children ordered by execution time so the query.N suffix is
stable across runs,
+ # breaking ties deterministically by job id.
+ start_time = cls._get_bigquery_job_datetime(properties, "startTime")
+ return (
+ # None is not comparable with datetime, so a missing startTime
maps to
+ # datetime.max to sort those children last instead of crashing the
sort.
+ start_time or datetime.max.replace(tzinfo=timezone.utc),
+ get_from_nullable_chain(properties, ["jobReference", "jobId"]) or
"",
+ )
+
def _get_inputs_and_outputs(self, properties: dict) ->
tuple[list[InputDataset], list[OutputDataset]]:
job_type = get_from_nullable_chain(properties, ["configuration",
"jobType"])
diff --git
a/providers/google/tests/unit/google/cloud/openlineage/test_mixins.py
b/providers/google/tests/unit/google/cloud/openlineage/test_mixins.py
index e7204368fd5..0106eb81eb3 100644
--- a/providers/google/tests/unit/google/cloud/openlineage/test_mixins.py
+++ b/providers/google/tests/unit/google/cloud/openlineage/test_mixins.py
@@ -20,10 +20,12 @@ import copy
import json
import logging
import os
+from datetime import datetime, timezone
from unittest.mock import MagicMock, patch
import pytest
from google.cloud.bigquery.table import Table
+from openlineage.client.event_v2 import RunState
from airflow.providers.common.compat.openlineage.facet import (
ColumnLineageDatasetFacet,
@@ -110,6 +112,31 @@ def read_common_json_file(rel: str):
return json.load(f)
+def make_task_instance():
+ logical_date = datetime(2024, 1, 1, tzinfo=timezone.utc)
+ dag_run = MagicMock(
+ logical_date=logical_date,
+ clear_number=0,
+ run_after=logical_date,
+ conf={},
+ )
+ ti = MagicMock(
+ dag_id="dag_id",
+ task_id="task_id",
+ try_number=1,
+ map_index=-1,
+ logical_date=logical_date,
+ )
+ ti.dag_run = dag_run
+ ti.get_template_context.return_value = {
+ "dag_run": dag_run,
+ "dag": MagicMock(),
+ "task": MagicMock(),
+ "task_instance": ti,
+ }
+ return ti
+
+
class TestBigQueryOpenLineageMixin:
def setup_method(self):
self.copy_job_details = read_common_json_file("copy_job_details.json")
@@ -395,7 +422,8 @@ class TestBigQueryOpenLineageMixin:
),
]
- def test_get_openlineage_facets_on_complete_script_job(self):
+ @patch("airflow.providers.openlineage.api.sql.emit_query_lineage")
+ def test_get_openlineage_facets_on_complete_script_job(self,
mock_emit_query_lineage):
self.client.get_job.side_effect = [
MagicMock(_properties=self.script_job_details),
MagicMock(_properties=self.query_job_details),
@@ -404,9 +432,10 @@ class TestBigQueryOpenLineageMixin:
Table.from_api_repr(read_common_json_file("table_details.json")),
Table.from_api_repr(read_common_json_file("out_table_details.json")),
]
- self.client.list_jobs.return_value = ["child_job_id"]
+ self.client.list_jobs.return_value = [MagicMock(job_id="child_job_id")]
+ mock_ti = make_task_instance()
- lineage = self.operator.get_openlineage_facets_on_complete(None)
+ lineage = self.operator.get_openlineage_facets_on_complete(mock_ti)
self.script_job_details["configuration"]["query"].pop("query")
assert lineage.run_facets == {
@@ -456,6 +485,361 @@ class TestBigQueryOpenLineageMixin:
},
),
]
+ mock_emit_query_lineage.assert_called_once()
+ assert (
+ mock_emit_query_lineage.call_args.kwargs["query_id"]
+ == self.query_job_details["jobReference"]["jobId"]
+ )
+ assert
mock_emit_query_lineage.call_args.kwargs["query_source_namespace"] == "bigquery"
+ assert mock_emit_query_lineage.call_args.kwargs["query_text"] is None
+ assert mock_emit_query_lineage.call_args.kwargs["is_successful"] is
True
+ assert mock_emit_query_lineage.call_args.kwargs["error_message"] is
None
+ assert mock_emit_query_lineage.call_args.kwargs["task_instance"] is
mock_ti
+ assert mock_emit_query_lineage.call_args.kwargs["job_name"] ==
"dag_id.task_id.query.1"
+ assert mock_emit_query_lineage.call_args.kwargs["start_time"] ==
datetime.fromtimestamp(
+ self.query_job_details["statistics"]["startTime"] / 1000,
tz=timezone.utc
+ )
+ assert mock_emit_query_lineage.call_args.kwargs["end_time"] ==
datetime.fromtimestamp(
+ self.query_job_details["statistics"]["endTime"] / 1000,
tz=timezone.utc
+ )
+
+ @patch.object(
+ _BigQueryInsertJobOperatorOpenLineageMixin,
+ "_get_inputs_and_outputs",
+ autospec=True,
+ )
+ @patch("airflow.providers.openlineage.api.sql.emit_query_lineage")
+ def
test_script_job_aggregates_parent_datasets_and_emits_child_query_lineage(
+ self, mock_emit_query_lineage, mock_get_inputs_and_outputs
+ ):
+ parent_job_details = copy.deepcopy(self.script_job_details)
+ parent_job_details["statistics"]["numChildJobs"] = "2"
+ child_job_1_details = {
+ "jobReference": {"jobId": "child_job_1"},
+ "configuration": {
+ "jobType": "QUERY",
+ "query": {"query": "CREATE TABLE output_table1 AS SELECT 1 AS
id"},
+ },
+ "statistics": {"query": {"cacheHit": False, "totalBytesBilled":
"10"}},
+ "status": {"state": "DONE"},
+ }
+ child_job_2_details = {
+ "jobReference": {"jobId": "child_job_2"},
+ "configuration": {
+ "jobType": "QUERY",
+ "query": {"query": "CREATE TABLE output_table2 AS SELECT 2 AS
id"},
+ },
+ "statistics": {"query": {"cacheHit": False, "totalBytesBilled":
"20"}},
+ "status": {"state": "DONE"},
+ }
+ input_table1 = InputDataset(namespace="bigquery",
name="project.dataset.input_table1")
+ output_table1 = OutputDataset(namespace="bigquery",
name="project.dataset.output_table1")
+ input_table2 = InputDataset(namespace="bigquery",
name="project.dataset.input_table2")
+ output_table2 = OutputDataset(namespace="bigquery",
name="project.dataset.output_table2")
+ self.client.get_job.side_effect = [
+ MagicMock(_properties=parent_job_details),
+ MagicMock(_properties=child_job_1_details),
+ MagicMock(_properties=child_job_2_details),
+ ]
+ self.client.list_jobs.return_value = [
+ MagicMock(job_id="child_job_1"),
+ MagicMock(job_id="child_job_2"),
+ ]
+ mock_ti = make_task_instance()
+
+ def get_inputs_and_outputs(_, properties):
+ query = properties["configuration"]["query"]["query"]
+ if "output_table1" in query:
+ return [input_table1], [output_table1]
+ return [input_table2], [output_table2]
+
+ mock_get_inputs_and_outputs.side_effect = get_inputs_and_outputs
+
+ lineage = self.operator.get_openlineage_facets_on_complete(mock_ti)
+
+ assert lineage.inputs == [input_table1, input_table2]
+ assert lineage.outputs == [output_table1, output_table2]
+ assert "bigQueryJob" in lineage.run_facets
+ assert "externalQuery" in lineage.run_facets
+ assert mock_emit_query_lineage.call_count == 2
+
+ first_call, second_call = mock_emit_query_lineage.call_args_list
+ assert first_call.kwargs["query_id"] == "child_job_1"
+ assert first_call.kwargs["query_source_namespace"] == "bigquery"
+ assert first_call.kwargs["inputs"] == [input_table1]
+ assert first_call.kwargs["outputs"] == [output_table1]
+ assert first_call.kwargs["task_instance"] is mock_ti
+ assert first_call.kwargs["job_name"] == "dag_id.task_id.query.1"
+
+ assert second_call.kwargs["query_id"] == "child_job_2"
+ assert second_call.kwargs["query_source_namespace"] == "bigquery"
+ assert second_call.kwargs["inputs"] == [input_table2]
+ assert second_call.kwargs["outputs"] == [output_table2]
+ assert second_call.kwargs["task_instance"] is mock_ti
+ assert second_call.kwargs["job_name"] == "dag_id.task_id.query.2"
+
+ @patch.object(
+ _BigQueryInsertJobOperatorOpenLineageMixin,
+ "_get_inputs_and_outputs",
+ autospec=True,
+ )
+
@patch("airflow.providers.openlineage.api.sql.resolve_task_emission_policy")
+ @patch("airflow.providers.openlineage.api.sql.is_openlineage_active",
return_value=True)
+ @patch("airflow.providers.openlineage.api.sql.emit")
+ def test_script_job_builds_child_query_events(
+ self,
+ mock_emit,
+ mock_is_openlineage_active,
+ mock_resolve_task_emission_policy,
+ mock_get_inputs_and_outputs,
+ ):
+ parent_job_details = copy.deepcopy(self.script_job_details)
+ parent_job_details["statistics"]["numChildJobs"] = "2"
+ child_job_1_details = {
+ "jobReference": {"jobId": "child_job_1"},
+ "configuration": {
+ "jobType": "QUERY",
+ "query": {"query": "CREATE TABLE output_table1 AS SELECT 1 AS
id"},
+ },
+ "statistics": {
+ "startTime": "1600000000000",
+ "endTime": "1600000005000",
+ "query": {"cacheHit": False, "totalBytesBilled": "10"},
+ },
+ "status": {"state": "DONE"},
+ }
+ child_job_2_details = {
+ "jobReference": {"jobId": "child_job_2"},
+ "configuration": {
+ "jobType": "QUERY",
+ "query": {"query": "CREATE TABLE output_table2 AS SELECT 2 AS
id"},
+ },
+ "statistics": {
+ "startTime": "1600000010000",
+ "endTime": "1600000015000",
+ "query": {"cacheHit": False, "totalBytesBilled": "20"},
+ },
+ "status": {"state": "DONE"},
+ }
+ input_table1 = InputDataset(namespace="bigquery",
name="project.dataset.input_table1")
+ output_table1 = OutputDataset(namespace="bigquery",
name="project.dataset.output_table1")
+ input_table2 = InputDataset(namespace="bigquery",
name="project.dataset.input_table2")
+ output_table2 = OutputDataset(namespace="bigquery",
name="project.dataset.output_table2")
+
+ class ChildJob:
+ def __init__(self, job_id):
+ self.job_id = job_id
+
+ job_details_by_id = {
+ "job_id": parent_job_details,
+ "child_job_1": child_job_1_details,
+ "child_job_2": child_job_2_details,
+ }
+ self.client.get_job.side_effect = lambda job_id:
MagicMock(_properties=job_details_by_id[job_id])
+ self.client.list_jobs.return_value = [ChildJob("child_job_2"),
ChildJob("child_job_1")]
+
+ def get_inputs_and_outputs(_, properties):
+ query = properties["configuration"]["query"]["query"]
+ if "output_table1" in query:
+ return [input_table1], [output_table1]
+ return [input_table2], [output_table2]
+
+ mock_get_inputs_and_outputs.side_effect = get_inputs_and_outputs
+ mock_resolve_task_emission_policy.return_value = MagicMock(emit=True)
+
+ lineage =
self.operator.get_openlineage_facets_on_complete(make_task_instance())
+
+ # Aggregation follows the sorted (execution-time) order, not the
list_jobs order.
+ assert lineage.inputs == [input_table1, input_table2]
+ assert lineage.outputs == [output_table1, output_table2]
+ assert mock_is_openlineage_active.call_count == 2
+ assert mock_emit.call_count == 4
+ child_1_start, child_1_complete, child_2_start, child_2_complete = [
+ call.args[0] for call in mock_emit.call_args_list
+ ]
+
+ assert child_1_start.eventType == RunState.START
+ assert child_1_complete.eventType == RunState.COMPLETE
+ assert child_1_complete.job.name == "dag_id.task_id.query.1"
+ assert child_1_complete.run.facets["externalQuery"].externalQueryId ==
"child_job_1"
+ assert child_1_complete.run.facets["externalQuery"].source ==
"bigquery"
+ assert child_1_complete.inputs == [input_table1]
+ assert child_1_complete.outputs == [output_table1]
+ assert child_1_start.eventTime == "2020-09-13T12:26:40+00:00"
+ assert child_1_complete.eventTime == "2020-09-13T12:26:45+00:00"
+ assert child_1_complete.run.facets["bigQueryJob"].billedBytes == 10
+ assert child_1_complete.job.facets["sql"].query == "CREATE TABLE
output_table1 AS SELECT 1 AS id"
+
+ assert child_2_start.eventType == RunState.START
+ assert child_2_complete.eventType == RunState.COMPLETE
+ assert child_2_complete.job.name == "dag_id.task_id.query.2"
+ assert child_2_complete.run.facets["externalQuery"].externalQueryId ==
"child_job_2"
+ assert child_2_complete.inputs == [input_table2]
+ assert child_2_complete.outputs == [output_table2]
+ assert child_2_start.eventTime == "2020-09-13T12:26:50+00:00"
+ assert child_2_complete.eventTime == "2020-09-13T12:26:55+00:00"
+
+ @patch.object(
+ _BigQueryInsertJobOperatorOpenLineageMixin,
+ "_get_inputs_and_outputs",
+ autospec=True,
+ )
+ @patch("airflow.providers.openlineage.api.sql.emit_query_lineage")
+ def test_script_job_continues_after_child_lineage_failure(
+ self, mock_emit_query_lineage, mock_get_inputs_and_outputs
+ ):
+ parent_job_details = copy.deepcopy(self.script_job_details)
+ parent_job_details["statistics"]["numChildJobs"] = "2"
+ child_job_1_details = {
+ "jobReference": {"jobId": "child_job_1"},
+ "configuration": {"jobType": "QUERY"},
+ "statistics": {"startTime": "1600000000000"},
+ "status": {"state": "DONE"},
+ }
+ child_job_2_details = {
+ "jobReference": {"jobId": "child_job_2"},
+ "configuration": {"jobType": "QUERY", "query": {"query": "SELECT
2"}},
+ "statistics": {"startTime": "1600000010000"},
+ "status": {"state": "DONE"},
+ }
+ input_table2 = InputDataset(namespace="bigquery",
name="project.dataset.input_table2")
+ output_table2 = OutputDataset(namespace="bigquery",
name="project.dataset.output_table2")
+ self.client.get_job.side_effect = [
+ MagicMock(_properties=parent_job_details),
+ MagicMock(_properties=child_job_1_details),
+ MagicMock(_properties=child_job_2_details),
+ ]
+ self.client.list_jobs.return_value = [
+ MagicMock(job_id="child_job_1"),
+ MagicMock(job_id="child_job_2"),
+ ]
+ mock_get_inputs_and_outputs.side_effect = [
+ RuntimeError("broken child"),
+ ([input_table2], [output_table2]),
+ ]
+
+ lineage =
self.operator.get_openlineage_facets_on_complete(make_task_instance())
+
+ assert lineage.inputs == [input_table2]
+ assert lineage.outputs == [output_table2]
+ assert "errorMessage" not in lineage.run_facets
+ mock_emit_query_lineage.assert_called_once()
+ assert mock_emit_query_lineage.call_args.kwargs["query_id"] ==
"child_job_2"
+ assert mock_emit_query_lineage.call_args.kwargs["inputs"] ==
[input_table2]
+ assert mock_emit_query_lineage.call_args.kwargs["outputs"] ==
[output_table2]
+ # The failing child keeps its positional slot so the query.N suffix
stays stable across runs.
+ assert mock_emit_query_lineage.call_args.kwargs["job_name"] ==
"dag_id.task_id.query.2"
+
+ @patch("airflow.providers.openlineage.api.sql.emit_query_lineage")
+ def
test_script_job_without_task_instance_does_not_emit_child_query_events(self,
mock_emit_query_lineage):
+ self.client.get_job.side_effect = [
+ MagicMock(_properties=self.script_job_details),
+ MagicMock(_properties=self.query_job_details),
+ ]
+ self.client.get_table.side_effect = [
+ Table.from_api_repr(read_common_json_file("table_details.json")),
+
Table.from_api_repr(read_common_json_file("out_table_details.json")),
+ ]
+ self.client.list_jobs.return_value = [MagicMock(job_id="child_job_id")]
+
+ lineage = self.operator.get_openlineage_facets_on_complete(None)
+
+ # Parent still aggregates child datasets; only the per-child emission
is skipped without a TI.
+ assert [i.name for i in lineage.inputs] ==
["airflow-openlineage.new_dataset.test_table"]
+ assert [o.name for o in lineage.outputs] ==
["airflow-openlineage.new_dataset.output_table"]
+ mock_emit_query_lineage.assert_not_called()
+ # Without the None-TI guard, building the child job_name dereferences
None and the failure
+ # surfaces as an errorMessage facet; its absence proves the guard
short-circuited cleanly.
+ assert "errorMessage" not in lineage.run_facets
+
+ @patch("airflow.providers.openlineage.api.sql.emit_query_lineage")
+ def test_child_query_lineage_without_query_omits_sql_job_facet(self,
mock_emit_query_lineage):
+ self.operator._emit_child_query_lineage(
+ task_instance=make_task_instance(),
+ child_index=1,
+ child_job_properties={
+ "configuration": {"jobType": "QUERY", "query": {}},
+ "statistics": {"query": {"cacheHit": False,
"totalBytesBilled": "10"}},
+ },
+ inputs=[],
+ outputs=[],
+ )
+
+ assert
mock_emit_query_lineage.call_args.kwargs["additional_job_facets"] is None
+
+ @pytest.mark.parametrize(
+ "statistics",
+ [
+ {"startTime": "invalid", "endTime": "1600000005000"},
+ {"startTime": "1600000000000", "endTime": "invalid"},
+ ],
+ )
+ @patch("airflow.providers.openlineage.api.sql.emit_query_lineage")
+ def test_child_query_lineage_omits_partial_timestamps(self,
mock_emit_query_lineage, statistics):
+ self.operator._emit_child_query_lineage(
+ task_instance=make_task_instance(),
+ child_index=1,
+ child_job_properties={
+ "configuration": {"jobType": "QUERY", "query": {}},
+ "statistics": statistics,
+ },
+ inputs=[],
+ outputs=[],
+ )
+
+ assert mock_emit_query_lineage.call_args.kwargs["start_time"] is None
+ assert mock_emit_query_lineage.call_args.kwargs["end_time"] is None
+
+ @patch("airflow.providers.openlineage.api.sql.emit_query_lineage")
+ def test_child_query_lineage_marks_failed_child_job(self,
mock_emit_query_lineage):
+ self.operator._emit_child_query_lineage(
+ task_instance=make_task_instance(),
+ child_index=1,
+ child_job_properties={
+ "jobReference": {"jobId": "child_job_id"},
+ "configuration": {"jobType": "QUERY", "query": {"query":
"SELECT 1"}},
+ "status": {"state": "DONE", "errorResult": {"reason":
"invalid", "message": "Syntax error"}},
+ },
+ inputs=[],
+ outputs=[],
+ )
+
+ assert mock_emit_query_lineage.call_args.kwargs["is_successful"] is
False
+ assert mock_emit_query_lineage.call_args.kwargs["error_message"] ==
"Syntax error"
+
+ @patch.dict("sys.modules", {"airflow.providers.openlineage.api.sql": None})
+ def test_child_query_lineage_skipped_with_old_openlineage_provider(self):
+ self.client.get_job.side_effect = [
+ MagicMock(_properties=self.script_job_details),
+ MagicMock(_properties=self.query_job_details),
+ ]
+ self.client.get_table.side_effect = [
+ Table.from_api_repr(read_common_json_file("table_details.json")),
+
Table.from_api_repr(read_common_json_file("out_table_details.json")),
+ ]
+ self.client.list_jobs.return_value = [MagicMock(job_id="child_job_id")]
+
+ lineage =
self.operator.get_openlineage_facets_on_complete(make_task_instance())
+
+ # Old OpenLineage providers lack emit_query_lineage: the parent event
must still
+ # aggregate child datasets and must not gain a misleading errorMessage
facet.
+ assert [i.name for i in lineage.inputs] ==
["airflow-openlineage.new_dataset.test_table"]
+ assert [o.name for o in lineage.outputs] ==
["airflow-openlineage.new_dataset.output_table"]
+ assert "errorMessage" not in lineage.run_facets
+
+ @pytest.mark.parametrize(
+ ("value", "expected"),
+ [
+ ("1600000000000", datetime(2020, 9, 13, 12, 26, 40,
tzinfo=timezone.utc)),
+ (None, None),
+ ("not-a-timestamp", None),
+ ],
+ )
+ def test_get_bigquery_job_datetime(self, value, expected):
+ assert (
+ self.operator._get_bigquery_job_datetime({"statistics":
{"startTime": value}}, "startTime")
+ == expected
+ )
def test_deduplicate_outputs(self):
outputs = [
