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..223fd3d05ee60 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,60 @@ 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. + + 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 + + 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 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) + 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, + 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 + if hasattr(relevant_ti, name) + } + values["dag_version_id"] = dag_version_id + return TIDataModel.model_validate(values) + def produce_dag_callback( self, dag: SerializedDAG, @@ -1515,23 +1575,18 @@ def produce_dag_callback( relevant_ti: TI | None = None, reason: str = "success", execute: bool = False, + *, + session: 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 +1599,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 +1609,25 @@ def produce_dag_callback( success=success, relevant_ti=relevant_ti, reason=reason, + session=session, ) return None 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, ): """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_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 ffa223508031b..850952ca182fe 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 @@ -4052,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 @@ -4084,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 @@ -4120,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 @@ -4145,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 @@ -4155,27 +4169,65 @@ 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) - ti = dr.get_task_instance("test_task") - # Simulate a task instance created before the dag_version table existed. + run_version_id = dr.created_dag_version_id + 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, session=session + ) session.flush() + assert latest_version.id != run_version_id - callback = dr.produce_dag_callback(dag=dag, success=False, relevant_ti=ti, reason="task_failure") + 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", session=session) + ti.dag_version_id = None + dr.created_dag_version_id = None + session.flush() + + # 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): @@ -4187,20 +4239,68 @@ 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() 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 == 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): + """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 == [] @pytest.mark.parametrize( ("multi_team", "team_name", "expected_tags"), @@ -4229,7 +4329,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)