mobuchowski commented on code in PR #65788:
URL: https://github.com/apache/airflow/pull/65788#discussion_r3155453782
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -164,6 +173,144 @@ def get_job_name(task: TaskInstance |
RuntimeTaskInstance) -> str:
return f"{task.dag_id}.{task.task_id}"
+def get_task_instance_from_context():
+ """Return the task instance from the currently executing Airflow task
context."""
+ try:
+ from airflow.sdk import get_current_context
+ except (ImportError, AttributeError):
+ from airflow.operators.python import get_current_context # type:
ignore[no-redef]
+
+ task_instance = get_current_context().get("task_instance")
+ if task_instance is None:
+ raise RuntimeError(
+ "No Airflow task instance found in the current context. "
+ "Call this function from within a running Airflow task, "
+ "or pass `task_instance=` explicitly."
+ )
+ return task_instance
+
+
+def get_dag_run_dag_and_task_from_ti(task_instance):
Review Comment:
nice :)
##########
providers/openlineage/src/airflow/providers/openlineage/api/sql.py:
##########
@@ -0,0 +1,172 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Public helpers for emitting OpenLineage events describing SQL query
executions."""
+
+from __future__ import annotations
+
+import logging
+from collections import defaultdict
+from typing import TYPE_CHECKING
+
+from openlineage.client.facet_v2 import error_message_run, external_query_run,
sql_job
+
+from airflow.providers.openlineage.api.core import emit, is_openlineage_active
+from airflow.providers.openlineage.plugins.adapter import _PRODUCER
+from airflow.providers.openlineage.plugins.macros import lineage_job_name,
lineage_run_id
+from airflow.providers.openlineage.utils.sql_hook_lineage import (
+ _create_ol_event_pair,
+ _parse_query_into_datasets,
+)
+from airflow.providers.openlineage.utils.utils import
get_task_instance_from_context
+
+if TYPE_CHECKING:
+ from datetime import datetime
+
+ from openlineage.client.event_v2 import Dataset
+ from openlineage.client.facet_v2 import JobFacet, RunFacet
+
+ from airflow.models.taskinstance import TaskInstance
+ from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance
+
+log = logging.getLogger(__name__)
+
+__all__ = ["emit_query_lineage"]
+
+# Per-task counters so repeated emit_query_lineage calls within one task get
unique job names.
+_QUERY_COUNTERS: dict[str, int] = defaultdict(int)
+
+
+def emit_query_lineage(
+ *,
Review Comment:
Maybe not now, but I'd see need in the future to have deterministic run id
generation - maybe by passing callable where current one is the default?
##########
providers/openlineage/src/airflow/providers/openlineage/api/sql.py:
##########
@@ -0,0 +1,172 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Public helpers for emitting OpenLineage events describing SQL query
executions."""
+
+from __future__ import annotations
+
+import logging
+from collections import defaultdict
+from typing import TYPE_CHECKING
+
+from openlineage.client.facet_v2 import error_message_run, external_query_run,
sql_job
+
+from airflow.providers.openlineage.api.core import emit, is_openlineage_active
+from airflow.providers.openlineage.plugins.adapter import _PRODUCER
+from airflow.providers.openlineage.plugins.macros import lineage_job_name,
lineage_run_id
+from airflow.providers.openlineage.utils.sql_hook_lineage import (
+ _create_ol_event_pair,
+ _parse_query_into_datasets,
+)
+from airflow.providers.openlineage.utils.utils import
get_task_instance_from_context
+
+if TYPE_CHECKING:
+ from datetime import datetime
+
+ from openlineage.client.event_v2 import Dataset
+ from openlineage.client.facet_v2 import JobFacet, RunFacet
+
+ from airflow.models.taskinstance import TaskInstance
+ from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance
+
+log = logging.getLogger(__name__)
+
+__all__ = ["emit_query_lineage"]
+
+# Per-task counters so repeated emit_query_lineage calls within one task get
unique job names.
+_QUERY_COUNTERS: dict[str, int] = defaultdict(int)
+
+
+def emit_query_lineage(
+ *,
+ query_id: str | None = None,
+ query_source_namespace: str | None = None,
+ query_text: str | None = None,
+ inputs: list[Dataset] | None = None,
+ outputs: list[Dataset] | None = None,
+ start_time: datetime | None = None,
+ end_time: datetime | None = None,
+ is_successful: bool = True,
+ error_message: str | None = None,
+ default_database: str | None = None,
+ default_schema: str | None = None,
+ task_instance: TaskInstance | RuntimeTaskInstance | None = None,
+ additional_run_facets: dict[str, RunFacet] | None = None,
+ additional_job_facets: dict[str, JobFacet] | None = None,
+) -> None:
+ """
+ Emit a START + COMPLETE/FAIL OpenLineage event pair describing a single
SQL query execution.
+
+ The emitted events carry a ``parent`` run facet pointing at the currently
executing Airflow task
+ run. Any OpenLineage root information present on the task instance is
propagated to the
+ emitted events so the entire run hierarchy stays connected. This helper
can be called multiple
+ times within a single task; each call produces a distinct query event pair
identified by a
+ sequential job name suffix (``<dag_id>.<task_id>.query.<n>``).
+
+ :param query_id: Unique identifier of the query in the given
``query_source_namespace``. When both
+ ``query_id`` and ``query_source_namespace`` are provided, an
``externalQuery`` run facet is
+ attached to the emitted events.
+ :param query_source_namespace: OpenLineage namespace of the system that
ran the query, e.g.
+ ``"snowflake://org-acct"``,
``"databricks://adb-<id>.azuredatabricks.net"``, ``"bigquery"``.
+ :param query_text: Raw SQL query text. When provided, it is attached via a
``sql`` JobFacet and
+ ``inputs``/``outputs`` explicitly supplied are enriched with datasets
retrieved from query parsing.
+ :param inputs: Additional input datasets.
+ :param outputs: Additional output datasets.
+ :param start_time: Event time of the START event. Defaults to the current
UTC time.
+ :param end_time: Event time of the COMPLETE/FAIL event. Defaults to the
current UTC time.
+ :param is_successful: Whether the query completed successfully (COMPLETE)
or failed (FAIL).
+ :param error_message: Optional error message attached as an
``errorMessage`` run facet.
+ :param default_database: Default database for resolving unqualified tables
in ``query_text``.
+ :param default_schema: Default schema for resolving unqualified tables in
``query_text``.
+ :param task_instance: The Airflow task instance to attribute the query to.
Defaults to the
+ currently executing task instance obtained from the execution context.
+ :param additional_run_facets: Extra run facets to merge into the emitted
events.
+ :param additional_job_facets: Extra job facets to merge into the emitted
events.
+
+ :raises RuntimeError: If ``task_instance`` is not provided and cannot be
resolved from the current
+ execution context.
+
+ Example:
+
+ .. code-block:: python
+
+ from airflow.providers.openlineage.api import emit_query_lineage
+
+
+ @task
+ def my_task():
+ emit_query_lineage(
+ query_id="acde070d-8c4c-4f0d-9d8a-162843c10333",
+
query_source_namespace="databricks://adb-498971240325220.10.azuredatabricks.net",
+ query_text="SELECT * FROM analytics.public.users",
+ )
+ """
+ if not is_openlineage_active():
+ log.info("OpenLineage is not active - emit_query_lineage will have no
effect.")
+ return
+
+ if task_instance is None:
+ log.debug("TaskInstance not provided, retrieving it from context.")
+ task_instance = get_task_instance_from_context()
+
+ all_inputs = inputs or []
+ all_outputs = outputs or []
+ # Parse SQL into datasets and enrich inputs/outputs.
+ if query_text and query_source_namespace:
+ query_inputs, query_outputs = _parse_query_into_datasets(
+ query_text=query_text,
+ query_source_namespace=query_source_namespace,
+ default_database=default_database,
+ default_schema=default_schema,
+ )
+ all_inputs.extend(query_inputs)
+ all_outputs.extend(query_outputs)
Review Comment:
should we mutate the original inputs/outputs lists?
##########
providers/openlineage/src/airflow/providers/openlineage/api/sql.py:
##########
@@ -0,0 +1,172 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Public helpers for emitting OpenLineage events describing SQL query
executions."""
+
+from __future__ import annotations
+
+import logging
+from collections import defaultdict
+from typing import TYPE_CHECKING
+
+from openlineage.client.facet_v2 import error_message_run, external_query_run,
sql_job
+
+from airflow.providers.openlineage.api.core import emit, is_openlineage_active
+from airflow.providers.openlineage.plugins.adapter import _PRODUCER
+from airflow.providers.openlineage.plugins.macros import lineage_job_name,
lineage_run_id
+from airflow.providers.openlineage.utils.sql_hook_lineage import (
+ _create_ol_event_pair,
+ _parse_query_into_datasets,
+)
+from airflow.providers.openlineage.utils.utils import
get_task_instance_from_context
+
+if TYPE_CHECKING:
+ from datetime import datetime
+
+ from openlineage.client.event_v2 import Dataset
+ from openlineage.client.facet_v2 import JobFacet, RunFacet
+
+ from airflow.models.taskinstance import TaskInstance
+ from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance
+
+log = logging.getLogger(__name__)
+
+__all__ = ["emit_query_lineage"]
+
+# Per-task counters so repeated emit_query_lineage calls within one task get
unique job names.
+_QUERY_COUNTERS: dict[str, int] = defaultdict(int)
+
+
+def emit_query_lineage(
+ *,
+ query_id: str | None = None,
+ query_source_namespace: str | None = None,
+ query_text: str | None = None,
+ inputs: list[Dataset] | None = None,
+ outputs: list[Dataset] | None = None,
+ start_time: datetime | None = None,
+ end_time: datetime | None = None,
+ is_successful: bool = True,
+ error_message: str | None = None,
+ default_database: str | None = None,
+ default_schema: str | None = None,
+ task_instance: TaskInstance | RuntimeTaskInstance | None = None,
+ additional_run_facets: dict[str, RunFacet] | None = None,
+ additional_job_facets: dict[str, JobFacet] | None = None,
+) -> None:
+ """
+ Emit a START + COMPLETE/FAIL OpenLineage event pair describing a single
SQL query execution.
+
+ The emitted events carry a ``parent`` run facet pointing at the currently
executing Airflow task
+ run. Any OpenLineage root information present on the task instance is
propagated to the
+ emitted events so the entire run hierarchy stays connected. This helper
can be called multiple
+ times within a single task; each call produces a distinct query event pair
identified by a
+ sequential job name suffix (``<dag_id>.<task_id>.query.<n>``).
+
+ :param query_id: Unique identifier of the query in the given
``query_source_namespace``. When both
+ ``query_id`` and ``query_source_namespace`` are provided, an
``externalQuery`` run facet is
+ attached to the emitted events.
+ :param query_source_namespace: OpenLineage namespace of the system that
ran the query, e.g.
+ ``"snowflake://org-acct"``,
``"databricks://adb-<id>.azuredatabricks.net"``, ``"bigquery"``.
+ :param query_text: Raw SQL query text. When provided, it is attached via a
``sql`` JobFacet and
+ ``inputs``/``outputs`` explicitly supplied are enriched with datasets
retrieved from query parsing.
+ :param inputs: Additional input datasets.
+ :param outputs: Additional output datasets.
+ :param start_time: Event time of the START event. Defaults to the current
UTC time.
+ :param end_time: Event time of the COMPLETE/FAIL event. Defaults to the
current UTC time.
+ :param is_successful: Whether the query completed successfully (COMPLETE)
or failed (FAIL).
+ :param error_message: Optional error message attached as an
``errorMessage`` run facet.
+ :param default_database: Default database for resolving unqualified tables
in ``query_text``.
+ :param default_schema: Default schema for resolving unqualified tables in
``query_text``.
+ :param task_instance: The Airflow task instance to attribute the query to.
Defaults to the
+ currently executing task instance obtained from the execution context.
+ :param additional_run_facets: Extra run facets to merge into the emitted
events.
+ :param additional_job_facets: Extra job facets to merge into the emitted
events.
+
+ :raises RuntimeError: If ``task_instance`` is not provided and cannot be
resolved from the current
+ execution context.
+
+ Example:
+
+ .. code-block:: python
+
+ from airflow.providers.openlineage.api import emit_query_lineage
+
+
+ @task
+ def my_task():
+ emit_query_lineage(
+ query_id="acde070d-8c4c-4f0d-9d8a-162843c10333",
+
query_source_namespace="databricks://adb-498971240325220.10.azuredatabricks.net",
+ query_text="SELECT * FROM analytics.public.users",
+ )
+ """
+ if not is_openlineage_active():
+ log.info("OpenLineage is not active - emit_query_lineage will have no
effect.")
+ return
+
+ if task_instance is None:
+ log.debug("TaskInstance not provided, retrieving it from context.")
+ task_instance = get_task_instance_from_context()
+
+ all_inputs = inputs or []
+ all_outputs = outputs or []
+ # Parse SQL into datasets and enrich inputs/outputs.
+ if query_text and query_source_namespace:
+ query_inputs, query_outputs = _parse_query_into_datasets(
+ query_text=query_text,
+ query_source_namespace=query_source_namespace,
+ default_database=default_database,
+ default_schema=default_schema,
+ )
+ all_inputs.extend(query_inputs)
+ all_outputs.extend(query_outputs)
+
+ # Run facets: user-provided first, internal facets overlaid so internals
always win.
+ run_facets: dict[str, RunFacet] = {**(additional_run_facets or {})}
+ if query_id and query_source_namespace:
+ run_facets["externalQuery"] = external_query_run.ExternalQueryRunFacet(
+ externalQueryId=query_id, source=query_source_namespace,
producer=_PRODUCER
+ )
+ if error_message:
+ run_facets["errorMessage"] = error_message_run.ErrorMessageRunFacet(
+ message=error_message, programmingLanguage="SQL",
producer=_PRODUCER
+ )
+
+ # Job facets: user-provided first, internal facets overlaid so internals
always win.
+ job_facets: dict[str, JobFacet] = {**(additional_job_facets or {})}
+ if query_text:
+ job_facets["sql"] = sql_job.SQLJobFacet(query=query_text,
producer=_PRODUCER)
+
+ ti_key = lineage_run_id(task_instance)
+ _QUERY_COUNTERS[ti_key] += 1
+ job_name =
f"{lineage_job_name(task_instance)}.query.{_QUERY_COUNTERS[ti_key]}"
Review Comment:
Wouldn't this naming scheme collide with sql hook lineage if someone somehow
uses both at the same time?
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -164,6 +173,144 @@ def get_job_name(task: TaskInstance |
RuntimeTaskInstance) -> str:
return f"{task.dag_id}.{task.task_id}"
+def get_task_instance_from_context():
+ """Return the task instance from the currently executing Airflow task
context."""
+ try:
+ from airflow.sdk import get_current_context
+ except (ImportError, AttributeError):
+ from airflow.operators.python import get_current_context # type:
ignore[no-redef]
+
+ task_instance = get_current_context().get("task_instance")
+ if task_instance is None:
+ raise RuntimeError(
+ "No Airflow task instance found in the current context. "
+ "Call this function from within a running Airflow task, "
+ "or pass `task_instance=` explicitly."
+ )
+ return task_instance
+
+
+def get_dag_run_dag_and_task_from_ti(task_instance):
+ """
+ Return a ``(dag_run, dag, task)`` tuple for the given task instance.
+
+ On Airflow 3+, ``task_instance.get_template_context()`` exposes all three
objects.
+ On Airflow 2, we fall back to the Ti attributes.
+ """
+ if AIRFLOW_V_3_0_PLUS:
+ context = task_instance.get_template_context()
+ return context["dag_run"], context["dag"], context["task"]
+ task = task_instance.task
+ return task_instance.dag_run, task.dag, task
+
+
+def build_task_event_run_facets(
Review Comment:
Maybe this is something we need in common with rest of OL provider - we're
doing pretty much the same facet construction in adapter.py `_build_job`
##########
providers/openlineage/src/airflow/providers/openlineage/api/datasets.py:
##########
@@ -0,0 +1,155 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Public helpers for emitting OpenLineage events carrying dataset lineage."""
+
+from __future__ import annotations
+
+import logging
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING
+
+from openlineage.client.event_v2 import Dataset, Job, Run, RunEvent, RunState
+
+from airflow.providers.openlineage.api.core import emit, is_openlineage_active
+from airflow.providers.openlineage.plugins.adapter import _PRODUCER,
OpenLineageAdapter
+from airflow.providers.openlineage.plugins.macros import (
+ _get_dag_run_clear_number,
+ _get_dag_run_conf,
+ _get_logical_date,
+ lineage_job_name,
+ lineage_job_namespace,
+ lineage_run_id,
+)
+from airflow.providers.openlineage.utils.utils import (
+ build_task_event_job_facets,
+ build_task_event_run_facets,
+ get_dag_run_dag_and_task_from_ti,
+ get_task_instance_from_context,
+)
+from airflow.utils.state import TaskInstanceState
+
+if TYPE_CHECKING:
+ from openlineage.client.event_v2 import InputDataset, OutputDataset
+ from openlineage.client.facet_v2 import JobFacet, RunFacet
+
+ from airflow.models.taskinstance import TaskInstance
+ from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance
+
+log = logging.getLogger(__name__)
+
+__all__ = ["emit_dataset_lineage"]
+
+
+def emit_dataset_lineage(
Review Comment:
We should carefully consider if `emit_dataset_lineage` and
`emit_query_lineage` should throw exceptions.
On one hand, we are very careful about breaking user tasks - we should not
do that _automatically_ - but, on the other hand this is explicit user action -
maybe this is okay? However, I'd make sure we're not raising whole bunch of
different exceptions in that case, but are building some singular one that
users can catch easily.
##########
providers/openlineage/src/airflow/providers/openlineage/api/sql.py:
##########
@@ -0,0 +1,172 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Public helpers for emitting OpenLineage events describing SQL query
executions."""
+
+from __future__ import annotations
+
+import logging
+from collections import defaultdict
+from typing import TYPE_CHECKING
+
+from openlineage.client.facet_v2 import error_message_run, external_query_run,
sql_job
+
+from airflow.providers.openlineage.api.core import emit, is_openlineage_active
+from airflow.providers.openlineage.plugins.adapter import _PRODUCER
+from airflow.providers.openlineage.plugins.macros import lineage_job_name,
lineage_run_id
+from airflow.providers.openlineage.utils.sql_hook_lineage import (
+ _create_ol_event_pair,
+ _parse_query_into_datasets,
+)
+from airflow.providers.openlineage.utils.utils import
get_task_instance_from_context
+
+if TYPE_CHECKING:
+ from datetime import datetime
+
+ from openlineage.client.event_v2 import Dataset
+ from openlineage.client.facet_v2 import JobFacet, RunFacet
+
+ from airflow.models.taskinstance import TaskInstance
+ from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance
+
+log = logging.getLogger(__name__)
+
+__all__ = ["emit_query_lineage"]
+
+# Per-task counters so repeated emit_query_lineage calls within one task get
unique job names.
+_QUERY_COUNTERS: dict[str, int] = defaultdict(int)
Review Comment:
are we sure this will be correctly rebuild for each task?
--
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]
