From 9984648bd7c1a513848b5722e5620d7e2373d383 Mon Sep 17 00:00:00 2001 From: nailo2c Date: Thu, 2 Jul 2026 14:15:15 +0800 Subject: [PATCH 1/7] Emit per-statement OpenLineage events for BigQuery script jobs --- providers/google/README.rst | 2 +- providers/google/pyproject.toml | 2 +- .../google/cloud/openlineage/mixins.py | 104 +++++- .../google/cloud/openlineage/test_mixins.py | 295 +++++++++++++++++- uv.lock | 4 +- 5 files changed, 395 insertions(+), 12 deletions(-) diff --git a/providers/google/README.rst b/providers/google/README.rst index 3a986ff3462b9..a8a01c75e27d9 100644 --- a/providers/google/README.rst +++ b/providers/google/README.rst @@ -193,7 +193,7 @@ Extra Dependencies ``microsoft.mssql`` ``apache-airflow-providers-microsoft-mssql`` ``mongo`` ``apache-airflow-providers-mongo`` ``mysql`` ``apache-airflow-providers-mysql`` -``openlineage`` ``apache-airflow-providers-openlineage`` +``openlineage`` ``apache-airflow-providers-openlineage>=2.16.0`` ``postgres`` ``apache-airflow-providers-postgres`` ``presto`` ``apache-airflow-providers-presto`` ``salesforce`` ``apache-airflow-providers-salesforce`` diff --git a/providers/google/pyproject.toml b/providers/google/pyproject.toml index 0bb1067364fdb..7bd6c88e48361 100644 --- a/providers/google/pyproject.toml +++ b/providers/google/pyproject.toml @@ -189,7 +189,7 @@ dependencies = [ "apache-airflow-providers-mysql" ] "openlineage" = [ - "apache-airflow-providers-openlineage" + "apache-airflow-providers-openlineage>=2.16.0" ] "postgres" = [ "apache-airflow-providers-postgres" 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 539c24f65332f..9f803486df748 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 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 @@ 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,50 @@ def get_openlineage_facets_on_complete(self, _): 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_to_emit = [] + for child_job in self._client.list_jobs(parent_job=self.job_id): + # BigQuery returns Job objects; keep raw job IDs supported for lightweight clients. + child_job_id = getattr(child_job, "job_id", child_job) + try: + 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) + except Exception as child_exception: + self.log.warning( + "Cannot retrieve lineage for BigQuery child job `%s`. %s", + child_job_id, + child_exception, + exc_info=True, + ) + continue inputs.extend(child_inputs) outputs.extend(child_outputs) + child_jobs_to_emit.append( + ( + str(child_job_id), + child_job_properties, + child_inputs, + child_outputs, + ) + ) + for child_index, ( + child_job_id, + child_job_properties, + child_inputs, + child_outputs, + ) in enumerate(sorted(child_jobs_to_emit, key=self._get_child_job_sort_key), start=1): + self._emit_child_query_lineage( + task_instance=task_instance, + child_index=child_index, + child_job_id=child_job_id, + child_job_properties=child_job_properties, + inputs=child_inputs, + outputs=child_outputs, + ) else: inputs, outputs = self._get_inputs_and_outputs(job_properties) @@ -139,6 +176,61 @@ def get_openlineage_facets_on_complete(self, _): 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_id: str, + child_job_properties: dict, + inputs: list[InputDataset], + outputs: list[OutputDataset], + ) -> None: + if task_instance is None: + self.log.debug("No task instance available. Skipping BigQuery child job OpenLineage event.") + return + + from airflow.providers.openlineage.api.sql import emit_query_lineage + 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 + start_time = self._get_bigquery_job_datetime(child_job_properties, "startTime") + end_time = self._get_bigquery_job_datetime(child_job_properties, "endTime") + emit_query_lineage( + query_id=child_job_id, + query_source_namespace=BIGQUERY_NAMESPACE, + inputs=inputs, + outputs=outputs, + start_time=start_time, + end_time=end_time, + 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, + ) + + @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, child_job) -> tuple[datetime, str]: + # Emit children ordered by execution time so the query.N suffix is stable across runs; + # children missing a startTime sort last, then break ties deterministically by job id. + child_job_id, child_job_properties, _, _ = child_job + start_time = cls._get_bigquery_job_datetime(child_job_properties, "startTime") + return ( + start_time or datetime.max.replace(tzinfo=timezone.utc), + child_job_id, + ) + 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 e7204368fd500..faab51e866196 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 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 @@ def test_get_openlineage_facets_on_complete_extract_job(self): ), ] - 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), @@ -405,8 +433,9 @@ def test_get_openlineage_facets_on_complete_script_job(self): Table.from_api_repr(read_common_json_file("out_table_details.json")), ] self.client.list_jobs.return_value = ["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,268 @@ def test_get_openlineage_facets_on_complete_script_job(self): }, ), ] + mock_emit_query_lineage.assert_called_once() + assert mock_emit_query_lineage.call_args.kwargs["query_id"] == "child_job_id" + assert mock_emit_query_lineage.call_args.kwargs["query_source_namespace"] == "bigquery" + 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 = { + "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 = { + "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 = ["child_job_1", "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 = { + "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 = { + "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()) + + assert lineage.inputs == [input_table2, input_table1] + assert lineage.outputs == [output_table2, output_table1] + 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" + 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={"configuration": {"jobType": "QUERY"}, "status": {"state": "DONE"}}), + MagicMock(_properties=copy.deepcopy(self.query_job_details)), + ] + self.client.list_jobs.return_value = ["child_job_1", "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] + + @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 = ["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() + + @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_id="child_job_id", + 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( + ("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 = [ diff --git a/uv.lock b/uv.lock index 43b6db586d290..d28d3c282d11a 100644 --- a/uv.lock +++ b/uv.lock @@ -5599,7 +5599,7 @@ mysql = [ { name = "apache-airflow-providers-mysql" }, ] openlineage = [ - { name = "apache-airflow-providers-openlineage" }, + { name = "apache-airflow-providers-openlineage", specifier = ">=2.16.0" }, ] oracle = [ { name = "apache-airflow-providers-oracle" }, @@ -5674,7 +5674,7 @@ requires-dist = [ { name = "apache-airflow-providers-microsoft-mssql", marker = "extra == 'microsoft-mssql'", editable = "providers/microsoft/mssql" }, { name = "apache-airflow-providers-mongo", marker = "extra == 'mongo'", editable = "providers/mongo" }, { name = "apache-airflow-providers-mysql", marker = "extra == 'mysql'", editable = "providers/mysql" }, - { name = "apache-airflow-providers-openlineage", marker = "extra == 'openlineage'", editable = "providers/openlineage" }, + { name = "apache-airflow-providers-openlineage", marker = "extra == 'openlineage'", editable = "providers/openlineage", specifier = ">=2.16.0" }, { name = "apache-airflow-providers-oracle", marker = "extra == 'oracle'", editable = "providers/oracle" }, { name = "apache-airflow-providers-postgres", marker = "extra == 'postgres'", editable = "providers/postgres" }, { name = "apache-airflow-providers-presto", marker = "extra == 'presto'", editable = "providers/presto" }, From 80a4282f09fad50bd815f1ba0c5016479b1e8f9e Mon Sep 17 00:00:00 2001 From: nailo2c Date: Thu, 2 Jul 2026 16:39:21 +0800 Subject: [PATCH 2/7] Fix CI MyPy type error --- .../airflow/providers/google/cloud/openlineage/mixins.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 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 9f803486df748..ec6fced9925ea 100644 --- a/providers/google/src/airflow/providers/google/cloud/openlineage/mixins.py +++ b/providers/google/src/airflow/providers/google/cloud/openlineage/mixins.py @@ -183,11 +183,11 @@ def _emit_child_query_lineage( child_index: int, child_job_id: str, child_job_properties: dict, - inputs: list[InputDataset], - outputs: list[OutputDataset], + 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.") + self.log.debug("No task instance available. Skipping BigQuery child job OpenLineage event.") # type: ignore[attr-defined] return from airflow.providers.openlineage.api.sql import emit_query_lineage @@ -207,7 +207,7 @@ def _emit_child_query_lineage( 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, + additional_job_facets=job_facets, # type: ignore[arg-type] ) @staticmethod From 8f56ffcb8f5001f8791cfcd62a6362c18fc37161 Mon Sep 17 00:00:00 2001 From: nailo2c Date: Thu, 2 Jul 2026 17:50:42 +0800 Subject: [PATCH 3/7] Modify the test to match the convention --- .../google/tests/unit/google/cloud/openlineage/test_mixins.py | 3 +++ 1 file changed, 3 insertions(+) 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 faab51e866196..c0a4fcdd4a0d5 100644 --- a/providers/google/tests/unit/google/cloud/openlineage/test_mixins.py +++ b/providers/google/tests/unit/google/cloud/openlineage/test_mixins.py @@ -717,6 +717,9 @@ def test_script_job_without_task_instance_does_not_emit_child_query_events(self, 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): From 195932c3b448d4c7b4e44998d6e4a30695bf5c2f Mon Sep 17 00:00:00 2001 From: nailo2c Date: Fri, 10 Jul 2026 15:54:45 +0800 Subject: [PATCH 4/7] Remove OpenLineage version pin and clarify child event emission --- providers/google/README.rst | 2 +- providers/google/pyproject.toml | 2 +- .../providers/google/cloud/openlineage/mixins.py | 10 ++++++++-- .../tests/unit/google/cloud/openlineage/test_mixins.py | 1 + uv.lock | 4 ++-- 5 files changed, 13 insertions(+), 6 deletions(-) diff --git a/providers/google/README.rst b/providers/google/README.rst index a8a01c75e27d9..3a986ff3462b9 100644 --- a/providers/google/README.rst +++ b/providers/google/README.rst @@ -193,7 +193,7 @@ Extra Dependencies ``microsoft.mssql`` ``apache-airflow-providers-microsoft-mssql`` ``mongo`` ``apache-airflow-providers-mongo`` ``mysql`` ``apache-airflow-providers-mysql`` -``openlineage`` ``apache-airflow-providers-openlineage>=2.16.0`` +``openlineage`` ``apache-airflow-providers-openlineage`` ``postgres`` ``apache-airflow-providers-postgres`` ``presto`` ``apache-airflow-providers-presto`` ``salesforce`` ``apache-airflow-providers-salesforce`` diff --git a/providers/google/pyproject.toml b/providers/google/pyproject.toml index 7bd6c88e48361..0bb1067364fdb 100644 --- a/providers/google/pyproject.toml +++ b/providers/google/pyproject.toml @@ -189,7 +189,7 @@ dependencies = [ "apache-airflow-providers-mysql" ] "openlineage" = [ - "apache-airflow-providers-openlineage>=2.16.0" + "apache-airflow-providers-openlineage" ] "postgres" = [ "apache-airflow-providers-postgres" 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 ec6fced9925ea..d57464976771d 100644 --- a/providers/google/src/airflow/providers/google/cloud/openlineage/mixins.py +++ b/providers/google/src/airflow/providers/google/cloud/openlineage/mixins.py @@ -200,6 +200,10 @@ def _emit_child_query_lineage( emit_query_lineage( query_id=child_job_id, 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, @@ -222,11 +226,13 @@ def _get_bigquery_job_datetime(properties: dict, field_name: str) -> datetime | @classmethod def _get_child_job_sort_key(cls, child_job) -> tuple[datetime, str]: - # Emit children ordered by execution time so the query.N suffix is stable across runs; - # children missing a startTime sort last, then break ties deterministically by job id. + # Emit children ordered by execution time so the query.N suffix is stable across runs, + # breaking ties deterministically by job id. child_job_id, child_job_properties, _, _ = child_job start_time = cls._get_bigquery_job_datetime(child_job_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), child_job_id, ) 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 c0a4fcdd4a0d5..50957383b807f 100644 --- a/providers/google/tests/unit/google/cloud/openlineage/test_mixins.py +++ b/providers/google/tests/unit/google/cloud/openlineage/test_mixins.py @@ -488,6 +488,7 @@ def test_get_openlineage_facets_on_complete_script_job(self, mock_emit_query_lin mock_emit_query_lineage.assert_called_once() assert mock_emit_query_lineage.call_args.kwargs["query_id"] == "child_job_id" 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["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( diff --git a/uv.lock b/uv.lock index d28d3c282d11a..43b6db586d290 100644 --- a/uv.lock +++ b/uv.lock @@ -5599,7 +5599,7 @@ mysql = [ { name = "apache-airflow-providers-mysql" }, ] openlineage = [ - { name = "apache-airflow-providers-openlineage", specifier = ">=2.16.0" }, + { name = "apache-airflow-providers-openlineage" }, ] oracle = [ { name = "apache-airflow-providers-oracle" }, @@ -5674,7 +5674,7 @@ requires-dist = [ { name = "apache-airflow-providers-microsoft-mssql", marker = "extra == 'microsoft-mssql'", editable = "providers/microsoft/mssql" }, { name = "apache-airflow-providers-mongo", marker = "extra == 'mongo'", editable = "providers/mongo" }, { name = "apache-airflow-providers-mysql", marker = "extra == 'mysql'", editable = "providers/mysql" }, - { name = "apache-airflow-providers-openlineage", marker = "extra == 'openlineage'", editable = "providers/openlineage", specifier = ">=2.16.0" }, + { name = "apache-airflow-providers-openlineage", marker = "extra == 'openlineage'", editable = "providers/openlineage" }, { name = "apache-airflow-providers-oracle", marker = "extra == 'oracle'", editable = "providers/oracle" }, { name = "apache-airflow-providers-postgres", marker = "extra == 'postgres'", editable = "providers/postgres" }, { name = "apache-airflow-providers-presto", marker = "extra == 'presto'", editable = "providers/presto" }, From e80b41a4e55fb30fb92ab58ba6805ecc2ada2284 Mon Sep 17 00:00:00 2001 From: nailo2c Date: Fri, 10 Jul 2026 18:22:36 +0800 Subject: [PATCH 5/7] Handle older OpenLineage providers and failed BigQuery child jobs --- .../google/cloud/openlineage/mixins.py | 60 +++++++------ .../google/cloud/openlineage/test_mixins.py | 86 ++++++++++++++++--- 2 files changed, 108 insertions(+), 38 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 d57464976771d..d39a7f0080c24 100644 --- a/providers/google/src/airflow/providers/google/cloud/openlineage/mixins.py +++ b/providers/google/src/airflow/providers/google/cloud/openlineage/mixins.py @@ -115,41 +115,36 @@ def get_openlineage_facets_on_complete(self, task_instance): # 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 - child_jobs_to_emit = [] + child_jobs_properties = [] for child_job in self._client.list_jobs(parent_job=self.job_id): - # BigQuery returns Job objects; keep raw job IDs supported for lightweight clients. - child_job_id = getattr(child_job, "job_id", child_job) try: - child_job_properties = self._client.get_job(job_id=child_job_id)._properties + 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 retrieve lineage for BigQuery child job `%s`. %s", - child_job_id, + "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) - child_jobs_to_emit.append( - ( - str(child_job_id), - child_job_properties, - child_inputs, - child_outputs, - ) - ) - for child_index, ( - child_job_id, - child_job_properties, - child_inputs, - child_outputs, - ) in enumerate(sorted(child_jobs_to_emit, key=self._get_child_job_sort_key), start=1): self._emit_child_query_lineage( task_instance=task_instance, child_index=child_index, - child_job_id=child_job_id, child_job_properties=child_job_properties, inputs=child_inputs, outputs=child_outputs, @@ -181,7 +176,6 @@ def _emit_child_query_lineage( *, task_instance, child_index: int, - child_job_id: str, child_job_properties: dict, inputs: list[InputDataset | Dataset], outputs: list[OutputDataset | Dataset], @@ -190,15 +184,24 @@ def _emit_child_query_lineage( self.log.debug("No task instance available. Skipping BigQuery child job OpenLineage event.") # type: ignore[attr-defined] return - from airflow.providers.openlineage.api.sql import emit_query_lineage + 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") emit_query_lineage( - query_id=child_job_id, + 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 @@ -208,6 +211,8 @@ def _emit_child_query_lineage( 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)}, @@ -225,16 +230,15 @@ def _get_bigquery_job_datetime(properties: dict, field_name: str) -> datetime | return None @classmethod - def _get_child_job_sort_key(cls, child_job) -> tuple[datetime, str]: + 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. - child_job_id, child_job_properties, _, _ = child_job - start_time = cls._get_bigquery_job_datetime(child_job_properties, "startTime") + 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), - child_job_id, + get_from_nullable_chain(properties, ["jobReference", "jobId"]) or "", ) def _get_inputs_and_outputs(self, properties: dict) -> tuple[list[InputDataset], list[OutputDataset]]: 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 50957383b807f..29d53108c9871 100644 --- a/providers/google/tests/unit/google/cloud/openlineage/test_mixins.py +++ b/providers/google/tests/unit/google/cloud/openlineage/test_mixins.py @@ -432,7 +432,7 @@ def test_get_openlineage_facets_on_complete_script_job(self, mock_emit_query_lin 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(mock_ti) @@ -486,9 +486,14 @@ def test_get_openlineage_facets_on_complete_script_job(self, mock_emit_query_lin ), ] mock_emit_query_lineage.assert_called_once() - assert mock_emit_query_lineage.call_args.kwargs["query_id"] == "child_job_id" + 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( @@ -510,6 +515,7 @@ def test_script_job_aggregates_parent_datasets_and_emits_child_query_lineage( 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"}, @@ -518,6 +524,7 @@ def test_script_job_aggregates_parent_datasets_and_emits_child_query_lineage( "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"}, @@ -534,7 +541,10 @@ def test_script_job_aggregates_parent_datasets_and_emits_child_query_lineage( MagicMock(_properties=child_job_1_details), MagicMock(_properties=child_job_2_details), ] - self.client.list_jobs.return_value = ["child_job_1", "child_job_2"] + 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): @@ -586,6 +596,7 @@ def test_script_job_builds_child_query_events( 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"}, @@ -598,6 +609,7 @@ def test_script_job_builds_child_query_events( "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"}, @@ -637,8 +649,9 @@ def get_inputs_and_outputs(_, properties): lineage = self.operator.get_openlineage_facets_on_complete(make_task_instance()) - assert lineage.inputs == [input_table2, input_table1] - assert lineage.outputs == [output_table2, output_table1] + # 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 = [ @@ -677,14 +690,29 @@ def test_script_job_continues_after_child_lineage_failure( ): 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={"configuration": {"jobType": "QUERY"}, "status": {"state": "DONE"}}), - MagicMock(_properties=copy.deepcopy(self.query_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"), ] - self.client.list_jobs.return_value = ["child_job_1", "child_job_2"] mock_get_inputs_and_outputs.side_effect = [ RuntimeError("broken child"), ([input_table2], [output_table2]), @@ -699,6 +727,8 @@ def test_script_job_continues_after_child_lineage_failure( 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): @@ -710,7 +740,7 @@ def test_script_job_without_task_instance_does_not_emit_child_query_events(self, 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")] lineage = self.operator.get_openlineage_facets_on_complete(None) @@ -727,7 +757,6 @@ def test_child_query_lineage_without_query_omits_sql_job_facet(self, mock_emit_q self.operator._emit_child_query_lineage( task_instance=make_task_instance(), child_index=1, - child_job_id="child_job_id", child_job_properties={ "configuration": {"jobType": "QUERY", "query": {}}, "statistics": {"query": {"cacheHit": False, "totalBytesBilled": "10"}}, @@ -738,6 +767,43 @@ def test_child_query_lineage_without_query_omits_sql_job_facet(self, mock_emit_q assert mock_emit_query_lineage.call_args.kwargs["additional_job_facets"] 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"), [ From f6ecce71c941751aceb7748a884d037a28dce381 Mon Sep 17 00:00:00 2001 From: nailo2c Date: Wed, 15 Jul 2026 16:57:36 +0800 Subject: [PATCH 6/7] Add handling for partial timestamps in child query lineage emission --- .../google/cloud/openlineage/mixins.py | 8 +++++-- .../google/cloud/openlineage/test_mixins.py | 23 +++++++++++++++++++ 2 files changed, 29 insertions(+), 2 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 d39a7f0080c24..3864c3a89da73 100644 --- a/providers/google/src/airflow/providers/google/cloud/openlineage/mixins.py +++ b/providers/google/src/airflow/providers/google/cloud/openlineage/mixins.py @@ -200,6 +200,11 @@ def _emit_child_query_lineage( 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") + timestamp_kwargs = ( + {"start_time": start_time, "end_time": end_time} + if start_time is not None and end_time is not None + else {} + ) emit_query_lineage( query_id=get_from_nullable_chain(child_job_properties, ["jobReference", "jobId"]), query_source_namespace=BIGQUERY_NAMESPACE, @@ -209,8 +214,7 @@ def _emit_child_query_lineage( query_text=None, inputs=inputs, outputs=outputs, - start_time=start_time, - end_time=end_time, + **timestamp_kwargs, is_successful=error_result is None, error_message=error_result.get("message") if error_result else None, task_instance=task_instance, 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 29d53108c9871..5a243b1e3fa8d 100644 --- a/providers/google/tests/unit/google/cloud/openlineage/test_mixins.py +++ b/providers/google/tests/unit/google/cloud/openlineage/test_mixins.py @@ -767,6 +767,29 @@ def test_child_query_lineage_without_query_omits_sql_job_facet(self, mock_emit_q 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 "start_time" not in mock_emit_query_lineage.call_args.kwargs + assert "end_time" not in mock_emit_query_lineage.call_args.kwargs + @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( From 0de556878c9f399d1db3628aa7609fc3480bbffa Mon Sep 17 00:00:00 2001 From: nailo2c Date: Thu, 16 Jul 2026 10:52:25 +0800 Subject: [PATCH 7/7] Refactor timestamp handling in BigQuery lineage emission to simplify logic and ensure None values are correctly passed --- .../providers/google/cloud/openlineage/mixins.py | 10 ++++------ .../tests/unit/google/cloud/openlineage/test_mixins.py | 4 ++-- 2 files changed, 6 insertions(+), 8 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 3864c3a89da73..3e5520721ed97 100644 --- a/providers/google/src/airflow/providers/google/cloud/openlineage/mixins.py +++ b/providers/google/src/airflow/providers/google/cloud/openlineage/mixins.py @@ -200,11 +200,8 @@ def _emit_child_query_lineage( 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") - timestamp_kwargs = ( - {"start_time": start_time, "end_time": end_time} - if start_time is not None and end_time is not None - else {} - ) + 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, @@ -214,7 +211,8 @@ def _emit_child_query_lineage( query_text=None, inputs=inputs, outputs=outputs, - **timestamp_kwargs, + 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, 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 5a243b1e3fa8d..0106eb81eb3c8 100644 --- a/providers/google/tests/unit/google/cloud/openlineage/test_mixins.py +++ b/providers/google/tests/unit/google/cloud/openlineage/test_mixins.py @@ -787,8 +787,8 @@ def test_child_query_lineage_omits_partial_timestamps(self, mock_emit_query_line outputs=[], ) - assert "start_time" not in mock_emit_query_lineage.call_args.kwargs - assert "end_time" not in mock_emit_query_lineage.call_args.kwargs + 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):