This is an automated email from the ASF dual-hosted git repository.
mobuchowski 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 2f9eb829973 Fix openlineage team name handling, add dag_version_data
and TI note (#70772)
2f9eb829973 is described below
commit 2f9eb829973a3f12f89468df25af36edf57856b9
Author: Kacper Muda <[email protected]>
AuthorDate: Fri Jul 31 15:11:09 2026 +0200
Fix openlineage team name handling, add dag_version_data and TI note
(#70772)
---
.../airflow/providers/openlineage/utils/utils.py | 35 +++--
.../tests/unit/openlineage/utils/test_utils.py | 154 +++++++++++++++------
2 files changed, 140 insertions(+), 49 deletions(-)
diff --git
a/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py
b/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py
index 5d58ecd9083..0afb3b347df 100644
--- a/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py
+++ b/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py
@@ -982,6 +982,9 @@ class DagRunInfo(InfoJsonEncodable):
"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),
}
@@ -1039,7 +1042,7 @@ class DagRunInfo(InfoJsonEncodable):
return {"alerts": result} if result else None
@classmethod
- def dag_version_info(cls, dagrun: DagRun, key: str) -> str | int | None:
+ 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", [])
@@ -1057,6 +1060,10 @@ class DagRunInfo(InfoJsonEncodable):
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
@@ -1073,18 +1080,29 @@ class DagRunInfo(InfoJsonEncodable):
if 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:
- bundle_name = cls.dag_version_info(dagrun, "bundle_name")
- if not isinstance(bundle_name, str):
+ # 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.dagbundle import DagBundleModel
+ from airflow.models.dag import DagModel
- return DagBundleModel.get_team_name(bundle_name)
+ return DagModel.get_team_name(dagrun.dag_id, session=session)
except Exception as e:
- log.warning(
+ log.info(
"OpenLineage failed to resolve the team name for dag `%s`:
%s.",
- safe_getattr(dagrun, "dag_id"),
+ dagrun.dag_id,
e,
)
log.debug("Exception details:", exc_info=True)
@@ -1094,9 +1112,10 @@ class DagRunInfo(InfoJsonEncodable):
class TaskInstanceInfo(InfoJsonEncodable):
"""Defines encoding TaskInstance object to JSON."""
- includes = ["duration", "try_number", "pool", "queued_dttm", "log_url"]
+ includes = ["duration", "log_url", "pool", "queued_dttm", "try_number"]
casts = {
"log_url": lambda ti: getattr(ti, "log_url", None),
+ "note": lambda ti: safe_getattr(ti, "note", None), # From manual
state changes only
"map_index": lambda ti: ti.map_index if getattr(ti, "map_index", -1)
!= -1 else None,
"rendered_map_index": lambda ti: (
getattr(ti, "rendered_map_index", None) if getattr(ti,
"map_index", -1) != -1 else None
diff --git a/providers/openlineage/tests/unit/openlineage/utils/test_utils.py
b/providers/openlineage/tests/unit/openlineage/utils/test_utils.py
index 2cd0f3dd3fa..0b5d4fe2f87 100644
--- a/providers/openlineage/tests/unit/openlineage/utils/test_utils.py
+++ b/providers/openlineage/tests/unit/openlineage/utils/test_utils.py
@@ -233,6 +233,7 @@ def test_get_airflow_dag_run_facet():
bundle_version="bundle_version",
id="version_id",
version_number="version_number",
+ version_data={"some": "data"},
)
]
dagrun_mock.deadlines = []
@@ -252,9 +253,14 @@ def test_get_airflow_dag_run_facet():
}
if hasattr(dag, "schedule_interval"): # Airflow 2 compat.
expected_dag_info["schedule_interval"] = "@once"
- note: str | None = None
+
+ optional_result = {}
if AIRFLOW_V_3_2_PLUS:
- note = "note"
+ optional_result["note"] = "note"
+
+ if AIRFLOW_V_3_3_PLUS:
+ optional_result["dag_version_data"] = {"some": "data"}
+
assert result == {
"airflowDagRun": AirflowDagRunFacet(
dag=expected_dag_info,
@@ -283,7 +289,9 @@ def test_get_airflow_dag_run_facet():
"partition_key": "some_partition_key",
"partition_date": "2024-06-01T02:03:34+00:00",
"triggered_by": "something",
- "note": note,
+ "note": None,
+ "dag_version_data": None,
+ **optional_result,
},
)
}
@@ -348,74 +356,92 @@ def test_dag_run_version(key):
@pytest.mark.db_test
[email protected](not AIRFLOW_V_3_3_PLUS, reason="multi-team requires
Airflow 3.3+")
-@patch("airflow.models.dagbundle.DagBundleModel.get_team_name")
-@patch("airflow.providers.openlineage.utils.utils.airflow_conf.getboolean",
return_value=True)
-def test_dag_run_team_name(
- mock_getboolean,
- mock_get_team_name,
-):
-
[email protected](not AIRFLOW_V_3_3_PLUS, reason="version_data requires
Airflow 3.3+")
+def test_dag_run_version_data():
dagrun_mock = MagicMock(DagRun)
- dagrun_mock.dag_versions = [
- MagicMock(
- bundle_name="bundle_name",
- bundle_version="bundle_version",
- id="version_id",
- version_number="version_number",
- )
- ]
+ dagrun_mock.dag_versions = [MagicMock(version_data={"schema": 1})]
+ assert DagRunInfo.dag_version_info(dagrun_mock, "version_data") ==
{"schema": 1}
- mock_get_team_name.return_value = "team_a"
- assert DagRunInfo.team_name(dagrun_mock) == "team_a"
[email protected]_test
+@patch("airflow.providers.openlineage.utils.utils.AIRFLOW_V_3_3_PLUS", False)
+def test_dag_run_version_data_below_3_3():
+ dagrun_mock = MagicMock(DagRun)
+ dagrun_mock.dag_versions = [MagicMock(version_data={"schema": 1})]
+ assert DagRunInfo.dag_version_info(dagrun_mock, "version_data") is None
- mock_get_team_name.assert_called_once_with("bundle_name")
+
[email protected](not AIRFLOW_V_3_3_PLUS, reason="version_data requires
Airflow 3.3+")
+def test_dag_run_version_data_detached_version_row():
+ """version_data is lazy-loaded and can hit a detached session like the
other columns."""
+ version = MagicMock()
+ type(version).version_data =
PropertyMock(side_effect=DetachedInstanceError)
+ dag_run = MagicMock()
+ dag_run.dag_versions = [version]
+ assert DagRunInfo.dag_version_info(dag_run, "version_data") is None
@pytest.mark.db_test
[email protected](not AIRFLOW_V_3_3_PLUS, reason="multi-team requires
Airflow 3.3+")
+@patch("airflow.models.dag.DagModel.get_team_name", return_value="team_a")
+@patch("sqlalchemy.orm.object_session")
+@patch("airflow.providers.openlineage.utils.utils.airflow_conf.getboolean",
return_value=True)
+def test_dag_run_team_name(mock_getboolean, mock_object_session,
mock_get_team_name):
+ """DB fallback uses the dagrun's existing session — no new session opened,
no HA lock risk."""
+ mock_session = MagicMock()
+ mock_object_session.return_value = mock_session
+
+ dagrun_mock = MagicMock(spec=DagRun)
+ dagrun_mock.dag_id = "test_dag"
+
+ assert DagRunInfo.team_name(dagrun_mock) == "team_a"
+ mock_get_team_name.assert_called_once_with("test_dag",
session=mock_session)
+
+
@pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="multi-team requires
Airflow 3.3+")
@pytest.mark.parametrize("team_name", ["team_a", None])
-@patch("airflow.models.dagbundle.DagBundleModel.get_team_name")
+@patch("sqlalchemy.orm.object_session")
@patch("airflow.providers.openlineage.utils.utils.airflow_conf.getboolean",
return_value=True)
-def test_dag_run_team_name_from_execution_api_dag_run(mock_getboolean,
mock_get_team_name, team_name):
- """The task runner has no DB session, so a DagRun carrying `team_name`
must be trusted as-is."""
- # A resolvable bundle plus a DB answer, so falling through to the lookup
would be observable.
- dagrun_mock = MagicMock(spec_set=["team_name", "dag_versions"])
+def test_dag_run_team_name_from_execution_api_dag_run(mock_getboolean,
mock_object_session, team_name):
+ """The task runner has no DB session, so a DagRun carrying `team_name`
must be trusted as-is.
+
+ The cascade must stop at the first step — the DB lookup (object_session)
must never be reached.
+ """
+ dagrun_mock = MagicMock(spec_set=["team_name"])
dagrun_mock.team_name = team_name
- dagrun_mock.dag_versions = [MagicMock(bundle_name="bundle_name")]
- mock_get_team_name.return_value = "from_db"
assert DagRunInfo.team_name(dagrun_mock) == team_name
- mock_get_team_name.assert_not_called()
+ mock_object_session.assert_not_called()
@pytest.mark.db_test
@pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="multi-team requires
Airflow 3.3+")
-@patch("airflow.models.dagbundle.DagBundleModel.get_team_name")
+@patch("airflow.models.dag.DagModel.get_team_name",
side_effect=RuntimeError("db gone"))
+@patch("sqlalchemy.orm.object_session")
@patch("airflow.providers.openlineage.utils.utils.airflow_conf.getboolean",
return_value=True)
-def test_dag_run_team_name_lookup_failure_does_not_raise(mock_getboolean,
mock_get_team_name):
+def test_dag_run_team_name_lookup_failure_does_not_raise(
+ mock_getboolean, mock_object_session, mock_get_team_name
+):
"""A failed lookup must degrade to None -- `_cast_fields` would otherwise
lose the whole event."""
- dagrun_mock = MagicMock(DagRun)
- dagrun_mock.dag_versions = [MagicMock(bundle_name="bundle_name")]
- mock_get_team_name.side_effect = RuntimeError("Session must be set
before!")
+ mock_object_session.return_value = MagicMock()
+ dagrun_mock = MagicMock(spec=DagRun)
+ dagrun_mock.dag_id = "test_dag"
assert DagRunInfo.team_name(dagrun_mock) is None
@pytest.mark.db_test
@pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="multi-team requires
Airflow 3.3+")
-@patch("airflow.models.dagbundle.DagBundleModel.get_team_name")
+@patch("sqlalchemy.orm.object_session", return_value=None)
@patch("airflow.providers.openlineage.utils.utils.airflow_conf.getboolean",
return_value=True)
-def test_dag_run_team_name_no_bundle(mock_getboolean, mock_get_team_name):
- dagrun_mock = MagicMock(DagRun)
- del dagrun_mock.dag_versions
+def test_dag_run_team_name_no_session(mock_getboolean, mock_object_session):
+ """When the dagrun has no attached session the lookup is skipped and None
is returned."""
+ dagrun_mock = MagicMock(spec=DagRun)
+ dagrun_mock.dag_id = "test_dag"
assert DagRunInfo.team_name(dagrun_mock) is None
- mock_get_team_name.assert_not_called()
-
@pytest.mark.db_test
@pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="multi-team requires
Airflow 3.3+")
@@ -430,6 +456,44 @@ def
test_dag_run_team_name_multi_team_disabled(mock_getboolean, mock_get_team_na
mock_get_team_name.assert_not_called()
[email protected]_test
+@patch("airflow.providers.openlineage.utils.utils.AIRFLOW_V_3_3_PLUS", False)
+def test_dag_run_team_name_below_airflow_3_3():
+ """Airflow < 3.3 has no multi-team support — team_name must return None
unconditionally."""
+ dagrun_mock = MagicMock(spec=DagRun)
+
+ assert DagRunInfo.team_name(dagrun_mock) is None
+
+
[email protected]_test
[email protected](not AIRFLOW_V_3_3_PLUS, reason="multi-team requires
Airflow 3.3+")
+@patch("sqlalchemy.orm.object_session")
+@patch("airflow.providers.openlineage.utils.utils.airflow_conf.getboolean",
return_value=True)
+def test_dag_run_team_name_from_scheduler_stamp(mock_getboolean,
mock_object_session):
+ """Scheduler stamps _team_name on ORM DagRun objects; the attribute is
read back as-is.
+
+ The cascade must stop at `_team_name` — the DB lookup (object_session)
must never be reached.
+ """
+ dagrun_mock = MagicMock(spec=DagRun)
+ dagrun_mock._team_name = "team_a"
+
+ assert DagRunInfo.team_name(dagrun_mock) == "team_a"
+
+ mock_object_session.assert_not_called()
+
+
[email protected]_test
[email protected](not AIRFLOW_V_3_3_PLUS, reason="multi-team requires
Airflow 3.3+")
+@patch("sqlalchemy.orm.object_session", return_value=None)
+@patch("airflow.providers.openlineage.utils.utils.airflow_conf.getboolean",
return_value=True)
+def test_dag_run_team_name_from_scheduler_stamp_non_str(mock_getboolean,
mock_object_session):
+ """Non-str _team_name (corrupted stamp) must not propagate — fall through
returns None."""
+ dagrun_mock = MagicMock(spec=DagRun)
+ dagrun_mock._team_name = 42
+
+ assert DagRunInfo.team_name(dagrun_mock) is None
+
+
def test_get_fully_qualified_class_name_serialized_operator():
op_module_path = BASH_OPERATOR_PATH
op_name = "BashOperator"
@@ -3025,6 +3089,7 @@ def test_dagrun_info_af3(mocked_dag_versions):
dv2.version_number = "version_number"
dv2.bundle_name = "bundle_name"
dv2.bundle_version = "bundle_version"
+ dv2.version_data = {"some": "data"}
optional_args = {}
if AIRFLOW_V_3_2_PLUS:
@@ -3059,12 +3124,16 @@ def test_dagrun_info_af3(mocked_dag_versions):
optional_result["partition_key"] = "some_partition_key"
optional_result["partition_date"] = "2024-06-01T00:00:00+00:00"
+ if AIRFLOW_V_3_3_PLUS:
+ optional_result["dag_version_data"] = {"some": "data"}
+
result = DagRunInfo(dagrun)
assert dict(result) == {
"conf": {"a": 1},
"clear_number": 0,
"dag_id": "dag_id",
"dag_team_name": None,
+ "dag_version_data": None,
"data_interval_end": "2024-06-01T00:00:00+00:00",
"data_interval_start": "2024-06-01T00:00:00+00:00",
"duration": 74.000546,
@@ -3127,6 +3196,7 @@ def test_dagrun_info_af2():
"dag_bundle_version": None,
"dag_version_id": None,
"dag_version_number": None,
+ "dag_version_data": None,
"note": None,
}
@@ -3169,6 +3239,7 @@ def test_taskinstance_info_af3():
assert dict(TaskInstanceInfo(runtime_ti)) == {
"log_url": runtime_ti.log_url,
"map_index": 2,
+ "note": None,
"rendered_map_index": None,
"try_number": 1,
"dag_bundle_version": "bundle_version",
@@ -3207,6 +3278,7 @@ def test_taskinstance_info_af2():
"log_url": "some_log_url",
"dag_bundle_name": None,
"dag_bundle_version": None,
+ "note": None,
}
# Also tested manually that it works well on AF2, hard to test hybrid
property so just mocking it here
