# 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.
from __future__ import annotations
import datetime
import json
import logging
from collections.abc import Callable, Collection, Mapping
from contextlib import suppress
from decimal import Decimal
from enum import Enum
from functools import wraps
from importlib import metadata
from typing import TYPE_CHECKING, Any, Literal
from uuid import UUID
import attrs
from openlineage.client.facet_v2 import (
documentation_job,
job_dependencies_run,
job_type_job,
nominal_time_run,
ownership_job,
parent_run,
tags_job,
)
from openlineage.client.utils import RedactMixin
from openlineage.client.uuid import generate_static_uuid
from sqlalchemy.orm.exc import DetachedInstanceError
from airflow import __version__ as AIRFLOW_VERSION
from airflow.models import DagRun, TaskInstance, TaskReschedule
from airflow.providers.common.compat.assets import Asset, AssetAlias
from airflow.providers.common.compat.module_loading import import_string
from airflow.providers.common.compat.sdk import (
DAG,
AirflowOptionalProviderFeatureException,
BaseOperator,
BaseSensorOperator,
MappedOperator,
conf as airflow_conf,
)
from airflow.providers.openlineage import (
__version__ as OPENLINEAGE_PROVIDER_VERSION,
conf,
)
from airflow.providers.openlineage.plugins.facets import (
AirflowDagRunFacet,
AirflowDebugRunFacet,
AirflowJobFacet,
AirflowMappedTaskRunFacet,
AirflowRunFacet,
AirflowStateRunFacet,
UnknownOperatorAttributeRunFacet,
UnknownOperatorInstance,
)
from airflow.providers.openlineage.utils.selective_enable import (
is_dag_lineage_enabled,
is_task_lineage_enabled,
)
from airflow.providers.openlineage.version_compat import (
AIRFLOW_V_3_0_PLUS,
AIRFLOW_V_3_2_2_PLUS,
AIRFLOW_V_3_2_PLUS,
AIRFLOW_V_3_3_PLUS,
AIRFLOW_V_3_4_PLUS,
get_base_airflow_version_tuple,
)
from airflow.serialization.serialized_objects import SerializedBaseOperator, SerializedDAG
try:
from airflow.serialization.definitions.mappedoperator import SerializedMappedOperator
except ImportError:
from airflow.models.mappedoperator import ( # type: ignore[attr-defined,no-redef]
MappedOperator as SerializedMappedOperator,
)
try:
from airflow.sdk.definitions.taskgroup import MappedTaskGroup
except ImportError: # Airflow 2
from airflow.utils.task_group import MappedTaskGroup # type: ignore[no-redef]
if not AIRFLOW_V_3_0_PLUS:
from airflow.utils.session import NEW_SESSION, provide_session
if TYPE_CHECKING:
from typing import TypeAlias
from openlineage.client.event_v2 import Dataset as OpenLineageDataset
from openlineage.client.facet_v2 import JobFacet, RunFacet, processing_engine_run
from sqlalchemy.orm import Session
from airflow.models.asset import AssetEvent
from airflow.sdk.execution_time.secrets_masker import (
Redactable,
Redacted,
SecretsMasker,
)
from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance
from airflow.utils.state import DagRunState, TaskInstanceState
[docs]
AnyOperator: TypeAlias = BaseOperator | MappedOperator | SerializedBaseOperator | SerializedMappedOperator
else:
try:
from airflow.sdk._shared.secrets_masker import (
Redactable,
Redacted,
SecretsMasker,
should_hide_value_for_key,
)
except ImportError:
try:
from airflow.sdk.execution_time.secrets_masker import (
Redactable,
Redacted,
SecretsMasker,
should_hide_value_for_key,
)
except ImportError:
from airflow.utils.log.secrets_masker import (
Redactable,
Redacted,
SecretsMasker,
should_hide_value_for_key,
)
[docs]
log = logging.getLogger(__name__)
_NOMINAL_TIME_FORMAT = "%Y-%m-%dT%H:%M:%S.%fZ"
_MAX_DOC_BYTES = 64 * 1024 # 64 kilobytes
_PRODUCER = f"https://github.com/apache/airflow/tree/providers-openlineage/{OPENLINEAGE_PROVIDER_VERSION}"
_JOB_TYPE_DAG: Literal["DAG"] = "DAG"
_JOB_TYPE_TASK: Literal["TASK"] = "TASK"
[docs]
def try_import_from_string(string: str) -> Any:
with suppress(ImportError):
return import_string(string)
[docs]
def get_operator_class(task: BaseOperator | MappedOperator) -> type:
if task.__class__.__name__ in ("DecoratedMappedOperator", "MappedOperator"):
return task.operator_class
return task.__class__
[docs]
def get_operator_provider_version(operator: AnyOperator) -> str | None:
"""Get the provider package version for the given operator."""
try:
class_path = get_fully_qualified_class_name(operator)
if not class_path.startswith("airflow.providers."):
return None
from airflow.providers_manager import ProvidersManager
providers_manager = ProvidersManager()
for package_name, provider_info in providers_manager.providers.items():
if package_name.startswith("apache-airflow-providers-"):
provider_module_path = package_name.replace(
"apache-airflow-providers-", "airflow.providers."
).replace("-", ".")
if class_path.startswith(provider_module_path + "."):
return provider_info.version
return None
except Exception:
return None
[docs]
def get_job_name(task: TaskInstance | RuntimeTaskInstance) -> str:
return f"{task.dag_id}.{task.task_id}"
[docs]
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
[docs]
def next_query_counter_from_context() -> str:
"""
Return the next sequential per-task query counter, suitable for use as a job-name suffix.
The counter is stored under a private key in the Airflow execution context dict. Airflow
creates a fresh context per task and discards it when the task ends, so the counter resets
automatically per task and per retry — there is no module-level state to leak.
When no Airflow execution context is available (for example, the helper was called from a
script with an explicit ``task_instance=`` outside a running task), a random suffix
is returned instead and a warning is logged.
"""
try:
from airflow.sdk import get_current_context
except (ImportError, AttributeError, ModuleNotFoundError):
from airflow.operators.python import get_current_context # type: ignore[no-redef]
try:
context = get_current_context()
query_counter_context_key = "_openlineage_manual_query_counter"
# `Context` is a TypedDict, so mypy refuses non-literal keys and types `.get()` of an
# unknown key as `object`. At runtime `Context` is a regular `dict`, so
# both operations below are safe — the type-ignores silence pure type-checker noise.
counter = int(context.get(query_counter_context_key, 0)) + 1 # type: ignore[call-overload]
context[query_counter_context_key] = counter # type: ignore[literal-required]
return str(counter)
except Exception as err:
from airflow.utils.strings import get_random_string
random_suffix = get_random_string()
log.info(
"OpenLineage encountered an error when retrieving query counter from context: `%s`. "
"Returning random suffix `%s`",
err,
random_suffix,
)
return random_suffix
[docs]
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
[docs]
def build_task_event_run_facets(
*,
task_instance,
dag_run,
dag,
task,
task_uuid: str,
ti_state: TaskInstanceState,
parent_run_id: str,
parent_job_name: str | None = None,
dr_conf: dict | None = None,
additional_run_facets: dict[str, RunFacet] | None = None,
include_full_task_info: bool = False,
) -> dict[str, RunFacet]:
"""Build the task-event run-facet dict."""
if dr_conf is None:
dr_conf = getattr(dag_run, "conf", {}) or {}
nominal_start = getattr(dag_run, "data_interval_start", None)
nominal_end = getattr(dag_run, "data_interval_end", None)
if isinstance(nominal_start, datetime.datetime):
nominal_start = nominal_start.isoformat()
if isinstance(nominal_end, datetime.datetime):
nominal_end = nominal_end.isoformat()
nominal_time_facet: dict[str, RunFacet] = (
{
"nominalTime": nominal_time_run.NominalTimeRunFacet(
nominalStartTime=nominal_start, nominalEndTime=nominal_end, producer=_PRODUCER
)
}
if nominal_start
else {}
)
return {
**(additional_run_facets or {}),
**get_user_provided_run_facets(task_instance, ti_state),
**get_task_parent_run_facet(
parent_run_id=parent_run_id,
parent_job_name=parent_job_name or dag.dag_id,
dr_conf=dr_conf,
),
**get_airflow_run_facet(
dag_run=dag_run,
dag=dag,
task_instance=task_instance,
task=task,
task_uuid=task_uuid,
include_full_task_info=include_full_task_info,
),
**get_airflow_debug_facet(),
**get_processing_engine_facet(),
**nominal_time_facet,
}
[docs]
def build_task_event_job_facets(
*,
task,
dag,
additional_job_facets: dict[str, JobFacet] | None = None,
) -> dict[str, JobFacet]:
"""Build the task-event job-facet dict."""
doc, doc_type = get_task_documentation(task)
if not doc:
doc, doc_type = get_dag_documentation(dag)
documentation_facet: dict[str, JobFacet] = (
{
"documentation": documentation_job.DocumentationJobFacet(
description=doc, contentType=doc_type, producer=_PRODUCER
)
}
if doc
else {}
)
owner_source = task if getattr(task, "owner", "airflow") != "airflow" else dag
owners = [o.strip() for o in getattr(owner_source, "owner", "").split(",") if o.strip()]
ownership_facet: dict[str, JobFacet] = (
{
"ownership": ownership_job.OwnershipJobFacet(
owners=[ownership_job.Owner(name=o) for o in sorted(owners)], producer=_PRODUCER
)
}
if owners
else {}
)
dag_tags = getattr(dag, "tags", None) or []
tags_facet: dict[str, JobFacet] = (
{
"tags": tags_job.TagsJobFacet(
tags=[
tags_job.TagsJobFacetFields(key=t, value=t, source="AIRFLOW") for t in sorted(dag_tags)
],
producer=_PRODUCER,
)
}
if dag_tags
else {}
)
return {
**(additional_job_facets or {}),
**_get_job_type_facets(get_airflow_job_type(_JOB_TYPE_TASK)),
**documentation_facet,
**ownership_facet,
**tags_facet,
}
[docs]
def get_airflow_job_type(job_type: Literal["DAG", "TASK"]) -> dict[str, str]:
"""Return the job type facet fields of an Airflow Dag or task."""
return {"processingType": "BATCH", "integration": "AIRFLOW", "jobType": job_type}
def _get_job_type_facets(job_type: dict[str, str] | None) -> dict[str, JobFacet]:
if not job_type:
return {}
return {
"jobType": job_type_job.JobTypeJobFacet(
processingType=job_type["processingType"],
integration=job_type["integration"],
jobType=job_type["jobType"],
producer=_PRODUCER,
)
}
def _get_parent_run_facet(
parent_run_id: str,
parent_job_name: str,
parent_job_namespace: str = conf.namespace(),
root_parent_run_id: str | None = None,
root_parent_job_name: str | None = None,
root_parent_job_namespace: str | None = None,
parent_job_type: dict[str, str] | None = None,
root_parent_job_type: dict[str, str] | None = None,
) -> dict[str, Any]:
"""Create the parent run facet from identifiers."""
if not any((root_parent_run_id, root_parent_job_name, root_parent_job_namespace)):
root_parent_job_type = parent_job_type
root_parent_run_id = root_parent_run_id or parent_run_id
root_parent_job_name = root_parent_job_name or parent_job_name
root_parent_job_namespace = root_parent_job_namespace or parent_job_namespace
return {
"parent": parent_run.ParentRunFacet(
run=parent_run.Run(runId=parent_run_id),
job=parent_run.Job(
namespace=parent_job_namespace,
name=parent_job_name,
facets=_get_job_type_facets(parent_job_type),
),
root=parent_run.Root(
run=parent_run.RootRun(runId=root_parent_run_id),
job=parent_run.RootJob(
namespace=root_parent_job_namespace,
name=root_parent_job_name,
facets=_get_job_type_facets(root_parent_job_type),
),
),
)
}
[docs]
def get_task_parent_run_facet(
parent_run_id: str,
parent_job_name: str,
parent_job_namespace: str = conf.namespace(),
root_parent_run_id: str | None = None,
root_parent_job_name: str | None = None,
root_parent_job_namespace: str | None = None,
dr_conf: dict | None = None,
) -> dict[str, Any]:
"""Retrieve the parent run facet."""
root_parent_job_type = None
all_root_info = (root_parent_run_id, root_parent_job_namespace, root_parent_job_name)
# If not all root identifiers are provided explicitly, try to get them from dagrun conf.
if not all(all_root_info):
# If some but not all root identifiers are provided warn and do not use any of them, we need all.
if any(all_root_info):
log.warning(
"Incomplete root OpenLineage information provided. "
"No root information will be used. Found values: "
"root_parent_run_id='%s', root_parent_job_namespace='%s', root_parent_job_name='%s'.",
root_parent_run_id,
root_parent_job_namespace,
root_parent_job_name,
)
root_parent_run_id, root_parent_job_namespace, root_parent_job_name = None, None, None
elif dr_conf:
# Check for root identifiers in dagrun conf
if root_info := get_root_information_from_dagrun_conf(dr_conf):
root_parent_run_id = root_info["root_parent_run_id"]
root_parent_job_namespace = root_info["root_parent_job_namespace"]
root_parent_job_name = root_info["root_parent_job_name"]
root_parent_job_type = root_info.get("root_parent_job_type")
# If not present, check for parent identifiers in dagrun conf and use them as root
elif parent_info := get_parent_information_from_dagrun_conf(dr_conf):
root_parent_run_id = parent_info["parent_run_id"]
root_parent_job_namespace = parent_info["parent_job_namespace"]
root_parent_job_name = parent_info["parent_job_name"]
root_parent_job_type = parent_info.get("parent_job_type")
return _get_parent_run_facet(
parent_run_id=parent_run_id,
parent_job_namespace=parent_job_namespace,
parent_job_name=parent_job_name,
root_parent_run_id=root_parent_run_id,
root_parent_job_namespace=root_parent_job_namespace,
root_parent_job_name=root_parent_job_name,
parent_job_type=get_airflow_job_type(_JOB_TYPE_DAG),
root_parent_job_type=root_parent_job_type,
)
def _get_openlineage_data_from_dagrun_conf(dr_conf: dict | None) -> dict:
"""Return the 'openlineage' section from a DAG run config if valid, otherwise an empty dict."""
if not dr_conf or not isinstance(dr_conf, dict):
return {}
ol_data = dr_conf.get("openlineage")
return ol_data if isinstance(ol_data, dict) else {}
def _get_job_type_from_ol_data(ol_data: dict, key: str) -> dict[str, str] | None:
"""Return the job type facet fields stored under ``key`` if they are all valid, otherwise None."""
job_type = ol_data.get(key)
if job_type is None:
return None
required_fields = ("processingType", "integration", "jobType")
# The spec allows partial job type, but only complete information is supported for now.
# Other fields (e.g. `_producer` of a serialized facet) are ignored.
if not isinstance(job_type, dict) or not all(
isinstance(job_type.get(field), str) and job_type[field] for field in required_fields
):
log.warning(
"Invalid OpenLineage %s '%s' in DAG run config - expected a dictionary with non-empty "
"string values for all of: %s.",
key,
job_type,
", ".join(required_fields),
)
return None
return {field: job_type[field] for field in required_fields}
[docs]
def get_dag_parent_run_facet(dr_conf: dict | None) -> dict[str, parent_run.ParentRunFacet]:
"""
Build the OpenLineage parent run facet from a DAG run configuration.
This function extracts parent run identifiers - run ID, job namespace, and job name -
from the DAG run configuration to construct an OpenLineage `ParentRunFacet`. It requires
a complete set of parent identifiers to proceed; if some but not all are present, or if
the run ID is invalid, the function returns an empty dictionary.
When valid parent identifiers are found, it also attempts to retrieve corresponding
root identifiers using `get_root_information_from_dagrun_conf`, which may fall back
to the parent identifiers if no explicit root data is available. The resulting facet
links both the immediate parent and root lineage information for the run.
Args:
dr_conf: The DAG run configuration dictionary.
Returns:
A dictionary containing a single entry mapping the facet name to the constructed
`ParentRunFacet`. Returns an empty dictionary if the configuration does not contain
a complete or valid set of parent identifiers.
"""
ol_data = _get_openlineage_data_from_dagrun_conf(dr_conf)
if not ol_data:
log.debug("No 'openlineage' data found in DAG run config.")
return {}
parent_info = get_parent_information_from_dagrun_conf(dr_conf)
if not parent_info: # Validation and logging is done in the above function
return {}
root_parent_run_id, root_parent_job_namespace, root_parent_job_name = None, None, None
root_parent_job_type = None
if root_info := get_root_information_from_dagrun_conf(dr_conf):
root_parent_run_id = root_info["root_parent_run_id"]
root_parent_job_namespace = root_info["root_parent_job_namespace"]
root_parent_job_name = root_info["root_parent_job_name"]
root_parent_job_type = root_info.get("root_parent_job_type")
else:
log.debug(
"Missing OpenLineage root identifiers in DAG run config, "
"parent identifiers will be used as root instead."
)
# Function already uses parent as root if root is missing, no need to explicitly pass it
return _get_parent_run_facet(
parent_run_id=parent_info["parent_run_id"],
parent_job_namespace=parent_info["parent_job_namespace"],
parent_job_name=parent_info["parent_job_name"],
root_parent_run_id=root_parent_run_id,
root_parent_job_namespace=root_parent_job_namespace,
root_parent_job_name=root_parent_job_name,
parent_job_type=parent_info.get("parent_job_type"),
root_parent_job_type=root_parent_job_type,
)
def _truncate_string_to_byte_size(s: str, max_size: int = _MAX_DOC_BYTES) -> str:
"""
Truncate a string to a maximum UTF-8 byte size, ensuring valid encoding.
This is used to safely limit the size of string content (e.g., for OpenLineage events)
without breaking multibyte characters. If truncation occurs, the result is a valid
UTF-8 string with any partial characters at the end removed.
Args:
s (str): The input string to truncate.
max_size (int): Maximum allowed size in bytes after UTF-8 encoding.
Returns:
str: A UTF-8-safe truncated string within the specified byte limit.
"""
encoded = s.encode("utf-8")
if len(encoded) <= max_size:
return s
log.debug(
"Truncating long string content for OpenLineage event. "
"Original size: %d bytes, truncated to: %d bytes (UTF-8 safe).",
len(encoded),
max_size,
)
truncated = encoded[:max_size]
# Make sure we don't cut a multibyte character in half
return truncated.decode("utf-8", errors="ignore")
[docs]
def get_task_documentation(operator: AnyOperator | None) -> tuple[str | None, str | None]:
"""Get task documentation and mime type, truncated to _MAX_DOC_BYTES bytes length, if present."""
if not operator:
return None, None
for attribute, mime_type in (
("doc", "text/plain"),
("doc_md", "text/markdown"),
("doc_json", "application/json"),
("doc_yaml", "application/x-yaml"),
("doc_rst", "text/x-rst"),
):
# The scheduler's mapped operator leaves the ones its task doesn't set undefined
if doc := getattr(operator, attribute, None):
return _truncate_string_to_byte_size(doc), mime_type
return None, None
[docs]
def get_dag_documentation(dag: DAG | SerializedDAG | None) -> tuple[str | None, str | None]:
"""Get dag documentation and mime type, truncated to _MAX_DOC_BYTES bytes length, if present."""
if not dag:
return None, None
doc, mime_type = None, None
if dag.doc_md:
doc = dag.doc_md
mime_type = "text/markdown"
elif dag.description:
doc = dag.description
mime_type = "text/plain"
if doc:
return _truncate_string_to_byte_size(doc), mime_type
return None, None
[docs]
def get_task_instance_map_index(task_instance: TaskInstance | RuntimeTaskInstance) -> int:
return _require_task_instance_position(task_instance).map_index
[docs]
def get_airflow_mapped_task_facet(task_instance: TaskInstance) -> dict[str, Any]:
# check for -1 comes from SmartSensor compatibility with dynamic task mapping
# this comes from Airflow code
log.debug(
"AirflowMappedTaskRunFacet is deprecated and will be removed. "
"Use information from AirflowRunFacet instead."
)
if (map_index := get_task_instance_map_index(task_instance)) != -1 and task_instance.task is not None:
return {
"airflow_mappedTask": AirflowMappedTaskRunFacet(
mapIndex=map_index,
operatorClass=get_fully_qualified_class_name(task_instance.task),
)
}
return {}
[docs]
def get_user_provided_run_facets(ti: TaskInstance, ti_state: TaskInstanceState) -> dict[str, RunFacet]:
custom_facets = {}
# Append custom run facets by executing the custom_run_facet functions.
for custom_facet_func in conf.custom_run_facets():
try:
func: Callable[[TaskInstance, TaskInstanceState], dict[str, RunFacet]] | None = (
try_import_from_string(custom_facet_func)
)
if not func:
log.warning(
"OpenLineage is unable to import custom facet function `%s`; will ignore it.",
custom_facet_func,
)
continue
facets: dict[str, RunFacet] | None = func(ti, ti_state)
if facets and isinstance(facets, dict):
duplicate_facet_keys = [facet_key for facet_key in facets if facet_key in custom_facets]
if duplicate_facet_keys:
log.warning(
"Duplicate OpenLineage custom facets key(s) found: `%s` from function `%s`; "
"this will overwrite the previous value.",
", ".join(duplicate_facet_keys),
custom_facet_func,
)
log.debug(
"Adding OpenLineage custom facet with key(s): `%s` from function `%s`.",
tuple(facets),
custom_facet_func,
)
custom_facets.update(facets)
except Exception as exc:
log.warning(
"Error processing custom facet function `%s`; will ignore it. Error was: %s: %s",
custom_facet_func,
type(exc).__name__,
exc,
)
return custom_facets
[docs]
def get_fully_qualified_class_name(operator: AnyOperator) -> str:
if isinstance(operator, (SerializedMappedOperator, SerializedBaseOperator)):
# as in airflow.api_connexion.schemas.common_schema.ClassReferenceSchema
return operator._task_module + "." + operator.task_type
op_class = get_operator_class(operator)
return op_class.__module__ + "." + op_class.__name__
[docs]
def is_operator_disabled(operator: AnyOperator) -> bool:
return get_fully_qualified_class_name(operator) in conf.disabled_operators()
[docs]
def is_selective_lineage_enabled(obj: DAG | SerializedDAG | AnyOperator) -> bool:
"""If selective enable is active check if DAG or Task is enabled to emit events."""
if not conf.selective_enable():
return True
if isinstance(obj, (DAG, SerializedDAG)):
return is_dag_lineage_enabled(obj)
if isinstance(obj, (BaseOperator, MappedOperator)):
return is_task_lineage_enabled(obj)
raise TypeError("is_selective_lineage_enabled can only be used on DAG or Operator objects")
if not AIRFLOW_V_3_0_PLUS:
@provide_session
[docs]
def is_ti_rescheduled_already(ti: TaskInstance, *, session=NEW_SESSION):
try:
from sqlalchemy import exists, select
except ImportError:
raise AirflowOptionalProviderFeatureException(
"sqlalchemy is required for checking task instance reschedule status. "
"Install it with: pip install 'apache-airflow-providers-openlineage[sqlalchemy]'"
)
if not isinstance(ti.task, BaseSensorOperator):
return False
if not ti.task.reschedule:
return False
return (
session.scalar(
select(
exists().where(
TaskReschedule.dag_id == ti.dag_id,
TaskReschedule.task_id == ti.task_id,
TaskReschedule.run_id == ti.run_id,
TaskReschedule.map_index == ti.map_index,
TaskReschedule.try_number == ti.try_number,
)
)
)
is True
)
[docs]
def safe_getattr(obj: Any, attr: str, default: Any = None) -> Any:
"""Get attribute from object, returning default if DetachedInstanceError is raised."""
try:
return getattr(obj, attr, default)
except DetachedInstanceError:
return default
[docs]
class InfoJsonEncodable(dict):
"""
Airflow objects might not be json-encodable overall.
The class provides additional attributes to control
what and how is encoded:
* renames: a dictionary of attribute name changes
* | casts: a dictionary consisting of attribute names
| and corresponding methods that should change
| object value
* includes: list of attributes to be included in encoding
* excludes: list of attributes to be excluded from encoding
Don't use both includes and excludes.
"""
[docs]
renames: dict[str, str] = {}
[docs]
casts: dict[str, Any] = {}
[docs]
includes: list[str] = []
[docs]
excludes: list[str] = []
def __init__(self, obj):
self._fields = []
self._cast_fields()
self._rename_fields()
self._include_fields()
self._extend_fields()
dict.__init__(
self,
**{field: InfoJsonEncodable._cast_basic_types(getattr(self, field)) for field in self._fields},
)
# Drop the build-time attributes, keeping the dict items set above. Attributes may hold live Airflow
# objects that can't be pickled for the listener's process pool, which would drop the event.
self.__dict__.clear()
@staticmethod
def _cast_basic_types(value):
if isinstance(value, datetime.datetime):
return value.isoformat()
if isinstance(value, datetime.timedelta):
return f"{value.total_seconds()} seconds"
if isinstance(value, (set, list, tuple)):
return str(list(value))
return value
def _extend_fields(self) -> None:
"""
Subclass hook to register extra fields not sourced from ``self.obj``.
Use this primarily for values passed into ``__init__`` (e.g. caller-supplied
context) or values computed on the subclass itself. For anything derived from
``self.obj``, prefer the declarative ``casts`` / ``renames`` / ``includes``
machinery — this hook is for fields that don't fit that model.
Runs after ``_cast_fields`` / ``_rename_fields`` / ``_include_fields`` and
before ``dict.__init__`` freezes ``self`` into its serialized form, so
implementations should both set the attribute on ``self`` and append its
name to ``self._fields``.
"""
def _rename_fields(self):
for field, renamed in self.renames.items():
if hasattr(self.obj, field):
setattr(self, renamed, getattr(self.obj, field))
self._fields.append(renamed)
def _cast_fields(self):
for field, func in self.casts.items():
setattr(self, field, func(self.obj))
self._fields.append(field)
def _include_fields(self):
if self.includes and self.excludes:
raise ValueError("Don't use both includes and excludes.")
if self.includes:
for field in self.includes:
if field not in self._fields and hasattr(self.obj, field):
setattr(self, field, getattr(self.obj, field))
self._fields.append(field)
else:
if hasattr(self.obj, "__dict__"):
obj_fields = self.obj.__dict__
elif attrs.has(self.obj.__class__): # e.g. attrs.define class with slots=True has no __dict__
obj_fields = {
field.name: getattr(self.obj, field.name) for field in attrs.fields(self.obj.__class__)
}
else:
raise ValueError(
"Cannot iterate over fields: "
f"The object of type {type(self.obj).__name__} neither has a __dict__ attribute "
"nor is defined as an attrs class."
)
for field, val in obj_fields.items():
# Keys that aren't identifiers are set by libraries rather than being attributes, e.g. the
# unpicklable ``__wire|<class>|<method>`` caches of ``methodtools.lru_cache``.
if not field.isidentifier():
continue
if field not in self._fields and field not in self.excludes and field not in self.renames:
# Any attribute may hold live Airflow objects, even nested ones that can't be pickled
setattr(self, field, _convert_to_json_safe(self._cast_basic_types(val), field))
self._fields.append(field)
[docs]
class DagInfo(InfoJsonEncodable):
"""Defines encoding DAG object to JSON."""
[docs]
includes = [
"dag_id",
"description",
"fileloc",
"owner",
"owner_links",
"schedule_interval", # For Airflow 2 only -> AF3 has timetable_summary
"start_date",
"tags",
]
[docs]
casts = {
"timetable": lambda dag: DagInfo.serialize_timetable(dag),
"timetable_summary": lambda dag: DagInfo.timetable_summary(dag),
}
[docs]
renames = {"_dag_id": "dag_id"}
@classmethod
@classmethod
[docs]
class DagRunInfo(InfoJsonEncodable):
"""Defines encoding DagRun object to JSON."""
[docs]
includes = [
"clear_number",
"conf",
"dag_id",
"data_interval_end",
"data_interval_start",
"end_date",
"execution_date", # Airflow 2
"external_trigger", # Removed in Airflow 3, use run_type instead
"logical_date", # Airflow 3
"partition_key", # Airflow 3.2+
"partition_date", # Airflow 3.2+
"run_after", # Airflow 3
"run_id",
"run_type",
"start_date",
"triggered_by",
"triggering_user_name", # Airflow 3
]
[docs]
casts = {
"note": lambda dagrun: safe_getattr(dagrun, "note") if AIRFLOW_V_3_2_PLUS else None,
"duration": lambda dagrun: DagRunInfo.duration(dagrun),
"dag_bundle_name": lambda dagrun: DagRunInfo.dag_version_info(dagrun, "bundle_name"),
"dag_bundle_version": lambda dagrun: DagRunInfo.dag_version_info(dagrun, "bundle_version"),
"dag_version_id": lambda dagrun: DagRunInfo.dag_version_info(dagrun, "version_id"),
"dag_version_number": lambda dagrun: DagRunInfo.dag_version_info(dagrun, "version_number"),
"dag_version_data": lambda dagrun: (
DagRunInfo.dag_version_info(dagrun, "version_data") if AIRFLOW_V_3_3_PLUS else None
),
"dag_team_name": lambda dagrun: DagRunInfo.team_name(dagrun) if AIRFLOW_V_3_3_PLUS else None,
"deadlines": lambda dagrun: DagRunInfo.deadlines(dagrun),
}
@classmethod
[docs]
def duration(cls, dagrun: DagRun) -> float | None:
if not getattr(dagrun, "end_date", None) or not isinstance(dagrun.end_date, datetime.datetime):
return None
if not getattr(dagrun, "start_date", None) or not isinstance(dagrun.start_date, datetime.datetime):
return None
return (dagrun.end_date - dagrun.start_date).total_seconds()
@classmethod
[docs]
def deadlines(cls, dagrun: DagRun) -> dict[str, Any] | None:
"""
Extract deadline state and alert definitions from a DagRun (on scheduler).
Returns a dict (not a list) so _cast_basic_types passes it through.
"""
try:
# AF2 DagRun and AF3 DagRun SDK model (on worker) do not have this information
deadlines = safe_getattr(dagrun, "deadlines")
if not deadlines:
return None
except Exception as err:
log.warning("OpenLineage failed to retrieve deadlines. Exception: %s", err)
return None
result = []
for d in deadlines:
try:
info: dict[str, Any] = {}
if deadline_time := getattr(d, "deadline_time", None):
info["deadline_time"] = deadline_time.isoformat()
if (missed := getattr(d, "missed", None)) is not None:
info["missed"] = missed
try:
# deadline_alert is a lazy-loaded ORM relationship that may
# trigger a DB query; keep it isolated so a detached-session
# error doesn't discard the rest of the deadline info.
if alert := safe_getattr(d, "deadline_alert"):
info.update(
{
k: v
for k in ("name", "description", "reference", "interval", "callback_def")
if (v := getattr(alert, k, None)) is not None
}
)
except Exception as err:
log.warning("OpenLineage could not load deadline_alert relationship for %s", err)
if info:
result.append(info)
except Exception as err:
log.warning("OpenLineage failed to serialize deadline: %s", err)
return {"alerts": result} if result else None
@classmethod
[docs]
def dag_version_info(cls, dagrun: DagRun, key: str) -> str | int | dict | None:
"""Extract DAG version info for given key, sourced from DagRun (on scheduler)."""
# AF2 DagRun and AF3 DagRun SDK model (on worker) do not have this information
dag_versions = safe_getattr(dagrun, "dag_versions", [])
if not dag_versions:
return None
# The DagVersion rows themselves are lazy-loaded, so reading their columns can still hit
# a detached session even though fetching the list above succeeded.
current_version = dag_versions[-1]
if key == "bundle_name":
return safe_getattr(current_version, "bundle_name")
if key == "bundle_version":
return safe_getattr(current_version, "bundle_version")
if key == "version_id":
version_id = safe_getattr(current_version, "id")
return str(version_id) if version_id is not None else None
if key == "version_number":
return safe_getattr(current_version, "version_number")
if key == "version_data":
if not AIRFLOW_V_3_3_PLUS:
return None
return safe_getattr(current_version, "version_data")
raise ValueError(f"Unsupported key: {key}`")
@classmethod
[docs]
def team_name(cls, dagrun: DagRun) -> str | None:
"""Extract the team name for the DagRun."""
if not AIRFLOW_V_3_3_PLUS or not airflow_conf.getboolean("core", "multi_team", fallback=False):
return None
# The Execution API delivers the team name on the DagRun payload it sends to the task
# runner, so task events resolve it from there rather than through the bundle lookup below,
# which needs a metadata DB session the task runner does not have.
# `hasattr` rather than a None check -- a team-less run legitimately carries None. The ORM
# DagRun exposes `team_name` too, but reading it lazy-loads the bundle, so scheduler-side
# runs stay on the guarded lookup below instead of risking a load on a detached instance.
if not isinstance(dagrun, DagRun) and hasattr(dagrun, "team_name"):
return dagrun.team_name
# Best-effort: the scheduler stamps `_team_name` on ORM DagRun objects before
# listener hooks fire. It's a private attribute with no stability guarantee,
# so guard with hasattr and an isinstance check.
if hasattr(dagrun, "_team_name"):
return dagrun._team_name if isinstance(dagrun._team_name, str) else None
try:
# Reuse the existing ORM session associated with the DagRun. Creating a new session here
# (via @provide_session) on get_team_name() can trigger an unexpected commit error.
from sqlalchemy.orm import object_session
session = object_session(dagrun)
if session is None:
return None
from airflow.models.dag import DagModel
return DagModel.get_team_name(dagrun.dag_id, session=session)
except Exception as e:
log.info(
"OpenLineage failed to resolve the team name for dag `%s`: %s.",
dagrun.dag_id,
e,
)
log.debug("Exception details:", exc_info=True)
return None
[docs]
class TaskInstanceInfo(InfoJsonEncodable):
"""Defines encoding TaskInstance object to JSON."""
[docs]
includes = ["log_url", "max_tries", "pool", "queued_dttm", "state", "try_number"]
[docs]
casts = {
"log_url": lambda ti: getattr(ti, "log_url", None),
"dag_version_id": lambda ti: (
str(dag_version_id) if (dag_version_id := getattr(ti, "dag_version_id", None)) else None
),
"dag_bundle_version": lambda ti: (
ti.bundle_instance.version if hasattr(ti, "bundle_instance") else None
),
"dag_bundle_name": lambda ti: ti.bundle_instance.name if hasattr(ti, "bundle_instance") else None,
}
@staticmethod
[docs]
def get_rendered_map_index(ti: TaskInstance | RuntimeTaskInstance) -> str | None:
"""Return the label a mapped task rendered with ``map_index_template``, as the worker sends it."""
if AIRFLOW_V_3_0_PLUS and isinstance(ti, TaskInstance):
# Its rendered_map_index falls back to the map index when no label was rendered, the worker's doesn't
return safe_getattr(ti, "_rendered_map_index")
return getattr(ti, "rendered_map_index", None)
def __init__(
self,
obj,
position: TaskInstancePosition | None = None,
dag: Any = None,
dag_run: Any = None,
end_date: datetime.datetime | None = None,
):
# Pass the position when the caller has already placed the task instance, e.g. from its stored region or
# with get_task_instance_position; it is then used as is.
self._position = position
# A stored task instance's bundle comes from these, as the worker's does, without a query
self._dag = dag
self._dag_run = dag_run
# The end date a terminal event is timed with, used when the task instance carries none, as a skipped
# Airflow 3 runtime task instance doesn't before Airflow 3.4.
self._end_date = end_date
super().__init__(obj)
def _get_start_date(self) -> datetime.datetime | None:
if not isinstance(start_date := getattr(self.obj, "start_date", None), datetime.datetime):
return None
# Before Airflow 3.2.2 a runtime task instance resumed from deferral starts at the resume time, while the
# stored start date stays the original one. Leaving it out keeps a wrong duration out of the facet too.
if not AIRFLOW_V_3_2_2_PLUS and self._is_resumed_from_deferral():
return None
return start_date
def _is_resumed_from_deferral(self) -> bool:
ti_context = getattr(self.obj, "_ti_context_from_server", None)
return bool(getattr(ti_context, "next_method", None))
def _get_end_date(self) -> datetime.datetime | None:
if isinstance(end_date := getattr(self.obj, "end_date", None), datetime.datetime):
return end_date
return self._end_date
def _get_duration(
self, start_date: datetime.datetime | None, end_date: datetime.datetime | None
) -> float | None:
# The stored duration is what the UI shows, so it wins. A runtime task instance has none -- the API
# server computes it after the listener ran. Airflow 2 keeps the previous try's duration when it clears
# end_date for the next try, so the stored duration only counts alongside an end_date.
if (duration := getattr(self.obj, "duration", None)) is not None and isinstance(
getattr(self.obj, "end_date", None), datetime.datetime
):
return duration
if start_date is None or end_date is None:
return None
try:
return (end_date - start_date).total_seconds()
except TypeError: # naive and aware datetimes can't be subtracted
return None
def _extend_fields(self) -> None:
start_date = self._get_start_date()
end_date = self._get_end_date()
for field, field_value in (
("start_date", start_date),
("end_date", end_date),
("duration", self._get_duration(start_date, end_date)),
):
if field_value is not None:
setattr(self, field, field_value)
self._fields.append(field)
position = self._position
if position is None:
# A stored task instance whose task isn't attached can't be placed without its region, and a facet
# must not fail the whole event, so then it describes the task instance without its position.
position = _get_task_instance_position(self.obj)
if position is not None:
position = _add_runtime_loop_iteration(self.obj, position)
map_index = position.map_index if position else -1
self.map_index = map_index if map_index != -1 else None
self.rendered_map_index = self.get_rendered_map_index(self.obj) if map_index != -1 else None
self._fields.extend(("map_index", "rendered_map_index"))
is_stored = isinstance(self.obj, TaskInstance)
if is_stored and AIRFLOW_V_3_0_PLUS:
# The worker gets the bundle name of the Dag and the bundle version of the Dag run
self.dag_bundle_name = getattr(self._dag, "bundle_name", None)
self.dag_bundle_version = getattr(self._dag_run, "bundle_version", None)
if position and position.expanded_ti_count is not None:
self.expanded_ti_count = position.expanded_ti_count
self._fields.append("expanded_ti_count")
# Before Airflow 3.4, a stored task instance failing with retries left already holds its next attempt's id
if (
AIRFLOW_V_3_0_PLUS
and (AIRFLOW_V_3_4_PLUS or not is_stored)
and (ti_id := getattr(self.obj, "id", None)) is not None
):
self.id = str(ti_id)
self._fields.append("id")
loop: dict[str, Any] = {}
if position and position.in_loop:
for name, value in (
("index", position.loop_index),
("region_id", position.loop_region_id),
("trajectory", position.loop_trajectory),
("forked_from_region_id", position.loop_forked_from_region_id),
("resumes_from_index", position.loop_resumes_from_index),
):
if value is not None:
loop[name] = str(value) if isinstance(value, UUID) else value
# Only a loop gate's runtime task instance has one, once it has evaluated ``until``
if (decision := getattr(self.obj, "loop_decision", None)) is not None:
loop["decision"] = {"stop": getattr(decision, "stop", None)}
if (reason := getattr(decision, "reason", None)) is not None:
loop["decision"]["reason"] = reason
if loop:
self.loop = loop
self._fields.append("loop")
# A runtime task instance in no region carries the sentinel region id instead of None
region_id = getattr(self.obj, "region_id", None)
if isinstance(region_id, UUID) and region_id.int != 0:
self.region_id = str(region_id)
self.region_index = getattr(self.obj, "region_index", None)
self._fields.extend(("region_id", "region_index"))
[docs]
class AssetInfo(InfoJsonEncodable):
"""
Defines encoding Airflow Asset object to JSON.
Output shape: ``{"uri", "extra", "type"[, "source_alias"]}``.
``type`` distinguishes how the asset showed up on the task:
* ``"asset"`` - statically declared on ``task.outlets`` / ``task.inlets``.
* ``"asset_event"`` — emitted dynamically via ``outlet_events[asset].extra``.
* ``"asset_event_from_alias"`` — resolved at runtime from an ``AssetAlias``
via ``outlet_events[alias].add(asset, ...)``; ``source_alias`` then carries the alias name.
``AssetAlias`` are intentionally NOT included - they are just helpers with no URI of their own.
"""
[docs]
includes = ["uri", "extra"]
def __init__(self, obj, asset_type: str = "asset", source_alias: str | None = None):
[docs]
self.source_alias = source_alias
super().__init__(obj)
def _extend_fields(self) -> None:
if self.type:
self._fields.append("type")
if self.source_alias:
self._fields.append("source_alias")
[docs]
def get_runtime_outlet_assets(task_instance) -> list[tuple[Asset, AssetAlias | None]]:
"""
Return ``(asset, alias)`` pairs for every runtime-attached outlet.
* ``outlet_events[AssetAlias].add(asset, ...)`` → ``(asset, alias)``.
* ``outlet_events[Asset].extra = {...}`` → ``(asset, None)``.
Empty list on AF2 or when ``task_instance`` is not an AF3 ``RuntimeTaskInstance``
(scheduler / API-server path has no usable worker template context).
"""
if not AIRFLOW_V_3_0_PLUS:
return []
try:
from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance
except ImportError:
return []
if not isinstance(task_instance, RuntimeTaskInstance):
return []
if not hasattr(task_instance, "get_template_context"):
return []
outlet_events = task_instance.get_template_context().get("outlet_events")
if outlet_events is None:
return []
results: list[tuple[Asset, AssetAlias | None]] = []
for key in outlet_events:
accessor = outlet_events[key]
if isinstance(key, AssetAlias):
for event in getattr(accessor, "asset_alias_events", []) or []:
# ``dest_asset_extra`` was introduced in Airflow 3.1.4 ; on
# earlier versions the alias event carries only ``extra``.
dest_asset_extra = getattr(event, "dest_asset_extra", None) or {}
results.append(
(
Asset(
name=event.dest_asset_key.name,
uri=event.dest_asset_key.uri,
# Per-event extra wins over asset-level extra.
extra={**dest_asset_extra, **(event.extra or {})},
),
key,
)
)
elif isinstance(key, Asset):
if extra := (getattr(accessor, "extra", None) or {}):
results.append((Asset(name=key.name, uri=key.uri, extra=dict(extra)), None))
return results
[docs]
class TaskInfo(InfoJsonEncodable):
"""Defines encoding BaseOperator/AbstractOperator object to JSON."""
[docs]
renames = {
"_BaseOperator__from_mapped": "mapped",
"_downstream_task_ids": "downstream_task_ids",
"_upstream_task_ids": "upstream_task_ids",
"_is_setup": "is_setup",
"_is_teardown": "is_teardown",
}
[docs]
includes = [
"deferrable",
"depends_on_past",
"downstream_task_ids",
"execution_timeout",
"executor_config",
"ignore_first_depends_on_past",
"max_active_tis_per_dag",
"max_active_tis_per_dagrun",
"max_retry_delay",
"multiple_outputs",
"owner",
"priority_weight",
"queue",
"retries",
"retry_delay",
"retry_exponential_backoff",
"run_as_user",
"sla",
"task_id",
"trigger_rule",
"upstream_task_ids",
"wait_for_downstream",
"wait_for_past_depends_before_skipping",
# Operator-specific useful attributes
"trigger_dag_id", # TriggerDagRunOperator
"trigger_run_id", # TriggerDagRunOperator
"note", # TriggerDagRunOperator
"external_dag_id", # ExternalTaskSensor and ExternalTaskMarker (if run, as it's EmptyOperator)
"external_task_id", # ExternalTaskSensor and ExternalTaskMarker (if run, as it's EmptyOperator)
"external_task_ids", # ExternalTaskSensor
"external_task_group_id", # ExternalTaskSensor
"external_dates_filter", # ExternalTaskSensor
"logical_date", # AF 3 ExternalTaskMarker (if run, as it's EmptyOperator)
"execution_date", # AF 2 ExternalTaskMarker (if run, as it's EmptyOperator)
"database", # BaseSQlOperator
"parameters", # SQLCheckOperator, SQLValueCheckOperator and BranchSQLOperator
"column_mapping", # SQLColumnCheckOperator
"pass_value", # SQLValueCheckOperator
"tol", # SQLValueCheckOperator
"metrics_thresholds", # SQLIntervalCheckOperator
"ratio_formula", # SQLIntervalCheckOperator
"ignore_zero", # SQLIntervalCheckOperator
"min_threshold", # SQLThresholdCheckOperator
"max_threshold", # SQLThresholdCheckOperator
"follow_task_ids_if_true", # BranchSQLOperator
"follow_task_ids_if_false", # BranchSQLOperator
"follow_branch", # BranchSQLOperator
"preoperator", # SQLInsertRowsOperator
"postoperator", # SQLInsertRowsOperator
"table_name_with_schema", # SQLInsertRowsOperator
"column_names", # SQLInsertRowsOperator
"hitl_summary", # All HITLOperator based operators
]
[docs]
casts = {
"operator_class": lambda task: task.task_type,
"operator_class_path": lambda task: get_fully_qualified_class_name(task),
"task_group": lambda task: (
TaskGroupInfo(task.task_group)
if hasattr(task, "task_group") and getattr(task.task_group, "_group_id", None)
else None
),
# task_group is the task's own group, without loop fields, and loop the loop around the task, so each is
# described once. For a loop "body", alone or in a task group "outer" (task_group's unique_group_id):
# task "body.work": task_group "body", the loop itself loop "body"
# task "body.inner.work": task_group "body.inner" loop "body"
# task "outer.body.work": task_group "outer.body", the loop loop "outer.body"
"loop": lambda task: get_task_loop_info(task),
"inlets": lambda task: [AssetInfo(i) for i in task.inlets if isinstance(i, Asset)],
"outlets": lambda task: [AssetInfo(o) for o in task.outlets if isinstance(o, Asset)],
"operator_provider_version": lambda task: get_operator_provider_version(task),
}
def __init__(self, obj, runtime_assets: list[tuple[Asset, AssetAlias | None]] | None = None):
self._runtime_assets = runtime_assets # Will not be included in final dict, not in self._fields.
super().__init__(obj)
def _extend_fields(self) -> None:
if self._runtime_assets and not hasattr(self, "outlets"): # Should never happen, just in case
self.outlets = []
self._fields.append("outlets")
for asset, alias in self._runtime_assets or ():
# Append runtime-resolved outlets (alias resolutions + dynamic asset events).
# Inlets are not touched - Airflow has no dynamic inlets.
self.outlets.append(
AssetInfo(asset, asset_type="asset_event")
if alias is None
else AssetInfo(asset, asset_type="asset_event_from_alias", source_alias=alias.name)
)
# ``inlets`` / ``outlets`` are a Python repr string that consumers cannot reliably parse
# (``extra`` can hold quotes, ``Decimal``, ``True``, ...). Expose them as real lists too.
# Set as items rather than fields so ``_cast_basic_types`` does not turn them into strings.
for field in ("inlets", "outlets"):
if hasattr(self, field):
self[f"{field}_list"] = _convert_to_json_safe(getattr(self, field), field)
[docs]
class TaskInfoComplete(TaskInfo):
"""Defines encoding BaseOperator/AbstractOperator object to JSON used when user enables full task info."""
[docs]
excludes = [
"_BaseOperator__instantiated",
"_BaseOperator__init_kwargs", # Causes recursion error on AF3, nothing useful there
"_dag",
"dag",
"expand_input",
"op_kwargs_expand_input",
"_hook",
"_log",
"_outlets",
"_inlets",
"_lock_for_execution",
"handler",
"params",
"python_callable",
]
[docs]
class TaskGroupInfo(InfoJsonEncodable):
"""Defines encoding TaskGroup object to JSON."""
[docs]
renames = {
"_group_id": "group_id",
}
[docs]
casts = {
# group_id is the group's own id, while the job facet and task.loop name groups by their unique id
"unique_group_id": lambda tg: getattr(tg, "group_id", None),
}
[docs]
includes = [
"downstream_group_ids",
"downstream_task_ids",
"prefix_group_id",
"tooltip",
"upstream_group_ids",
"upstream_task_ids",
]
[docs]
def get_airflow_dag_run_facet(dag_run: DagRun) -> dict[str, RunFacet]:
if not dag_run.dag:
return {}
return {
"airflowDagRun": AirflowDagRunFacet(
dag=DagInfo(dag_run.dag),
dagRun=DagRunInfo(dag_run),
)
}
@conf.cache
def _get_all_packages_installed() -> dict[str, str]:
"""
Retrieve a dictionary of all installed packages and their versions.
This operation involves scanning the system's installed packages, which can be a heavy operation.
It is recommended to cache the result to avoid repeated, expensive lookups.
"""
return {dist.metadata["Name"]: dist.version for dist in metadata.distributions()}
[docs]
def get_processing_engine_facet() -> dict[str, processing_engine_run.ProcessingEngineRunFacet]:
from openlineage.client.facet_v2 import processing_engine_run
return {
"processing_engine": processing_engine_run.ProcessingEngineRunFacet(
version=AIRFLOW_VERSION,
name="Airflow",
openlineageAdapterVersion=OPENLINEAGE_PROVIDER_VERSION,
)
}
[docs]
def get_airflow_debug_facet() -> dict[str, AirflowDebugRunFacet]:
if not conf.debug_mode():
return {}
log.warning("OpenLineage debug_mode is enabled. Be aware that this may log and emit extensive details.")
return {
"debug": AirflowDebugRunFacet(
packages=_get_all_packages_installed(),
)
}
[docs]
def get_airflow_run_facet(
dag_run: DagRun,
dag: DAG | SerializedDAG,
task_instance: TaskInstance,
task: BaseOperator,
task_uuid: str,
include_full_task_info: bool = False,
ti_position: TaskInstancePosition | None = None,
end_date: datetime.datetime | None = None,
) -> dict[str, AirflowRunFacet]:
runtime_assets = get_runtime_outlet_assets(task_instance)
return {
"airflow": AirflowRunFacet(
dag=DagInfo(dag),
dagRun=DagRunInfo(dag_run),
taskInstance=TaskInstanceInfo(
task_instance, position=ti_position, dag=dag, dag_run=dag_run, end_date=end_date
),
task=(
TaskInfoComplete(task, runtime_assets=runtime_assets)
if include_full_task_info
else TaskInfo(task, runtime_assets=runtime_assets)
),
taskUuid=task_uuid,
)
}
[docs]
def get_airflow_job_facet(dag_run: DagRun) -> dict[str, AirflowJobFacet]:
if not dag_run.dag:
return {}
try:
edge_map = _build_labeled_edge_map(dag_run.dag)
return {
"airflow": AirflowJobFacet(
taskTree={}, # caused OOM errors, to be removed, see #41587
taskGroups=_get_task_groups_details(dag=dag_run.dag, edge_map=edge_map),
tasks=_get_tasks_details(dag=dag_run.dag, edge_map=edge_map),
)
}
except Exception as e:
log.warning("Failed to build AirflowJobFacet for DagRun %s/%s: %s.", dag_run.dag, dag_run.run_id, e)
log.debug("Exception details:", exc_info=True)
return {}
[docs]
def get_airflow_state_run_facet(
dag_id: str, run_id: str, task_ids: list[str], dag_run_state: DagRunState
) -> dict[str, AirflowStateRunFacet]:
try:
tis = DagRun.fetch_task_instances(dag_id=dag_id, run_id=run_id, task_ids=task_ids)
def get_task_duration(ti):
if ti.duration is not None:
return ti.duration
if ti.end_date is not None and ti.start_date is not None:
return (ti.end_date - ti.start_date).total_seconds()
# Fallback to 0.0 for tasks with missing timestamps (e.g., skipped/terminated tasks)
return 0.0
return {
"airflowState": AirflowStateRunFacet(
dagRunState=dag_run_state,
# TODO: keyed by task_id only, so the mapped instances and loop iterations of one task collapse
# into a single entry (the last one by region wins). Needs the map index and loop iteration here.
tasksState={ti.task_id: ti.state for ti in tis},
tasksDuration={ti.task_id: get_task_duration(ti) for ti in tis},
)
}
except Exception as e:
log.warning("Failed to build AirflowStateRunFacet for DagRun %s/%s: %s.", dag_id, run_id, e)
log.debug("Exception details:", exc_info=True)
return {}
[docs]
def is_dag_run_asset_triggered(
dag_run: DagRun,
):
"""Return whether the given DAG run was triggered by an asset."""
if AIRFLOW_V_3_0_PLUS:
from airflow.utils.types import DagRunTriggeredByType
return dag_run.triggered_by == DagRunTriggeredByType.ASSET
# AF 2 Path
from airflow.models.dagrun import DagRunType
return dag_run.run_type == DagRunType.DATASET_TRIGGERED # type: ignore[attr-defined] # This attr is available on AF2, but mypy can't see it
@attrs.frozen
[docs]
class TaskInstancePosition:
"""
Where a task instance runs: in a loop or not, the map index users see, and its loop iteration.
The loop fields are ``None`` outside a loop, and also when they can't be known. A runtime task instance
tells all of them, and ``expanded_ti_count`` for a mapped task, through its ``loop`` and
``expanded_ti_count``, read by ``get_task_instance_position``. A stored one tells only what its task and
region rows do: the iteration of a task that isn't mapped, from its ``region_index``.
``AirflowRunFacet.json`` describes what these fields mean in events.
"""
[docs]
loop_index: int | None = None
[docs]
loop_region_id: UUID | None = None
[docs]
loop_trajectory: int | None = None
[docs]
loop_forked_from_region_id: UUID | None = None
[docs]
loop_resumes_from_index: int | None = None
[docs]
expanded_ti_count: int | None = None
def _is_loop_task_group(task_group: Any) -> bool:
if not AIRFLOW_V_3_4_PLUS:
return False
# Loops only exist from Airflow 3.4 on
from airflow.sdk.definitions._internal.loop import LoopTaskGroup
from airflow.serialization.definitions.taskgroup import SerializedLoopTaskGroup
return isinstance(task_group, (LoopTaskGroup, SerializedLoopTaskGroup))
def _find_enclosing_loop(task: Any) -> tuple[bool, Any | None]:
"""Tell whether a task runs in a loop, and return that loop's group when it is known."""
from airflow.sdk.definitions.taskgroup import TaskGroup
from airflow.serialization.definitions.taskgroup import SerializedTaskGroup
group = getattr(task, "task_group", None)
while isinstance(group, (TaskGroup, SerializedTaskGroup)):
if _is_loop_task_group(group):
return True, group
group = group.parent_group
if group is None:
return False, None
# A group of an unknown type could be a loop. Assuming one only gives each iteration its own runId, while
# missing one would merge the runs of all its iterations, which share every coordinate but the UUID.
log.warning(
"OpenLineage can't tell whether task `%s` runs in a loop, as its task group has an unknown type `%s`; "
"assuming it does",
getattr(task, "task_id", None),
type(group).__name__,
)
return True, None
[docs]
def get_task_loop_info(task: Any) -> dict[str, Any] | None:
"""Describe the loop around a task as ``task.loop``, or return None for a task outside a known loop."""
if not AIRFLOW_V_3_4_PLUS:
return None
if (loop_group := _find_enclosing_loop(task)[1]) is None:
return None
return {
field: getattr(loop_group, field, None)
for field in ("group_id", "max_iterations", "has_until", "gate_task_id", "terminal_task_id")
}
def _get_task_instance_position(
task_instance: TaskInstance | RuntimeTaskInstance,
) -> TaskInstancePosition | None:
"""Place a task instance from its own fields and task, or return None when only its stored region can tell."""
is_stored = isinstance(task_instance, TaskInstance)
if not AIRFLOW_V_3_4_PLUS:
return TaskInstancePosition(in_loop=False, map_index=getattr(task_instance, "map_index", -1))
map_index: int = getattr(task_instance, "region_index" if is_stored else "map_index", -1)
region_id = getattr(task_instance, "region_id", None)
if not isinstance(region_id, UUID) or region_id.int == 0:
return TaskInstancePosition(in_loop=False, map_index=map_index)
if (task := getattr(task_instance, "task", None)) is None:
return None
# A worker parses a bundle without versions as it is when the attempt starts, which may no longer have the
# task in the loop of the Dag version the attempt pins, while Airflow places the attempt with that version
in_loop = _find_enclosing_loop(task)[0] or getattr(task_instance, "loop", None) is not None
loop_index: int | None = None
if is_stored and not task.get_needs_expansion():
# The stored region_index means a map index only for a mapped task. A task that isn't mapped has none,
# and inside a loop its region_index is the iteration instead.
if in_loop:
loop_index = map_index
map_index = -1
return TaskInstancePosition(in_loop=in_loop, map_index=map_index, loop_index=loop_index)
def _require_task_instance_position(
task_instance: TaskInstance | RuntimeTaskInstance,
) -> TaskInstancePosition:
if (position := _get_task_instance_position(task_instance)) is None:
raise ValueError(
f"Task instance {task_instance.id} has no task attached, so only its stored region can place it; "
"use get_orm_task_instance_positions with the session it was loaded in"
)
return position
def _add_runtime_loop_iteration(
task_instance: TaskInstance | RuntimeTaskInstance, position: TaskInstancePosition
) -> TaskInstancePosition:
if isinstance(task_instance, TaskInstance):
return position
# Only runtime task instances of Airflow 3.4+ expose these
position = attrs.evolve(position, expanded_ti_count=getattr(task_instance, "expanded_ti_count", None))
if not position.in_loop or position.loop_index is not None:
return position
if (loop := getattr(task_instance, "loop", None)) is None:
return position
return attrs.evolve(
position,
loop_index=loop.index,
loop_region_id=loop.region_id,
loop_trajectory=loop.trajectory,
loop_forked_from_region_id=loop.forked_from_region_id,
loop_resumes_from_index=loop.resumes_from_index,
)
[docs]
def get_task_instance_position(task_instance: TaskInstance | RuntimeTaskInstance) -> TaskInstancePosition:
"""Place a task instance with its loop iteration, once per event, to pass to ``get_airflow_run_facet``."""
return _add_runtime_loop_iteration(task_instance, _require_task_instance_position(task_instance))
def _holds_another_dag_version(task_instance: TaskInstance) -> bool:
"""Tell whether a stored task instance in a region holds its task from another Dag version than its own."""
region_id = getattr(task_instance, "region_id", None)
if not isinstance(region_id, UUID) or region_id.int == 0:
return False
dag = getattr(getattr(task_instance, "task", None), "dag", None)
held_dag_version_id = getattr(dag, "dag_version_id", None)
return held_dag_version_id is not None and held_dag_version_id != task_instance.dag_version_id
[docs]
def get_orm_task_instance_positions(
task_instances: Collection[TaskInstance], *, session: Session | None
) -> dict[UUID, TaskInstancePosition]:
"""
Place stored task instances, reading the regions of those without their task in one query.
With a session, an instance whose task comes from another Dag version than its own, as the scheduler holds
the latest one for a Dag edited since, is placed from its region too.
Server side only: workers never hold a session, and an instance that needs its region read without one
raises instead of opening a connection.
"""
positions: dict[UUID, TaskInstancePosition] = {}
unplaced: list[TaskInstance] = []
for task_instance in task_instances:
position = None
# Another version may have moved the task in or out of a loop, while its region rows can't have changed
if session is None or not _holds_another_dag_version(task_instance):
position = _get_task_instance_position(task_instance)
if position is None:
unplaced.append(task_instance)
else:
positions[task_instance.id] = position
if not unplaced:
return positions
if session is None:
raise ValueError("Placing a stored task instance without its task needs the session it was loaded in")
from sqlalchemy import select
from airflow.models.dynamic_region import DynamicRegion
regions = {
row.id: row
for row in session.execute(
select(DynamicRegion.id, DynamicRegion.node_id, DynamicRegion.parent_region_id).where(
DynamicRegion.id.in_({task_instance.region_id for task_instance in unplaced})
)
)
}
for task_instance in unplaced:
region = regions.get(task_instance.region_id)
if region is not None and region.node_id == task_instance.task_id:
positions[task_instance.id] = TaskInstancePosition(
in_loop=region.parent_region_id is not None, map_index=task_instance.region_index
)
else:
positions[task_instance.id] = TaskInstancePosition(in_loop=True, map_index=-1)
return positions
[docs]
def build_task_instance_ol_run_id(
dag_id: str,
task_id: str,
try_number: int,
logical_date: datetime.datetime,
map_index: int,
):
"""
Generate a deterministic OpenLineage run ID for a task instance.
Args:
dag_id: The DAG identifier.
task_id: The task identifier.
try_number: The task try number.
logical_date: The logical execution date from dagrun.
map_index: The task map index.
Returns:
A deterministic OpenLineage run ID for the task instance.
"""
return str(
generate_static_uuid(
instant=logical_date,
data=f"{conf.namespace()}.{dag_id}.{task_id}.{try_number}.{map_index}".encode(),
)
)
[docs]
def get_task_instance_run_id(
task_instance: TaskInstance | RuntimeTaskInstance,
*,
logical_date: datetime.datetime,
position: TaskInstancePosition | None = None,
) -> str:
"""
Return the OpenLineage ``runId`` of a task instance's attempt, which all its events must share.
In a loop the iterations share every coordinate but the attempt UUID, so that is the ``runId``. Otherwise
it is built from the coordinates.
"""
if position is None:
position = _require_task_instance_position(task_instance)
if position.in_loop:
return str(task_instance.id)
return build_task_instance_ol_run_id(
dag_id=task_instance.dag_id,
task_id=task_instance.task_id,
try_number=task_instance.try_number,
logical_date=logical_date,
map_index=position.map_index,
)
[docs]
def is_valid_uuid(uuid_string: str | None) -> bool:
"""Validate that a string is a valid UUID format."""
if uuid_string is None:
return False
try:
from uuid import UUID
UUID(uuid_string)
return True
except (ValueError, TypeError):
return False
[docs]
def build_dag_run_ol_run_id(dag_id: str, logical_date: datetime.datetime, clear_number: int) -> str:
"""
Generate a deterministic OpenLineage run ID for a DAG run.
Args:
dag_id: The DAG identifier.
logical_date: The logical execution date.
clear_number: The DAG run clear number.
Returns:
A deterministic OpenLineage run ID for the DAG run.
"""
return str(
generate_static_uuid(
instant=logical_date,
data=f"{conf.namespace()}.{dag_id}.{clear_number}".encode(),
)
)
def _get_eagerly_loaded_dagrun_consumed_asset_events(dag_id: str, dag_run_id: str) -> list[AssetEvent]:
"""
Retrieve consumed asset events for a DagRun with relationships eagerly loaded.
Downstream code accesses source_task_instance, source_dag_run, and asset on each AssetEvent.
These relationships are lazy-loaded by default, which could cause N+1 query problem
(2 + 3*N queries for N events). Using `joinedload` fetches everything in a single query.
The returned AssetEvent objects have all needed relationships pre-populated in memory,
so they can be safely used after the session is closed.
Returns:
AssetEvent objects with populated relationships, or empty list if DagRun not found.
"""
# This should only be used on scheduler, so DB access is allowed
from sqlalchemy import select
from sqlalchemy.orm import joinedload
from airflow.utils.session import create_session
if AIRFLOW_V_3_0_PLUS:
from airflow.models.asset import AssetEvent
options = (
joinedload(DagRun.consumed_asset_events).joinedload(AssetEvent.source_dag_run),
joinedload(DagRun.consumed_asset_events).joinedload(AssetEvent.source_task_instance),
joinedload(DagRun.consumed_asset_events).joinedload(AssetEvent.asset),
)
else: # AF2 path
from airflow.models.dataset import DatasetEvent
options = (
joinedload(DagRun.consumed_dataset_events).joinedload(DatasetEvent.source_dag_run),
joinedload(DagRun.consumed_dataset_events).joinedload(DatasetEvent.source_task_instance),
joinedload(DagRun.consumed_dataset_events).joinedload(DatasetEvent.dataset),
)
with create_session() as session:
dag_run_with_events = session.scalar(
select(DagRun).where(DagRun.dag_id == dag_id).where(DagRun.run_id == dag_run_id).options(*options)
)
if not dag_run_with_events:
return []
if AIRFLOW_V_3_0_PLUS:
events = dag_run_with_events.consumed_asset_events
else: # AF2 path
events = dag_run_with_events.consumed_dataset_events
return events
def _place_asset_event_sources(events: list[AssetEvent]) -> dict[UUID, TaskInstancePosition]:
"""Place the source task instances of consumed asset events; only the scheduler builds these facets."""
sources = [ti for event in events if isinstance(ti := event.source_task_instance, TaskInstance)]
positions = {ti.id: position for ti in sources if (position := _get_task_instance_position(ti))}
if unplaced := [ti for ti in sources if ti.id not in positions]:
from airflow.utils.session import create_session
with create_session() as session:
positions.update(get_orm_task_instance_positions(unplaced, session=session))
return positions
def _extract_ol_info_from_asset_event(
asset_event: AssetEvent, positions: Mapping[UUID, TaskInstancePosition] | None = None
) -> dict[str, str] | None:
"""
Extract OpenLineage job information from an AssetEvent.
Information is gathered from multiple potential sources, checked in priority
order:
1. TaskInstance (primary): Provides the most complete and reliable context.
2. AssetEvent source fields (fallback): Offers basic `dag_id.task_id` metadata.
3. `asset_event.extra["openlineage"]` (last resort): May include user provided OpenLineage details.
Args:
asset_event: The AssetEvent from which to extract job information.
Returns:
A dictionary containing `job_name`, `job_namespace`, and optionally
`run_id`, or `None` if insufficient information is available.
"""
ti = asset_event.source_task_instance
if ti:
result = {
"job_name": get_job_name(ti),
"job_namespace": conf.namespace(),
}
position = (positions or {}).get(ti.id) or _require_task_instance_position(ti)
if position.in_loop:
result["run_id"] = str(ti.id)
return result
source_dr = asset_event.source_dag_run
if source_dr:
logical_date = source_dr.logical_date # Get logical date from DagRun for OL run_id generation
if AIRFLOW_V_3_0_PLUS and logical_date is None:
logical_date = source_dr.run_after
if logical_date is not None:
result["run_id"] = build_task_instance_ol_run_id(
dag_id=ti.dag_id,
task_id=ti.task_id,
try_number=ti.try_number,
logical_date=logical_date,
map_index=position.map_index,
)
return result
# Then, check AssetEvent source_* fields
if asset_event.source_dag_id and asset_event.source_task_id:
return {
"job_name": f"{asset_event.source_dag_id}.{asset_event.source_task_id}",
"job_namespace": conf.namespace(),
# run_id cannot be constructed from these fields alone
}
# Lastly, check asset_event.extra["openlineage"]
if asset_event.extra:
ol_info_from_extra = asset_event.extra.get("openlineage")
if isinstance(ol_info_from_extra, dict):
job_name = ol_info_from_extra.get("parentJobName")
job_namespace = ol_info_from_extra.get("parentJobNamespace")
run_id = ol_info_from_extra.get("parentRunId")
if job_name and job_namespace:
result = {
"job_name": str(job_name),
"job_namespace": str(job_namespace),
}
if run_id:
if not is_valid_uuid(str(run_id)):
log.warning(
"Invalid runId in AssetEvent.extra; ignoring value. event_id=%s, run_id=%s",
asset_event.id,
run_id,
)
else:
result["run_id"] = str(run_id)
return result
return None
def _get_ol_job_dependencies_from_asset_events(events: list[AssetEvent]) -> list[dict[str, Any]]:
"""
Extract and deduplicate OpenLineage job dependencies from asset events.
This function processes a list of asset events, extracts OpenLineage dependency information
from all relevant sources, and deduplicates the results based on the tuple (job_namespace, job_name, run_id)
to prevent emitting duplicate dependencies. Multiple asset events from the same job but different
source runs/assets are aggregated into a single dependency entry with all source information preserved.
Args:
events: List of AssetEvent objects to process.
Returns:
A list of deduplicated dictionaries containing OpenLineage job dependency information.
Each dictionary includes job_name, job_namespace, optional run_id, and an asset_events
list containing source information from all aggregated events.
"""
# Use a dictionary keyed by (namespace, job_name, run_id) to deduplicate
# Multiple asset events from the same task instance should only create one dependency
deduplicated_jobs: dict[tuple[str, str, str | None], dict[str, Any]] = {}
positions = _place_asset_event_sources(events)
for asset_event in events:
# Extract OpenLineage information
ol_info = _extract_ol_info_from_asset_event(asset_event, positions)
# Skip if we don't have minimum required info (job_name and namespace)
if not ol_info:
log.debug(
"Insufficient OpenLineage information, skipping asset event: %s",
str(asset_event),
)
continue
# Create deduplication key: (namespace, job_name, run_id)
# We deduplicate on job identity (namespace + name + run_id), not on source dag_run_id
# Multiple asset events from the same job but different source runs/assets are aggregated
dedup_key = (
ol_info["job_namespace"],
ol_info["job_name"],
ol_info.get("run_id"),
)
# Collect source information for this asset event
source_info = {
"dag_run_id": asset_event.source_run_id,
"asset_event_id": asset_event.id,
"asset_event_extra": asset_event.extra or None,
"asset_id": asset_event.asset_id if AIRFLOW_V_3_0_PLUS else asset_event.dataset_id,
"asset_uri": asset_event.uri,
"partition_key": getattr(asset_event, "partition_key", None),
}
if dedup_key not in deduplicated_jobs:
# First occurrence: create the job entry with initial source info
deduplicated_jobs[dedup_key] = {**ol_info, "asset_events": [source_info]}
else:
# Already seen: append source info to existing entry
deduplicated_jobs[dedup_key]["asset_events"].append(source_info)
result = list(deduplicated_jobs.values())
return result
def _build_job_dependency_facet(
dag_id: str, dag_run_id: str
) -> dict[str, job_dependencies_run.JobDependenciesRunFacet]:
"""
Build the JobDependenciesRunFacet for a DagRun.
Args:
dag_id: The DAG identifier.
dag_run_id: The DagRun identifier.
Returns:
A dictionary containing the JobDependenciesRunFacet, or an empty dictionary.
"""
log.info(
"Building OpenLineage JobDependenciesRunFacet for DagRun(dag_id=%s, run_id=%s).",
dag_id,
dag_run_id,
)
events = _get_eagerly_loaded_dagrun_consumed_asset_events(dag_id, dag_run_id)
if not events:
log.info("DagRun %s/%s has no consumed asset events", dag_id, dag_run_id)
return {}
ol_dependencies = _get_ol_job_dependencies_from_asset_events(events=events)
if not ol_dependencies:
log.info(
"No OpenLineage job dependencies generated from asset events consumed by DagRun %s/%s.",
dag_id,
dag_run_id,
)
return {}
upstream_dependencies = []
for job in ol_dependencies:
job_identifier = job_dependencies_run.JobIdentifier(
namespace=job["job_namespace"],
name=job["job_name"],
)
run_identifier = None
if job.get("run_id"):
run_identifier = job_dependencies_run.RunIdentifier(runId=job["run_id"])
job_dependency = job_dependencies_run.JobDependency(
job=job_identifier,
run=run_identifier,
dependency_type="IMPLICIT_ASSET_DEPENDENCY",
).with_additional_properties(airflow={"asset_events": job.get("asset_events")}) # type: ignore[arg-type] # Fixed in OL client 1.42, waiting for release
upstream_dependencies.append(job_dependency)
return {
"jobDependencies": job_dependencies_run.JobDependenciesRunFacet(
upstream=upstream_dependencies,
)
}
[docs]
def get_dag_job_dependency_facet(
dag_id: str, dag_run_id: str
) -> dict[str, job_dependencies_run.JobDependenciesRunFacet]:
"""
Safely retrieve the asset-triggered job dependency facet for a DagRun.
This function collects information about the asset events that triggered the specified DagRun,
including details about the originating DAG runs and task instances. If the DagRun was not triggered
by assets, or if any error occurs during lookup or processing, the function logs the error and returns
an empty dictionary. This guarantees that facet generation never raises exceptions and does not
interfere with event emission processes.
Args:
dag_id: The DAG identifier.
dag_run_id: The DagRun identifier.
Returns:
A dictionary with JobDependenciesRunFacet, or an empty dictionary
if the DagRun was not asset-triggered or if an error occurs.
"""
try:
return _build_job_dependency_facet(dag_id=dag_id, dag_run_id=dag_run_id)
except Exception as e:
log.warning("Failed to build JobDependenciesRunFacet for DagRun %s/%s: %s.", dag_id, dag_run_id, e)
log.debug("Exception details:", exc_info=True)
return {}
def _build_labeled_edge_map(dag: DAG | SerializedDAG) -> dict[str, tuple[dict, dict]]:
"""
Build a mapping of classified labeled edges for every task and group in the DAG.
Translates Airflow's internal ``dag.edge_info`` (which uses virtual
``downstream_join_id`` / ``upstream_join_id`` keys) into a clean mapping
keyed by real ``task_id`` or ``group_id``. Each value is a tuple of
``(downstream_task_edges, downstream_group_edges)`` sorted dicts.
``downstream_task_edges``
Target is a **specific task**. Keyed by ``task_id``, regardless of
whether source and target are in the same group, different groups,
or at root level.
``downstream_group_edges``
Target is a **TaskGroup as a whole**. Keyed by ``group_id``.
Only produced by ``task >> TaskGroup`` or ``TaskGroup >> TaskGroup``.
**How to read the result (for OL event consumers)**
Look up ``edge_map.get(task_id)`` or ``edge_map.get(group_id)``.
Missing keys mean no labeled edges from that source. Unlabeled edges
are not included — use ``downstream_task_ids`` for the full dependency
graph.
Key points:
- Cross-group task-to-task labels (e.g. ``step_2 >> Label >> final``
where the tasks are in different groups) appear under the **source
group's** key, not the originating task's key, because Airflow promotes
labels to the group boundary.
- ``task >> TaskGroup`` produces entries in **both** dicts: per-root-task
entries in ``downstream_task_edges`` and a group-level entry in
``downstream_group_edges``, all with the same label.
See ``test_get_tasks_details_with_edge_labels`` for a comprehensive example
covering all edge patterns (same-group, cross-group, task-to-TaskGroup,
TaskGroup-to-TaskGroup, nested groups, etc.).
"""
edge_info = dag.edge_info or {}
if not edge_info:
return {}
upstream_join_id_to_group_id: dict[str, str] = {}
downstream_join_id_to_group_id: dict[str, str] = {}
for tg_id, tg in dag.task_group_dict.items():
upstream_join_id_to_group_id[tg.upstream_join_id] = tg_id
downstream_join_id_to_group_id[tg.downstream_join_id] = tg_id
result: dict[str, tuple[dict, dict]] = {}
for source_id, edges in edge_info.items():
# Resolve downstream_join_id sources to their owning group_id
group_id = downstream_join_id_to_group_id.get(source_id)
logical_key = group_id if group_id is not None else source_id
task_edges: dict = {}
group_edges: dict = {}
for target_id, info in sorted(edges.items()):
resolved_group = upstream_join_id_to_group_id.get(target_id)
if resolved_group is not None: # If target is a group, classify as group edge
group_edges[resolved_group] = info
else:
task_edges[target_id] = info
result[logical_key] = (task_edges, group_edges)
return result
def _get_tasks_details(dag: DAG | SerializedDAG, edge_map: dict[str, tuple[dict, dict]]) -> dict:
tasks = {}
for single_task in sorted(dag.tasks, key=lambda x: x.task_id):
task_edges, group_edges = edge_map.get(single_task.task_id, ({}, {}))
task_loop = get_task_loop_info(single_task)
get_needs_expansion = getattr(single_task, "get_needs_expansion", None)
is_mapped = bool(get_needs_expansion()) if callable(get_needs_expansion) else None
trigger_rule = getattr(single_task, "trigger_rule", None)
tasks[single_task.task_id] = {
"operator": get_fully_qualified_class_name(single_task),
"task_group": single_task.task_group.group_id if single_task.task_group else None,
# The loop around the task, which is not task_group for a task in a task group inside the loop
"loop_group": task_loop["group_id"] if task_loop else None,
# Expanded by itself or by a mapped task group around it
"is_mapped": is_mapped,
"parse_time_mapped_count": _get_parse_time_mapped_count(single_task) if is_mapped else None,
"trigger_rule": getattr(trigger_rule, "value", trigger_rule),
"depends_on_past": getattr(single_task, "depends_on_past", None),
"retries": getattr(single_task, "retries", None),
"emits_ol_events": _emits_ol_events(single_task),
"ui_color": single_task.ui_color,
"ui_fgcolor": single_task.ui_fgcolor,
"ui_label": single_task.label,
"is_setup": single_task.is_setup,
"is_teardown": single_task.is_teardown,
"downstream_task_ids": sorted(single_task.downstream_task_ids),
"downstream_task_edges": task_edges,
"downstream_group_edges": group_edges,
}
return tasks
def _get_task_groups_details(dag: DAG | SerializedDAG, edge_map: dict[str, tuple[dict, dict]]) -> dict:
result = {}
for tg_id, tg in dag.task_group_dict.items():
task_edges, group_edges = edge_map.get(tg_id, ({}, {}))
is_mapped = _is_mapped_task_group(tg)
is_loop = _is_loop_task_group(tg)
result[tg_id] = {
"parent_group": tg.parent_group.group_id if tg.parent_group else None,
# Independent, so a group could be both
"is_mapped": is_mapped,
"is_loop": is_loop,
"parse_time_mapped_count": _get_parse_time_mapped_count(tg) if is_mapped else None,
"ui_color": tg.ui_color,
"ui_fgcolor": tg.ui_fgcolor,
"ui_label": tg.label,
"downstream_task_edges": task_edges,
"downstream_group_edges": group_edges,
**{
field: getattr(tg, field)
for field in ("gate_task_id", "has_until", "max_iterations", "terminal_task_id")
if is_loop and hasattr(tg, field)
},
}
return result
def _is_mapped_task_group(task_group: Any) -> bool:
if (is_mapped := getattr(task_group, "is_mapped", None)) is not None:
return bool(is_mapped)
# Only the scheduler's serialized task groups tell it themselves
return isinstance(task_group, MappedTaskGroup)
def _get_parse_time_mapped_count(node: Any) -> int | None:
"""Return how many instances a mapped task or task group expands to, when it's known before the Dag runs."""
# The scheduler's serialized Dag and Airflow 2's count, while a Task SDK Dag doesn't
if (get_count := getattr(node, "get_parse_time_mapped_ti_count", None)) is None:
return None
try:
return get_count()
except RuntimeError: # NotFullyPopulated, as the size depends on the results of upstream tasks
return None
def _emits_ol_events(task: AnyOperator) -> bool:
from airflow.providers.openlineage.utils.emission_policy import resolve_task_emission_policy
# resolve_task_emission_policy already incorporates the selective_enable check.
controls = resolve_task_emission_policy(
operator=task,
dag_id=task.dag_id,
task_id=task.task_id,
)
is_task_schedulable_method = getattr(TaskInstance, "is_task_schedulable", None) # Added in 3.2.0 #56039
if is_task_schedulable_method and callable(is_task_schedulable_method):
is_skipped_as_empty_operator = not is_task_schedulable_method(task)
else:
# For older Airflow versions, re-create Airflow core internal logic as
# empty operators without callbacks/outlets are skipped for optimization by Airflow
# in airflow.models.taskinstance.TaskInstance._schedule_downstream_tasks or
# airflow.models.dagrun.DagRun.schedule_tis, depending on Airflow version
is_skipped_as_empty_operator = all(
(
task.inherits_from_empty_operator,
not getattr(task, "on_execute_callback", None),
not getattr(task, "on_success_callback", None),
not task.outlets,
not (task.inlets and get_base_airflow_version_tuple() >= (3, 0, 2)), # Added in 3.0.2 #50773
not (
getattr(task, "has_on_execute_callback", None) # Added in 3.1.0 #54569
and get_base_airflow_version_tuple() >= (3, 1, 0)
),
not (
getattr(task, "has_on_success_callback", None) # Added in 3.1.0 #54569
and get_base_airflow_version_tuple() >= (3, 1, 0)
),
)
)
emits_ol_events = all(
(
controls.emit,
not is_skipped_as_empty_operator,
)
)
return emits_ol_events
[docs]
def get_unknown_source_attribute_run_facet(task: BaseOperator, name: str | None = None):
if not name:
name = get_operator_class(task).__name__
log.debug(
"UnknownOperatorAttributeRunFacet is deprecated and will be removed. "
"Use information from AirflowRunFacet instead."
)
return {
"unknownSourceAttribute": attrs.asdict(
UnknownOperatorAttributeRunFacet(
unknownItems=[
UnknownOperatorInstance(
name=name,
properties=TaskInfo(task),
)
]
)
)
}
[docs]
class OpenLineageRedactor(SecretsMasker):
"""
This class redacts sensitive data similar to SecretsMasker in Airflow logs.
The difference is that our default max recursion depth is way higher - due to
the structure of OL events we need more depth.
Additionally, we allow data structures to specify data that needs not to be
redacted by specifying _skip_redact list by deriving RedactMixin.
"""
@classmethod
[docs]
def from_masker(cls, other: SecretsMasker) -> OpenLineageRedactor:
instance = cls()
instance.patterns = other.patterns
instance.replacer = other.replacer
for attr in ["sensitive_variables_fields", "min_length_to_mask", "secret_mask_adapter"]:
if hasattr(other, attr):
setattr(instance, attr, getattr(other, attr))
return instance
def _should_hide_value_for_key(self, name):
"""Compatibility helper for should_hide_value_for_key across Airflow versions."""
try:
return self.should_hide_value_for_key(name)
except AttributeError:
# fallback to module level function
return should_hide_value_for_key(name)
def _redact(self, item: Redactable, name: str | None, depth: int, max_depth: int, **kwargs) -> Redacted: # type: ignore[override]
if AIRFLOW_V_3_0_PLUS:
# Keep compatibility for Airflow 2.x, remove when Airflow 3.0 is the minimum version
class AirflowContextDeprecationWarning(UserWarning):
pass
else:
from airflow.utils.context import ( # type: ignore[attr-defined,no-redef]
AirflowContextDeprecationWarning,
)
if depth > max_depth:
return item
try:
# It's impossible to check the type of variable in a dict without accessing it, and
# this already causes warning - so suppress it
with suppress(AirflowContextDeprecationWarning):
if type(item).__name__ == "Proxy":
# Those are deprecated values in _DEPRECATION_REPLACEMENTS
# in airflow.utils.context.Context
return "<<non-redactable: Proxy>>"
if name and self._should_hide_value_for_key(name):
return self._redact_all(item, depth, max_depth)
if attrs.has(type(item)):
# TODO: FIXME when mypy gets compatible with new attrs
for dict_key, subval in attrs.asdict(
item, # type: ignore[arg-type]
recurse=False,
).items():
if _is_name_redactable(dict_key, item):
setattr(
item,
dict_key,
self._redact(
subval,
name=dict_key,
depth=(depth + 1),
max_depth=max_depth,
),
)
return item
# Dicts go to the base masker, which walks their items; the attributes of a dict
# subclass (e.g. InfoJsonEncodable) aren't what gets emitted.
if is_json_serializable(item) and hasattr(item, "__dict__") and not isinstance(item, dict):
for dict_key, subval in item.__dict__.items():
if type(subval).__name__ == "Proxy":
return "<<non-redactable: Proxy>>"
if _is_name_redactable(dict_key, item):
setattr(
item,
dict_key,
self._redact(
subval,
name=dict_key,
depth=(depth + 1),
max_depth=max_depth,
),
)
return item
return super()._redact(item, name, depth, max_depth, **kwargs)
except Exception as exc:
log.warning("Unable to redact %r. Error was: %s: %s", item, type(exc).__name__, exc)
return item
def _get_json_native_form(value: Any) -> Any:
# Attrs objects become dicts so the redactor can still mask fields by their names, e.g. ``password``.
# Other objects are not cast with ``str()``, as their ``repr`` may show such fields unmasked.
if attrs.has(type(value)):
return attrs.asdict(value, recurse=False)
if isinstance(value, Enum):
return value.value
if isinstance(value, (set, frozenset)):
return list(value)
if isinstance(value, (datetime.date, datetime.time, datetime.timedelta, Decimal, UUID)):
return str(value)
return f"<<non-serializable: {type(value).__qualname__}>>"
def _convert_to_json_safe(value: Any, name: str) -> Any:
"""
Return ``value`` with only JSON-native types.
Non-serializable values are converted with ``_get_json_native_form``; dict keys JSON can't encode
(e.g. tuples) are dropped.
"""
try:
return json.loads(json.dumps(value, default=_get_json_native_form, skipkeys=True))
except (TypeError, ValueError) as e:
log.warning("Unable to convert `%s` to JSON-native types: %s", name, e)
return None
[docs]
def is_json_serializable(item):
try:
json.dumps(item)
return True
except (TypeError, ValueError):
return False
def _is_name_redactable(name, redacted):
if not issubclass(redacted.__class__, RedactMixin):
return not name.startswith("_")
return name not in redacted.skip_redact
[docs]
def print_warning(log):
def decorator(f):
@wraps(f)
def wrapper(*args, **kwargs):
try:
return f(*args, **kwargs)
except Exception:
log.warning(
"OpenLineage event emission failed. "
"Exception below is being caught but it's printed for visibility. "
"This has no impact on actual task execution status.",
exc_info=True,
)
return wrapper
return decorator
[docs]
def get_filtered_unknown_operator_keys(operator: BaseOperator) -> dict:
not_required_keys = {"dag", "task_group"}
return {attr: value for attr, value in operator.__dict__.items() if attr not in not_required_keys}
[docs]
def should_use_external_connection(hook) -> bool:
# If we're at Airflow 2.10, the execution is process-isolated, so we can safely run those again.
return True
def _get_providers_manager_instance():
"""
Return the ProvidersManager instance to use for OpenLineage converter lookup.
Picks the task-runtime-safe entry point on Airflow 3.2+ and falls back to the legacy
``airflow.providers_manager.ProvidersManager``.
Kept as a module-level function so tests can patch this single seam with a plain
``mock.patch(..., return_value=fake_pm)`` instead of juggling ``sys.modules`` or
dotted paths that may not exist on the target Airflow version.
"""
try:
# Task-runtime-safe entry point — only available in Airflow 3.2+.
from airflow.sdk.plugins_manager import ProvidersManagerTaskRuntime
return ProvidersManagerTaskRuntime()
except (ImportError, ModuleNotFoundError, AttributeError):
from airflow.providers_manager import ProvidersManager
return ProvidersManager()
[docs]
def translate_airflow_asset(asset: Asset, lineage_context) -> OpenLineageDataset | None:
"""
Convert an Asset with an AIP-60 compliant URI to an OpenLineageDataset.
This function returns None if no URI normalizer is defined, no asset converter is found or
some core Airflow changes are missing and ImportError is raised.
"""
if AIRFLOW_V_3_0_PLUS:
from airflow.sdk.definitions.asset import _get_normalized_scheme
else:
try:
from airflow.datasets import _get_normalized_scheme # type: ignore[no-redef]
except ImportError:
log.debug("Could not import objects required for asset translation.")
return None
providers_manager = _get_providers_manager_instance()
if providers_manager is None:
log.debug("Could not resolve a ProvidersManager instance for asset translation.")
return None
ol_converters = getattr(providers_manager, "asset_to_openlineage_converters", None)
if not ol_converters:
ol_converters = getattr(providers_manager, "dataset_to_openlineage_converters", None)
if not ol_converters:
log.debug("Could not retrieve openlineage converters.")
return None
if not (normalized_uri := getattr(asset, "normalized_uri", None)):
log.debug("Normalized uri is not accessible or empty.")
return None
if not (normalized_scheme := _get_normalized_scheme(normalized_uri)):
log.debug("Normalized scheme is empty for URI `%s`.", normalized_uri)
return None
if (airflow_to_ol_converter := ol_converters.get(normalized_scheme)) is None:
log.debug("No converter found for scheme `%s`.", normalized_scheme)
return None
try:
return airflow_to_ol_converter(Asset(uri=normalized_uri, extra=asset.extra), lineage_context)
except Exception as err:
log.debug("Airflow asset `%s` conversion to OL Dataset failed with err `%s`", normalized_uri, err)
return None