kacpermuda commented on code in PR #69234:
URL: https://github.com/apache/airflow/pull/69234#discussion_r3549668610
##########
providers/google/src/airflow/providers/google/cloud/openlineage/mixins.py:
##########
@@ -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)
Review Comment:
Did this loop not work before? We iterated over list_jobs result as it was
job_ids, not job_instances. I think I've tested this when implementing, was I
wrong or the client changed over time?
##########
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
Review Comment:
As mentioned above, this will not work with older OL provider (api was added
recently), so we need to try/except and early return.
##########
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:
Let's add an inline comment here why we're using max, and not f.e. None
##########
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
Review Comment:
Can we also add a comment on why we're not passing this sql as query_text to
the `emit_query_lineage` and that this is intentional? Also let's pass
query_text=None explicitly to the function, since we don't want to use sql
parsing capability.
##########
providers/google/src/airflow/providers/google/cloud/openlineage/mixins.py:
##########
@@ -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),
Review Comment:
Can't we just use a dict here with job_id as key?
##########
providers/google/src/airflow/providers/google/cloud/openlineage/mixins.py:
##########
@@ -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(
Review Comment:
Let's move this whole piece with emitting child events to a separate method
maybe, but for sure let's guard it with the try/except on "from
airflow.providers.openlineage.api.sql import emit_query_lineage" as this code
is not available in older OL clients. This way, we keep the older behavior
intact for all, and for those with newer OL clients, we offer this new
functionality.
##########
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(
Review Comment:
Just curious, do we have error from bigquery if the job failed? We could
also pass it here for each sql. Also, if the script stopped, f.e. only first of
all 10 sqls executed, do we know about it ?
##########
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
Review Comment:
This is very hard to debug and read (this mysterious list of tuples), let's
not pass what we do not need to this function or figure out a more readable way
to store this child-job info.
##########
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"
Review Comment:
We should probably not add this version constraint, since the new feature is
an addition to current functionality, not breaking the entire Ol integration.
--
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]
