kaxil commented on code in PR #72721:
URL: https://github.com/apache/airflow/pull/72721#discussion_r3963857280
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1511,30 +1517,64 @@ 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
Review Comment:
`_build_callback_last_ti` deliberately requires `session` because
`create_session()` hands back the caller's own scoped session and then commits
and closes it, but `@provide_session` on the two wrappers puts that default
right back for anyone who omits it. All four production call sites now pass
`session`, so a required keyword-only arg on both would match the helper. Was
`@provide_session` here just for the existing `execute_dag_callbacks` callers
in `test_dag.py`?
##########
airflow-core/tests/unit/models/test_dagrun.py:
##########
@@ -4117,18 +4157,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
Review Comment:
`Session.close` is patched inside the `with` block, so nothing could have
expunged `ti` and this assertion holds no matter what the call did.
`mock_close.mock_calls == []` on the line above is already the real check.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
