diff --git a/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py b/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py index f40e9b12a0be7..2cdda396adc76 100644 --- a/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py +++ b/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py @@ -38,6 +38,7 @@ ) from openlineage.client.utils import RedactMixin from openlineage.client.uuid import generate_static_uuid +from sqlalchemy.orm import object_session from sqlalchemy.orm.exc import DetachedInstanceError from airflow import __version__ as AIRFLOW_VERSION @@ -1068,7 +1069,15 @@ def team_name(cls, dagrun: DagRun) -> str | None: if not isinstance(bundle_name, str): return None - return DagBundleModel.get_team_name(bundle_name) + # Reuse the existing ORM session associated with the DagRun. Creating a new + # session here (via @provide_session) can trigger an unexpected commit within + # the scheduler callback, which breaks HA lock semantics. + session = object_session(dagrun) + + if session is None: + return None + + return DagBundleModel.get_team_name(bundle_name=bundle_name, session=session) class TaskInstanceInfo(InfoJsonEncodable): diff --git a/providers/openlineage/tests/unit/openlineage/utils/test_utils.py b/providers/openlineage/tests/unit/openlineage/utils/test_utils.py index e659e4f3de3da..290c342626260 100644 --- a/providers/openlineage/tests/unit/openlineage/utils/test_utils.py +++ b/providers/openlineage/tests/unit/openlineage/utils/test_utils.py @@ -335,13 +335,18 @@ def test_dag_run_version(key): @pytest.mark.db_test @pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="multi-team requires Airflow 3.3+") +@patch("airflow.providers.openlineage.utils.utils.object_session") @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, + mock_object_session, ): + session = MagicMock() + mock_object_session.return_value = session + dagrun_mock = MagicMock(DagRun) dagrun_mock.dag_versions = [ MagicMock( @@ -356,7 +361,10 @@ def test_dag_run_team_name( assert DagRunInfo.team_name(dagrun_mock) == "team_a" - mock_get_team_name.assert_called_once_with("bundle_name") + mock_get_team_name.assert_called_once_with( + bundle_name="bundle_name", + session=session, + ) @pytest.mark.db_test