From d59a1efd302a0992164811ad62203445d6a691c6 Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Tue, 8 Sep 2026 14:06:26 +0100 Subject: [PATCH 1/3] Keep the task instance in dag callbacks when its Dag version is missing Since 3.3.0 the scheduler dropped a task instance with no dag_version_id from the callback context rather than crash on it, which cost users who emit tracing events from on_success_callback / on_failure_callback the ti they correlate each event on, leaving them no way to tie a callback back to its task. Handle a missing dag_version less destructively, so the callbacks receive the full task instance. --- .../src/airflow/jobs/scheduler_job_runner.py | 1 + airflow-core/src/airflow/models/dagrun.py | 86 +++++++++--- airflow-core/tests/unit/models/test_dagrun.py | 125 +++++++++++++++--- 3 files changed, 174 insertions(+), 38 deletions(-) diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index 90aefb30be99d..089c78f3827e7 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py @@ -3062,6 +3062,7 @@ def _schedule_dag_run( relevant_ti=last_unfinished_ti, reason="timed_out", execute=False, + session=session, ) # Team name should be added before listeners are called in notify_dagrun_state_changed() diff --git a/airflow-core/src/airflow/models/dagrun.py b/airflow-core/src/airflow/models/dagrun.py index c22fbc4a75a16..b39824b4c8994 100644 --- a/airflow-core/src/airflow/models/dagrun.py +++ b/airflow-core/src/airflow/models/dagrun.py @@ -120,7 +120,10 @@ from sqlalchemy.sql.elements import Case, ColumnElement from sqlalchemy.sql.selectable import Select - from airflow.api_fastapi.execution_api.datamodels.taskinstance import DagRun as DRDataModel + from airflow.api_fastapi.execution_api.datamodels.taskinstance import ( + DagRun as DRDataModel, + TaskInstance as TIDataModel, + ) from airflow.models.dag_version import DagVersion from airflow.models.taskinstancekey import TaskInstanceKey from airflow.sdk import DAG as SDKDAG @@ -1333,6 +1336,7 @@ def recalculate(self) -> _UnfinishedStates: relevant_ti=ti_causing_failure, reason="task_failure", execute=execute_callbacks, + session=session, ) # Check if the max_consecutive_failed_dag_runs has been provided and not 0 @@ -1363,6 +1367,7 @@ def recalculate(self) -> _UnfinishedStates: relevant_ti=last_succeeded_ti, reason="success", execute=execute_callbacks, + session=session, ) if dag.deadline: @@ -1400,6 +1405,7 @@ def recalculate(self) -> _UnfinishedStates: relevant_ti=blocking_ti, reason="all_tasks_deadlocked", execute=execute_callbacks, + session=session, ) # finally, if the leaves aren't done, the dag is still running @@ -1508,6 +1514,45 @@ def notify_dagrun_state_changed(self, msg: str): # we can't get all the state changes on SchedulerJob, # or LocalTaskJob, so we don't want to "falsely advertise" we notify about that + def _build_callback_last_ti(self, relevant_ti: TI, *, session: Session) -> TIDataModel | None: + """ + Build a callback context's ``last_ti``, standing in a Dag version if the record has none. + + ``session`` is required, not defaulted: ``settings.Session`` is a ``scoped_session``, + so acquiring one here would hand back the caller's own and then commit and close it. + """ + from airflow.api_fastapi.execution_api.datamodels.taskinstance import TaskInstance as TIDataModel + from airflow.models.dag_version import DagVersion + + dag_version_id = relevant_ti.dag_version_id + if dag_version_id is None: + dag_version_id = self.created_dag_version_id + if dag_version_id is None: + latest_dag_version = DagVersion.get_latest_version(self.dag_id, session=session) + dag_version_id = latest_dag_version.id if latest_dag_version else None + if dag_version_id is None: + self.log.warning( + "Task instance %s has no dag_version_id and Dag %s has no version to stand in; " + "omitting last_ti from the dag callback context.", + relevant_ti, + self.dag_id, + ) + return None + self.log.warning( + "Task instance %s has no dag_version_id (pre-versioning record); " + "reporting it as %s in the dag callback context.", + relevant_ti, + dag_version_id, + ) + values = { + name: getattr(relevant_ti, name) + for name in TIDataModel.model_fields + if hasattr(relevant_ti, name) + } + values["dag_version_id"] = dag_version_id + return TIDataModel.model_validate(values) + + @provide_session def produce_dag_callback( self, dag: SerializedDAG, @@ -1515,23 +1560,18 @@ def produce_dag_callback( relevant_ti: TI | None = None, reason: str = "success", execute: bool = False, + *, + session: Session = NEW_SESSION, ) -> DagCallbackRequest | None: """Create a callback request for the DAG, or execute the callbacks directly if instructed, and return None.""" - # Historical task instances created before the dag_version table existed (migration - # 0047_3_0_0_add_dag_versioning) have dag_version_id=None. The TaskInstance datamodel used to - # build the callback context requires a non-null UUID, so passing such a TI as last_ti would - # raise a ValidationError and crash the scheduler. Drop last_ti in that case; the callback still - # fires with a minimal context (dag, run_id, reason). - if relevant_ti is not None and relevant_ti.dag_version_id is None: - self.log.warning( - "Task instance %s has no dag_version_id (pre-versioning record); " - "omitting last_ti from the dag callback context.", - relevant_ti, - ) - relevant_ti = None if not execute: from airflow.models.dag_version import _resolve_version_data + last_ti = ( + self._build_callback_last_ti(relevant_ti, session=session) + if relevant_ti is not None + else None + ) # Only carry version_data for pinned runs so the callback initializes the bundle # against the same version the run used. version_data = _resolve_version_data(self.created_dag_version, self.bundle_version) @@ -1544,7 +1584,7 @@ def produce_dag_callback( version_data=version_data, context_from_server=DagRunContext( dag_run=self, - last_ti=relevant_ti, + last_ti=last_ti, ), is_failure_callback=(not success), msg=reason, @@ -1554,22 +1594,26 @@ def produce_dag_callback( success=success, relevant_ti=relevant_ti, reason=reason, + session=session, ) return None + @provide_session def execute_dag_callbacks( - self, dag: SDKDAG, success: bool = True, relevant_ti: TI | None = None, reason: str = "success" + self, + dag: SDKDAG, + success: bool = True, + relevant_ti: TI | None = None, + reason: str = "success", + *, + session: Session = NEW_SESSION, ): """Only needed for `dag.test` where `execute_callbacks=True` is passed to `update_state`.""" - from airflow.api_fastapi.execution_api.datamodels.taskinstance import ( - TaskInstance as TIDataModel, - TIRunContext, - ) + from airflow.api_fastapi.execution_api.datamodels.taskinstance import TIRunContext from airflow.models.dag import DagModel from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance - if relevant_ti: - last_ti_model = TIDataModel.model_validate(relevant_ti, from_attributes=True) + if relevant_ti and (last_ti_model := self._build_callback_last_ti(relevant_ti, session=session)): task = dag.get_task(relevant_ti.task_id) runtime_ti = RuntimeTaskInstance.model_construct( diff --git a/airflow-core/tests/unit/models/test_dagrun.py b/airflow-core/tests/unit/models/test_dagrun.py index ffa223508031b..d48e68e397ab6 100644 --- a/airflow-core/tests/unit/models/test_dagrun.py +++ b/airflow-core/tests/unit/models/test_dagrun.py @@ -40,7 +40,7 @@ select, update, ) -from sqlalchemy.orm import joinedload +from sqlalchemy.orm import Session, joinedload from sqlalchemy.orm.exc import StaleDataError from airflow import settings @@ -97,8 +97,6 @@ from unit.models import DEFAULT_DATE as _DEFAULT_DATE if TYPE_CHECKING: - from sqlalchemy.orm.session import Session - from airflow.serialization.definitions.dag import SerializedDAG pytestmark = [pytest.mark.db_test, pytest.mark.need_serialized_dag] @@ -465,7 +463,7 @@ def on_success_callable(context): with mock.patch.object(dag_run, "execute_dag_callbacks") as execute_dag_callbacks: _, callback = dag_run.update_state() assert execute_dag_callbacks.mock_calls == [ - mock.call(dag=dag, success=True, relevant_ti=ANY, reason="success") + mock.call(dag=dag, success=True, relevant_ti=ANY, reason="success", session=ANY) ] # Make sure the correct TI is passed on success call_args = execute_dag_callbacks.call_args @@ -499,7 +497,7 @@ def on_failure_callable(context): with mock.patch.object(dag_run, "execute_dag_callbacks") as execute_dag_callbacks: _, callback = dag_run.update_state() assert execute_dag_callbacks.mock_calls == [ - mock.call(dag=dag, success=False, relevant_ti=ANY, reason="task_failure") + mock.call(dag=dag, success=False, relevant_ti=ANY, reason="task_failure", session=ANY) ] # Make sure the correct TI is passed on failure call_args = execute_dag_callbacks.call_args @@ -541,7 +539,13 @@ def on_failure_callable(context): with mock.patch.object(dr, "execute_dag_callbacks") as execute_dag_callbacks: _, callback = dr.update_state(execute_callbacks=True) assert execute_dag_callbacks.mock_calls == [ - mock.call(dag=serialized_dag, success=False, relevant_ti=ti_middle, reason="all_tasks_deadlocked") + mock.call( + dag=serialized_dag, + success=False, + relevant_ti=ti_middle, + reason="all_tasks_deadlocked", + session=ANY, + ) ] # Make sure the correct TI is passed on deadlock call_args = execute_dag_callbacks.call_args @@ -1436,7 +1440,7 @@ def on_success_callable(context): with mock.patch.object(dag_run, "execute_dag_callbacks") as execute_dag_callbacks: _, callback = dag_run.update_state() assert execute_dag_callbacks.mock_calls == [ - mock.call(dag=scheduler_dag, success=True, relevant_ti=ANY, reason="success") + mock.call(dag=scheduler_dag, success=True, relevant_ti=ANY, reason="success", session=ANY) ] # Make sure the correct TI is passed on success call_args = execute_dag_callbacks.call_args @@ -4155,27 +4159,63 @@ def on_failure(context): assert context_received["ti"].dag_id == "test_dag" assert context_received["ti"].run_id == dr.run_id - def test_produce_dag_callback_drops_last_ti_without_dag_version(self, dag_maker, session): - """A historical TI with dag_version_id=None must not crash callback construction.""" + @pytest.mark.parametrize("run_keeps_version", [True, False]) + def test_produce_dag_callback_stands_in_version_for_versionless_last_ti( + self, dag_maker, session, run_keeps_version + ): + """A historical TI with dag_version_id=None still reaches the callback, under a stand-in version.""" + with dag_maker("test_dag", session=session) as dag: + BashOperator(task_id="test_task", bash_command="echo 1") + + dr = dag_maker.create_dagrun() + dr.dag_model = DagModel.get_dagmodel(dag.dag_id, session=session) + run_version_id = dr.created_dag_version_id + ti = dr.get_task_instance("test_task") + ti.dag_version_id = None + if not run_keeps_version: + dr.created_dag_version_id = None + # Newer than the run's, so the two fallback sources are distinguishable. + latest_version = DagVersion.write_dag(dag_id=dr.dag_id, bundle_name="dag_maker", version_number=2) + session.flush() + assert latest_version.id != run_version_id + + expected_version_id = run_version_id if run_keeps_version else latest_version.id + + callback = dr.produce_dag_callback( + dag=dag, success=False, relevant_ti=ti, reason="task_failure", session=session + ) + + assert callback is not None + assert callback.context_from_server is not None + last_ti = callback.context_from_server.last_ti + assert last_ti is not None + assert last_ti.task_id == "test_task" + assert last_ti.dag_version_id == expected_version_id + + def test_produce_dag_callback_drops_last_ti_when_dag_has_no_version(self, dag_maker, session): + """With no version anywhere to stand in, the callback still fires without last_ti.""" with dag_maker("test_dag", session=session) as dag: BashOperator(task_id="test_task", bash_command="echo 1") dr = dag_maker.create_dagrun() dr.dag_model = DagModel.get_dagmodel(dag.dag_id, session=session) ti = dr.get_task_instance("test_task") - # Simulate a task instance created before the dag_version table existed. ti.dag_version_id = None + dr.created_dag_version_id = None session.flush() - callback = dr.produce_dag_callback(dag=dag, success=False, relevant_ti=ti, reason="task_failure") + # Patched here, not as a decorator: dag_maker's setup needs the real lookup. + with mock.patch.object(DagVersion, "get_latest_version", autospec=True, return_value=None): + callback = dr.produce_dag_callback( + dag=dag, success=False, relevant_ti=ti, reason="task_failure", session=session + ) assert callback is not None - # last_ti is dropped so the non-null UUID datamodel validation never fires. assert callback.context_from_server is not None assert callback.context_from_server.last_ti is None def test_execute_dag_callbacks_without_dag_version(self, dag_maker, session): - """The execute=True path must also tolerate a TI with dag_version_id=None.""" + """The execute=True path must also carry a TI with dag_version_id=None into the context.""" context_received = None def on_failure(context): @@ -4189,18 +4229,69 @@ def on_failure(context): dr.dag_model = DagModel.get_dagmodel(dag.dag_id, session=session) ti = dr.get_task_instance("test_task") ti.dag_version_id = None + dr.created_dag_version_id = None session.flush() dag.on_failure_callback = on_failure dag.has_on_failure_callback = True - dr.produce_dag_callback(dag=dag, success=False, relevant_ti=ti, reason="task_failure", execute=True) + dr.produce_dag_callback( + dag=dag, + success=False, + relevant_ti=ti, + reason="task_failure", + execute=True, + session=session, + ) - # Callback still fires with the minimal fallback context (no last_ti template vars). assert context_received is not None assert context_received["reason"] == "task_failure" - assert "ti" not in context_received - assert context_received["run_id"] == dr.run_id + assert context_received["ti"].task_id == "test_task" + assert context_received["ti"].run_id == dr.run_id + assert ( + context_received["ti"].dag_version_id + == DagVersion.get_latest_version(dag.dag_id, session=session).id + ) + + @pytest.mark.parametrize("strip_dag_version", [False, True]) + def test_produce_dag_callback_preserves_callers_transaction(self, dag_maker, session, strip_dag_version): + """Executing callbacks must not commit or close the session the caller handed in.""" + + def on_failure(context): + pass + + with dag_maker("test_dag", session=session, on_failure_callback=on_failure) as dag: + BashOperator(task_id="test_task", bash_command="echo 1") + + dr = dag_maker.create_dagrun() + dr.dag_model = DagModel.get_dagmodel(dag.dag_id, session=session) + # get_task_instance would otherwise default the session and close this one. + ti = dr.get_task_instance("test_task", session=session) + if strip_dag_version: + ti.dag_version_id = None + dr.created_dag_version_id = None + session.flush() + + dag.on_failure_callback = on_failure + dag.has_on_failure_callback = True + + with ( + mock.patch.object(Session, "commit", autospec=True) as mock_commit, + mock.patch.object(Session, "close", autospec=True) as mock_close, + ): + dr.produce_dag_callback( + dag=dag, + success=False, + relevant_ti=ti, + reason="task_failure", + execute=True, + session=session, + ) + + assert mock_commit.mock_calls == [] + assert mock_close.mock_calls == [] + # close() expunges everything, so an attached instance proves the session survived. + assert ti in session @pytest.mark.parametrize( ("multi_team", "team_name", "expected_tags"), From 7dee6ffd98679b15e79880a91581635e5cf752ec Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Thu, 10 Sep 2026 09:41:58 +0100 Subject: [PATCH 2/3] Require the caller's session when producing Dag callbacks settings.Session is a scoped_session, so a defaulted session hands back the caller's own and then commits and closes it out from under them. Every production caller already holds a session, so the default bought nothing and only stood to end a transaction nobody asked to end. --- airflow-core/src/airflow/models/dagrun.py | 41 ++++++++++--------- airflow-core/tests/unit/models/test_dag.py | 8 ++-- airflow-core/tests/unit/models/test_dagrun.py | 22 ++++++---- 3 files changed, 40 insertions(+), 31 deletions(-) diff --git a/airflow-core/src/airflow/models/dagrun.py b/airflow-core/src/airflow/models/dagrun.py index b39824b4c8994..5f5493b49a77a 100644 --- a/airflow-core/src/airflow/models/dagrun.py +++ b/airflow-core/src/airflow/models/dagrun.py @@ -1524,26 +1524,29 @@ def _build_callback_last_ti(self, relevant_ti: TI, *, session: Session) -> TIDat from airflow.api_fastapi.execution_api.datamodels.taskinstance import TaskInstance as TIDataModel from airflow.models.dag_version import DagVersion - dag_version_id = relevant_ti.dag_version_id + if relevant_ti.dag_version_id is not None: + return TIDataModel.model_validate(relevant_ti, from_attributes=True) + + dag_version_id = self.created_dag_version_id + if dag_version_id is None: + latest_dag_version = DagVersion.get_latest_version(self.dag_id, session=session) + dag_version_id = latest_dag_version.id if latest_dag_version else None if dag_version_id is None: - dag_version_id = self.created_dag_version_id - if dag_version_id is None: - latest_dag_version = DagVersion.get_latest_version(self.dag_id, session=session) - dag_version_id = latest_dag_version.id if latest_dag_version else None - if dag_version_id is None: - self.log.warning( - "Task instance %s has no dag_version_id and Dag %s has no version to stand in; " - "omitting last_ti from the dag callback context.", - relevant_ti, - self.dag_id, - ) - return None self.log.warning( - "Task instance %s has no dag_version_id (pre-versioning record); " - "reporting it as %s in the dag callback context.", + "Task instance %s has no dag_version_id and Dag %s has no version to stand in; " + "omitting last_ti from the Dag callback context.", relevant_ti, - dag_version_id, + self.dag_id, ) + return None + self.log.warning( + "Task instance %s has no dag_version_id (pre-versioning record); " + "reporting it as %s in the Dag callback context.", + relevant_ti, + dag_version_id, + ) + # The datamodel rejects a null version, so the stand-in has to be spliced in + # before validation rather than copied onto a validated model. values = { name: getattr(relevant_ti, name) for name in TIDataModel.model_fields @@ -1552,7 +1555,6 @@ def _build_callback_last_ti(self, relevant_ti: TI, *, session: Session) -> TIDat values["dag_version_id"] = dag_version_id return TIDataModel.model_validate(values) - @provide_session def produce_dag_callback( self, dag: SerializedDAG, @@ -1561,7 +1563,7 @@ def produce_dag_callback( reason: str = "success", execute: bool = False, *, - session: Session = NEW_SESSION, + session: Session, ) -> DagCallbackRequest | None: """Create a callback request for the DAG, or execute the callbacks directly if instructed, and return None.""" if not execute: @@ -1598,7 +1600,6 @@ def produce_dag_callback( ) return None - @provide_session def execute_dag_callbacks( self, dag: SDKDAG, @@ -1606,7 +1607,7 @@ def execute_dag_callbacks( relevant_ti: TI | None = None, reason: str = "success", *, - session: Session = NEW_SESSION, + session: Session, ): """Only needed for `dag.test` where `execute_callbacks=True` is passed to `update_state`.""" from airflow.api_fastapi.execution_api.datamodels.taskinstance import TIRunContext diff --git a/airflow-core/tests/unit/models/test_dag.py b/airflow-core/tests/unit/models/test_dag.py index ff0b49e806ea3..b22227b695be4 100644 --- a/airflow-core/tests/unit/models/test_dag.py +++ b/airflow-core/tests/unit/models/test_dag.py @@ -1226,8 +1226,8 @@ def test_dag_handle_callback_crash(self, mock_incr, testing_dag_bundle): ) # should not raise any exception - dag_run.execute_dag_callbacks(dag=dag, success=False) - dag_run.execute_dag_callbacks(dag=dag, success=True) + dag_run.execute_dag_callbacks(dag=dag, success=False, session=session) + dag_run.execute_dag_callbacks(dag=dag, success=True, session=session) mock_incr.assert_called_with( "dag.callback_exceptions", @@ -1267,8 +1267,8 @@ def test_dag_handle_callback_with_removed_task(self, dag_maker, session, testing assert dag_run.get_task_instance(task_removed.task_id).state == TaskInstanceState.REMOVED # should not raise any exception - dag_run.execute_dag_callbacks(dag=dag, success=False) - dag_run.execute_dag_callbacks(dag=dag, success=True) + dag_run.execute_dag_callbacks(dag=dag, success=False, session=session) + dag_run.execute_dag_callbacks(dag=dag, success=True, session=session) @time_machine.travel(timezone.datetime(2025, 11, 11)) @pytest.mark.parametrize( diff --git a/airflow-core/tests/unit/models/test_dagrun.py b/airflow-core/tests/unit/models/test_dagrun.py index d48e68e397ab6..3546b9ad470f2 100644 --- a/airflow-core/tests/unit/models/test_dagrun.py +++ b/airflow-core/tests/unit/models/test_dagrun.py @@ -4056,7 +4056,11 @@ def on_success(context): dag.has_on_success_callback = True dr.execute_dag_callbacks( - dag, success=True, relevant_ti=dr.get_task_instance("test_task"), reason="test_success" + dag, + success=True, + relevant_ti=dr.get_task_instance("test_task", session=session), + reason="test_success", + session=session, ) assert called is True @@ -4088,7 +4092,11 @@ def on_failure(context): dag.has_on_failure_callback = True dr.execute_dag_callbacks( - dag, success=False, relevant_ti=dr.get_task_instance("test_task"), reason="test_failure" + dag, + success=False, + relevant_ti=dr.get_task_instance("test_task", session=session), + reason="test_failure", + session=session, ) assert called is True @@ -4124,8 +4132,9 @@ def on_failure_2(context): dr.execute_dag_callbacks( dag, success=False, - relevant_ti=dr.get_task_instance("test_task"), + relevant_ti=dr.get_task_instance("test_task", session=session), reason="test_failure", + session=session, ) assert call_count == 2 @@ -4149,8 +4158,9 @@ def on_failure(context): dr.execute_dag_callbacks( dag, success=False, - relevant_ti=dr.get_task_instance("test_task"), + relevant_ti=dr.get_task_instance("test_task", session=session), reason="test_failure", + session=session, ) assert context_received is not None @@ -4290,8 +4300,6 @@ def on_failure(context): assert mock_commit.mock_calls == [] assert mock_close.mock_calls == [] - # close() expunges everything, so an attached instance proves the session survived. - assert ti in session @pytest.mark.parametrize( ("multi_team", "team_name", "expected_tags"), @@ -4320,7 +4328,7 @@ def failing_callback(context): conf_vars({("core", "multi_team"): multi_team}), mock.patch("airflow.models.dag.DagModel.get_team_name", return_value=team_name), ): - dr.execute_dag_callbacks(dag, success=False) + dr.execute_dag_callbacks(dag, success=False, session=session) mock_incr.assert_any_call("dag.callback_exceptions", tags=expected_tags) From 9f189de0ed1996a9f44ff494c9c62172ff6b8020 Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Mon, 14 Sep 2026 15:01:42 +0100 Subject: [PATCH 3/3] Say which Dag version a callback stood in, and log it at the right level One warning covered both stand-in tiers, so an ordinary Airflow 3 task instance whose run still has a version was reported to operators as a record predating Dag versioning. Both tiers also warned even though the callback receives its full context and nothing is lost; the sibling helper in the scheduler reserves warning for the case it cannot resolve. The tests covering the fallback let get_task_instance default its session, which committed and closed the fixture's own. The objects came back detached, so the flush never reached the database and no versionless row was ever under test. --- airflow-core/src/airflow/models/dagrun.py | 41 ++++++++++++------- airflow-core/tests/unit/models/test_dagrun.py | 17 ++++---- 2 files changed, 36 insertions(+), 22 deletions(-) diff --git a/airflow-core/src/airflow/models/dagrun.py b/airflow-core/src/airflow/models/dagrun.py index 5f5493b49a77a..223fd3d05ee60 100644 --- a/airflow-core/src/airflow/models/dagrun.py +++ b/airflow-core/src/airflow/models/dagrun.py @@ -1520,6 +1520,12 @@ def _build_callback_last_ti(self, relevant_ti: TI, *, session: Session) -> TIDat ``session`` is required, not defaulted: ``settings.Session`` is a ``scoped_session``, so acquiring one here would hand back the caller's own and then commit and close it. + + Unlike ``_ensure_ti_has_dag_version_id`` in the scheduler, the stand-in is reported + but never written back: this path only describes a task instance to a callback, and + healing the row belongs to the paths that own it. The consequence is that a purge + which heals the row to the latest version first wins, and the run's own version is + then never reported here. """ from airflow.api_fastapi.execution_api.datamodels.taskinstance import TaskInstance as TIDataModel from airflow.models.dag_version import DagVersion @@ -1528,23 +1534,30 @@ def _build_callback_last_ti(self, relevant_ti: TI, *, session: Session) -> TIDat return TIDataModel.model_validate(relevant_ti, from_attributes=True) dag_version_id = self.created_dag_version_id - if dag_version_id is None: + if dag_version_id is not None: + self.log.info( + "Task instance %s has no dag_version_id; reporting it under the run's version %s " + "in the Dag callback context.", + relevant_ti, + dag_version_id, + ) + else: latest_dag_version = DagVersion.get_latest_version(self.dag_id, session=session) - dag_version_id = latest_dag_version.id if latest_dag_version else None - if dag_version_id is None: - self.log.warning( - "Task instance %s has no dag_version_id and Dag %s has no version to stand in; " - "omitting last_ti from the Dag callback context.", + if latest_dag_version is None: + self.log.warning( + "Task instance %s has no dag_version_id and Dag %s has no version to stand in; " + "omitting last_ti from the Dag callback context.", + relevant_ti, + self.dag_id, + ) + return None + dag_version_id = latest_dag_version.id + self.log.info( + "Task instance %s and its run both have no dag_version_id (pre-versioning records); " + "reporting it under the latest version %s in the Dag callback context.", relevant_ti, - self.dag_id, + dag_version_id, ) - return None - self.log.warning( - "Task instance %s has no dag_version_id (pre-versioning record); " - "reporting it as %s in the Dag callback context.", - relevant_ti, - dag_version_id, - ) # The datamodel rejects a null version, so the stand-in has to be spliced in # before validation rather than copied onto a validated model. values = { diff --git a/airflow-core/tests/unit/models/test_dagrun.py b/airflow-core/tests/unit/models/test_dagrun.py index 3546b9ad470f2..850952ca182fe 100644 --- a/airflow-core/tests/unit/models/test_dagrun.py +++ b/airflow-core/tests/unit/models/test_dagrun.py @@ -4180,12 +4180,14 @@ def test_produce_dag_callback_stands_in_version_for_versionless_last_ti( dr = dag_maker.create_dagrun() dr.dag_model = DagModel.get_dagmodel(dag.dag_id, session=session) run_version_id = dr.created_dag_version_id - ti = dr.get_task_instance("test_task") + ti = dr.get_task_instance("test_task", session=session) ti.dag_version_id = None if not run_keeps_version: dr.created_dag_version_id = None # Newer than the run's, so the two fallback sources are distinguishable. - latest_version = DagVersion.write_dag(dag_id=dr.dag_id, bundle_name="dag_maker", version_number=2) + latest_version = DagVersion.write_dag( + dag_id=dr.dag_id, bundle_name="dag_maker", version_number=2, session=session + ) session.flush() assert latest_version.id != run_version_id @@ -4209,7 +4211,7 @@ def test_produce_dag_callback_drops_last_ti_when_dag_has_no_version(self, dag_ma dr = dag_maker.create_dagrun() dr.dag_model = DagModel.get_dagmodel(dag.dag_id, session=session) - ti = dr.get_task_instance("test_task") + ti = dr.get_task_instance("test_task", session=session) ti.dag_version_id = None dr.created_dag_version_id = None session.flush() @@ -4237,7 +4239,9 @@ def on_failure(context): dr = dag_maker.create_dagrun() dr.dag_model = DagModel.get_dagmodel(dag.dag_id, session=session) - ti = dr.get_task_instance("test_task") + # dag_maker writes exactly one version, so the run's is also the latest one. + expected_version_id = dr.created_dag_version_id + ti = dr.get_task_instance("test_task", session=session) ti.dag_version_id = None dr.created_dag_version_id = None session.flush() @@ -4258,10 +4262,7 @@ def on_failure(context): assert context_received["reason"] == "task_failure" assert context_received["ti"].task_id == "test_task" assert context_received["ti"].run_id == dr.run_id - assert ( - context_received["ti"].dag_version_id - == DagVersion.get_latest_version(dag.dag_id, session=session).id - ) + assert context_received["ti"].dag_version_id == expected_version_id @pytest.mark.parametrize("strip_dag_version", [False, True]) def test_produce_dag_callback_preserves_callers_transaction(self, dag_maker, session, strip_dag_version):