This is an automated email from the ASF dual-hosted git repository.
hussein-awala pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new d561f37ea54 Fix OpenLineage task events lost for Dags with mapped
tasks (#75860)
d561f37ea54 is described below
commit d561f37ea5471f4e67bd7e5924aea425ded0bec4
Author: Kacper Muda <[email protected]>
AuthorDate: Thu Oct 29 09:18:13 2026 +0100
Fix OpenLineage task events lost for Dags with mapped tasks (#75860)
* Fix OpenLineage task events lost for Dags with mapped tasks
The listener pickles its facets to send them to a process pool. Info objects
kept the serialized Dag and other build-time objects as attributes, and a
serialized Dag with a mapped operator can't be pickled, so task events
emitted
from the scheduler or API server (e.g. FAIL after an executor-reported
failure)
were silently dropped. Full task info also included the stored operator's
Dag
and methodtools caches, which can't be pickled either.
Co-Authored-By: Claude <[email protected]>
* Ignore test-defined InfoJsonEncodable subclasses in pickling test
Other tests define local subclasses, which stay listed by __subclasses__
when
they run first in the same process, as in the provider compatibility jobs.
Co-Authored-By: Claude <[email protected]>
* Keep only JSON-native values in OpenLineage full task info
Full task info copies operator attributes, and any of them may reach live
Airflow objects that can't be pickled. A mapped task fed from an upstream
task's output does, through its expand input, so its events were still
dropped when full task info was on.
Co-Authored-By: Claude <[email protected]>
* Redact OpenLineage info objects by their dict items
The redactor sent every JSON-serializable object with a `__dict__` to a
branch that only masks attributes. For `InfoJsonEncodable` dict subclasses
that skipped the emitted items, so full task info leaked rendered templated
fields and Dag run `conf` leaked secrets. Dicts now go to the base masker.
Operator attribute values that aren't JSON-native no longer become their
`repr`, which could show secret fields: attrs objects become dicts, so they
stay maskable by key name, and unknown objects get the client's
`<<non-serializable: ...>>` placeholder.
Co-Authored-By: Claude <[email protected]>
---------
Co-authored-by: Claude <[email protected]>
---
.../airflow/providers/openlineage/utils/utils.py | 61 ++++++++---
.../tests/unit/openlineage/plugins/test_utils.py | 48 +++++++--
.../tests/unit/openlineage/utils/test_utils.py | 118 +++++++++++++++++++--
3 files changed, 192 insertions(+), 35 deletions(-)
diff --git
a/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py
b/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py
index 320a324da8d..e7d6303bda3 100644
--- a/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py
+++ b/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py
@@ -22,6 +22,8 @@ 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
@@ -891,7 +893,9 @@ class InfoJsonEncodable(dict):
self,
**{field: InfoJsonEncodable._cast_basic_types(getattr(self,
field)) for field in self._fields},
)
- del self.obj
+ # 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):
@@ -951,8 +955,13 @@ class InfoJsonEncodable(dict):
"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:
- setattr(self, field, val)
+ # 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)
@@ -1491,20 +1500,12 @@ class TaskInfo(InfoJsonEncodable):
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)
- # ``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 after ``super().__init__`` 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_list(getattr(self, field), field)
def _extend_fields(self) -> None:
- if not self._runtime_assets:
- return
- if not hasattr(self, "outlets"): # Should never happen, just in case
+ 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:
+ 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(
@@ -1512,6 +1513,12 @@ class TaskInfo(InfoJsonEncodable):
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)
class TaskInfoComplete(TaskInfo):
@@ -1522,6 +1529,9 @@ class TaskInfoComplete(TaskInfo):
"_BaseOperator__instantiated",
"_BaseOperator__init_kwargs", # Causes recursion error on AF3,
nothing useful there
"_dag",
+ "dag",
+ "expand_input",
+ "op_kwargs_expand_input",
"_hook",
"_log",
"_outlets",
@@ -2540,7 +2550,9 @@ class OpenLineageRedactor(SecretsMasker):
),
)
return item
- if is_json_serializable(item) and hasattr(item, "__dict__"):
+ # 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>>"
@@ -2562,16 +2574,31 @@ class OpenLineageRedactor(SecretsMasker):
return item
-def _convert_to_json_safe_list(value: list, name: str) -> list | None:
+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 cast with ``str()``; dict keys JSON can't
encode (e.g. tuples) are dropped.
+ 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=str, skipkeys=True))
+ return json.loads(json.dumps(value, default=_get_json_native_form,
skipkeys=True))
except (TypeError, ValueError) as e:
- log.warning("Unable to serialize task `%s` to a JSON list: %s", name,
e)
+ log.warning("Unable to convert `%s` to JSON-native types: %s", name, e)
return None
diff --git a/providers/openlineage/tests/unit/openlineage/plugins/test_utils.py
b/providers/openlineage/tests/unit/openlineage/plugins/test_utils.py
index 88006cf74dc..02ab5a8f0d8 100644
--- a/providers/openlineage/tests/unit/openlineage/plugins/test_utils.py
+++ b/providers/openlineage/tests/unit/openlineage/plugins/test_utils.py
@@ -20,6 +20,7 @@ import datetime
import json
import uuid
from json import JSONEncoder
+from types import SimpleNamespace
from typing import Any
from unittest.mock import MagicMock, patch
@@ -33,8 +34,10 @@ from airflow.providers.common.compat.sdk import timezone
from airflow.providers.openlineage.plugins.facets import AirflowDebugRunFacet
from airflow.providers.openlineage.utils.utils import (
DagInfo,
+ DagRunInfo,
InfoJsonEncodable,
OpenLineageRedactor,
+ TaskInfoComplete,
_get_all_packages_installed,
_is_name_redactable,
get_airflow_debug_facet,
@@ -303,6 +306,31 @@ def test_redact_with_exclusions(monkeypatch):
assert redacted_nested == NestedMixined("***", NestedMixined("passwd",
None))
[email protected]_redact
+def test_redact_masks_items_of_info_json_encodables():
+ sm = SecretsMasker()
+ if AIRFLOW_V_3_1_PLUS:
+ sm.sensitive_variables_fields = list(DEFAULT_SENSITIVE_FIELDS)
+ sm.add_mask("s3cr3tvalue")
+ redactor = OpenLineageRedactor.from_masker(sm)
+
+ @define
+ class Creds:
+ user: str
+ password: str
+
+ task = BashOperator(task_id="t", bash_command="curl -u user:s3cr3tvalue x")
+ task.creds = Creds(user="u", password="unregistered")
+ dag_run = SimpleNamespace(
+ dag_id="d", run_id="r", conf={"password": "unregistered", "cmd": "x
s3cr3tvalue"}
+ )
+
+ redacted_task = redactor.redact(TaskInfoComplete(task))
+ assert redacted_task["bash_command"] == "curl -u user:*** x"
+ assert redacted_task["creds"] == {"user": "u", "password": "***"}
+ assert redactor.redact(DagRunInfo(dag_run))["conf"] == {"password": "***",
"cmd": "x ***"}
+
+
def test_get_fully_qualified_class_name():
from airflow.providers.openlineage.plugins.adapter import
OpenLineageAdapter
@@ -486,7 +514,7 @@ def test_serialize_timetable_complex_with_alias():
}
if AIRFLOW_V_3_4_PLUS:
expected["batch_asset_events"] = True
- assert dag_info.timetable == expected
+ assert dag_info["timetable"] == expected
@pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="This test checks
serialization only in 3.0 conditions")
@@ -504,7 +532,7 @@ def test_serialize_timetable_single_asset():
}
if AIRFLOW_V_3_4_PLUS:
expected["batch_asset_events"] = False
- assert dag_info.timetable == expected
+ assert dag_info["timetable"] == expected
@pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="This test checks
serialization only in 3.0 conditions")
@@ -522,7 +550,7 @@ def test_serialize_timetable_list_of_assets():
}
if AIRFLOW_V_3_4_PLUS:
expected["batch_asset_events"] = True
- assert dag_info.timetable == expected
+ assert dag_info["timetable"] == expected
@pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="This test checks
serialization only in 3.0 conditions")
@@ -581,7 +609,7 @@ def
test_serialize_timetable_with_complex_logical_condition():
}
if AIRFLOW_V_3_4_PLUS:
expected["batch_asset_events"] = True
- assert dag_info.timetable == expected
+ assert dag_info["timetable"] == expected
@pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="This test checks
serialization only in 3.0 conditions")
@@ -655,7 +683,7 @@ def
test_serialize_timetable_with_dataset_or_time_schedule():
}
if AIRFLOW_V_3_4_PLUS:
expected["batch_asset_events"] = True
- assert dag_info.timetable == expected
+ assert dag_info["timetable"] == expected
@pytest.mark.skipif(
@@ -677,7 +705,7 @@ def test_serialize_timetable_2_10_complex_with_alias():
dag.timetable = DatasetTriggeredTimetable(asset)
dag_info = DagInfo(dag)
- assert dag_info.timetable == {
+ assert dag_info["timetable"] == {
"dataset_condition": {
"__type": DagAttributeTypes.DATASET_ANY,
"objects": [
@@ -703,7 +731,7 @@ def test_serialize_timetable_2_10_complex_with_alias():
def test_serialize_timetable_2_10_single_asset():
dag = DAG(dag_id="test", start_date=datetime.datetime(2025, 1, 1),
schedule=Asset("a"))
dag_info = DagInfo(dag)
- assert dag_info.timetable == {
+ assert dag_info["timetable"] == {
"dataset_condition": {"__type": DagAttributeTypes.DATASET, "uri": "a",
"extra": None}
}
@@ -715,7 +743,7 @@ def test_serialize_timetable_2_10_single_asset():
def test_serialize_timetable_2_10_list_of_assets():
dag = DAG(dag_id="test", start_date=datetime.datetime(2025, 1, 1),
schedule=[Asset("a"), Asset("b")])
dag_info = DagInfo(dag)
- assert dag_info.timetable == {
+ assert dag_info["timetable"] == {
"dataset_condition": {
"__type": DagAttributeTypes.DATASET_ALL,
"objects": [
@@ -738,7 +766,7 @@ def
test_serialize_timetable_2_10_with_complex_logical_condition():
& (Asset("ds3") | Asset("ds4", extra={"another_extra": 345})),
)
dag_info = DagInfo(dag)
- assert dag_info.timetable == {
+ assert dag_info["timetable"] == {
"dataset_condition": {
"__type": DagAttributeTypes.DATASET_ALL,
"objects": [
@@ -779,7 +807,7 @@ def
test_serialize_timetable_2_10_with_dataset_or_time_schedule():
),
)
dag_info = DagInfo(dag)
- assert dag_info.timetable == {
+ assert dag_info["timetable"] == {
"timetable": {
"__type": "airflow.timetables.trigger.CronTriggerTimetable",
"__var": {"expression": "0 0 * 3 *", "timezone": "UTC",
"interval": 0.0},
diff --git a/providers/openlineage/tests/unit/openlineage/utils/test_utils.py
b/providers/openlineage/tests/unit/openlineage/utils/test_utils.py
index 9c9518b97bc..756d86fc7e0 100644
--- a/providers/openlineage/tests/unit/openlineage/utils/test_utils.py
+++ b/providers/openlineage/tests/unit/openlineage/utils/test_utils.py
@@ -20,12 +20,15 @@ from __future__ import annotations
import datetime
import enum
import pathlib
+import pickle
+import threading
from decimal import Decimal
from types import SimpleNamespace
from unittest import mock
from unittest.mock import MagicMock, PropertyMock, patch
from uuid import UUID, uuid4
+import attrs
import pendulum
import pytest
from openlineage.client.facet_v2 import parent_run
@@ -51,7 +54,7 @@ from airflow.providers.openlineage.utils.utils import (
TaskInstanceInfo,
TaskInstancePosition,
_build_labeled_edge_map,
- _convert_to_json_safe_list,
+ _convert_to_json_safe,
_extract_ol_info_from_asset_event,
_get_ol_job_dependencies_from_asset_events,
_get_openlineage_data_from_dagrun_conf,
@@ -2615,11 +2618,11 @@ def test_daginfo_timetable_summary():
# timetable is enough to get summary
dag.timetable = NullTimetable()
dag.timetable_summary = None
- assert DagInfo(dag).timetable_summary == "None"
+ assert DagInfo(dag)["timetable_summary"] == "None"
# but if summary is present, it's preferred
dag.timetable_summary = "explicit_summary"
- assert DagInfo(dag).timetable_summary == "explicit_summary"
+ assert DagInfo(dag)["timetable_summary"] == "explicit_summary"
@pytest.mark.skipif(AIRFLOW_V_3_0_PLUS, reason="Airflow 2 tests")
@@ -4642,6 +4645,87 @@ def
test_task_instance_info_of_stored_instance_without_task_does_not_raise():
assert "loop" not in info
+def _get_info_json_encodable_subclasses(cls: type[InfoJsonEncodable] =
InfoJsonEncodable):
+ for subclass in cls.__subclasses__():
+ # Leave out the subclasses other tests define
+ if subclass.__module__ == InfoJsonEncodable.__module__:
+ yield subclass
+ yield from _get_info_json_encodable_subclasses(subclass)
+
+
[email protected]_test
+def
test_info_json_encodables_of_stored_objects_pickle_to_their_items_only(dag_maker):
+ """
+ The scheduler and API server send facets to a process pool, while what
they are built from may not pickle.
+
+ A serialized Dag with a mapped operator doesn't, so no encodable may keep
it or any other built-from object.
+ """
+ asset = Asset("s3://bucket/key")
+ with dag_maker(serialized=True):
+ with TaskGroup("group"):
+ plain = EmptyOperator(task_id="plain", inlets=[asset],
outlets=[asset])
+ plain >> PythonOperator.partial(task_id="mapped",
python_callable=list).expand(op_kwargs=[{}, {}])
+ PythonOperator.partial(task_id="mapped_from_upstream",
python_callable=list).expand(
+ op_kwargs=plain.output
+ )
+ run = dag_maker.create_dagrun()
+ stored = next(ti for ti in run.task_instances if ti.task_id ==
"group.plain")
+ dag = stored.task.dag
+
+ infos = [
+ DagInfo(dag),
+ DagRunInfo(run),
+ TaskInstanceInfo(stored, dag=dag, dag_run=run),
+ AssetInfo(asset),
+ TaskInfo(stored.task, runtime_assets=[(asset, None)]),
+ TaskInfoComplete(stored.task),
+ TaskInfo(dag.get_task("mapped")),
+ TaskInfoComplete(dag.get_task("mapped")),
+ TaskInfo(dag.get_task("mapped_from_upstream")),
+ TaskInfoComplete(dag.get_task("mapped_from_upstream")),
+ TaskGroupInfo(stored.task.task_group),
+ ]
+
+ # A new encodable must be added here too
+ assert {type(info) for info in infos} ==
set(_get_info_json_encodable_subclasses())
+ for info in infos:
+ assert vars(info) == {}
+ unpickled = pickle.loads(pickle.dumps(info))
+ assert type(unpickled) is type(info)
+ assert unpickled == info
+
+
[email protected]_test
+def
test_task_info_complete_of_stored_mapped_task_leaves_out_its_expand_input(dag_maker):
+ with dag_maker(serialized=True):
+
+ @task
+ def produce():
+ return [1, 2]
+
+ @task
+ def consume(x):
+ return x
+
+ produced = produce()
+ consume.expand(x=produced)
+ PythonOperator.partial(task_id="mapped",
python_callable=list).expand(op_kwargs=produced)
+ dag = dag_maker.create_dagrun().task_instances[0].task.dag
+
+ for task_id in ("consume", "mapped"):
+ assert not {"expand_input", "op_kwargs_expand_input"} &
TaskInfoComplete(dag.get_task(task_id)).keys()
+
+
+def
test_task_info_complete_keeps_only_json_native_values_of_operator_attributes():
+ operator = BashOperator(task_id="t", bash_command="echo 1")
+ operator.custom = {"when": datetime.datetime(2026, 1, 2), "lock":
threading.Lock()}
+
+ info = TaskInfoComplete(operator)
+
+ assert info["custom"] == {"when": "2026-01-02 00:00:00", "lock":
"<<non-serializable: lock>>"}
+ pickle.dumps(info)
+
+
def test_task_instance_info_uses_the_position_it_is_given():
# The task instance's own loop and expansion differ, so the info can only
match the given position
ti = SimpleNamespace(
@@ -6085,6 +6169,16 @@ class _DatasetName(str, enum.Enum):
SYNC = "user_sync"
+class _Color(enum.Enum):
+ RED = 1
+
+
[email protected]
+class _Creds:
+ user: str
+ password: str
+
+
@pytest.mark.parametrize(
("value", "expected"),
[
@@ -6100,19 +6194,27 @@ class _DatasetName(str, enum.Enum):
pytest.param(
[{"nested": {"items": ["a", None, 1.0]}}], [{"nested": {"items":
["a", None, 1.0]}}], id="nested"
),
+ pytest.param([{"c": _Color.RED}], [{"c": 1}], id="enum_uses_value"),
+ pytest.param([{"s": {"a"}}], [{"s": ["a"]}], id="set"),
+ pytest.param(
+ [{"creds": _Creds("u", "p")}], [{"creds": {"user": "u",
"password": "p"}}], id="attrs_as_dict"
+ ),
+ pytest.param(
+ [{"lock": threading.Lock()}], [{"lock": "<<non-serializable:
lock>>"}], id="other_object"
+ ),
],
)
-def test_convert_to_json_safe_list(value, expected):
- assert _convert_to_json_safe_list(value, "outlets") == expected
+def test_convert_to_json_safe(value, expected):
+ assert _convert_to_json_safe(value, "outlets") == expected
@patch("airflow.providers.openlineage.utils.utils.log", autospec=True)
-def test_convert_to_json_safe_list_returns_none_on_failure(mock_log):
+def test_convert_to_json_safe_returns_none_on_failure(mock_log):
circular: list = []
circular.append(circular)
- assert _convert_to_json_safe_list(circular, "outlets") is None
+ assert _convert_to_json_safe(circular, "outlets") is None
mock_log.warning.assert_called_once_with(
- "Unable to serialize task `%s` to a JSON list: %s", "outlets", mock.ANY
+ "Unable to convert `%s` to JSON-native types: %s", "outlets", mock.ANY
)
assert str(mock_log.warning.call_args.args[2]) == "Circular reference
detected"
