nailo2c commented on code in PR #69234:
URL: https://github.com/apache/airflow/pull/69234#discussion_r3556817336
##########
providers/google/src/airflow/providers/google/cloud/openlineage/mixins.py:
##########
@@ -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 | 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
+
+ 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, # 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, 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),
Review Comment:
Okay, let me add it.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
