kaxil commented on code in PR #72721:
URL: https://github.com/apache/airflow/pull/72721#discussion_r3979587882
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1511,30 +1517,66 @@ 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
+
+ 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
Review Comment:
`_ensure_ti_has_dag_version_id` (`jobs/scheduler_job_runner.py:225`) already
solves this for `TaskCallbackRequest` and `EmailRequest` with the same
docstring reasoning, and it has four call sites in the file this diff also
touches. Your tier order looks right, it matches the existing convention at
`models/taskinstance.py:449-452`, but that helper differs on the other two
axes: it resolves latest-only, and it persists the backfill onto the row. The
persistence is the part that bites, because whichever callback fires first for
a legacy TI then decides what the other sees: once a heartbeat-purge callback
heals the row to `latest`, this branch never fires and the run's own version is
never reported, and the grid picks which serialized Dag to draw from
`ti.dag_version_id` (`api_fastapi/core_api/routes/ui/grid.py:417`). Converge on
one helper, or add a line saying why the Dag-callback path deliberately does
not persist?
##########
airflow-core/tests/unit/models/test_dagrun.py:
##########
@@ -4172,18 +4222,67 @@ 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
Review Comment:
This recomputes the expression the production fallback itself uses at
`dagrun.py:1535`, so a bug inside `get_latest_version` would keep the test
green. `dag_maker` creates exactly one version here, so capturing
`dr.created_dag_version_id` before line 4225 nulls it and asserting that
literal would pin the value independently.
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1511,30 +1517,66 @@ 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
+
+ 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:
+ 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(
Review Comment:
This message calls the row a pre-versioning record, but it covers both tiers
and they are different populations. Tier 2 (`get_latest_version`) is the
genuine Airflow 2 migration case, whereas tier 1 fires when the run has a
version and the TI does not, which is an Airflow 3 state that
`models/taskinstance.py:449-452` and `_pin_versionless_tis_to_run_version` both
exist to handle. So an operator reading the scheduler log gets told a current
row predates Dag versioning; splitting the two messages would keep the
diagnosis honest.
##########
airflow-core/tests/unit/models/test_dagrun.py:
##########
@@ -4138,27 +4152,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")
Review Comment:
`get_task_instance` without `session=` is the defaulting you call out at
line 4262 and fixed at 4044, 4080, 4118 and 4144. Here it commits and closes
the fixture session, so `ti` and `dr` come back detached: I ran this shape in
breeze and `ti in session` is `False`, and after `session.flush()` the row
still holds the run's version while the in-memory attribute is `None`. The
assertions pass because `_build_callback_last_ti` reads the attribute, which
means the tests carrying the PR's new coverage never put a versionless row in
the database at all. Same call shape at 4195 and 4223, and `write_dag` at 4171
closes it a second time, so `session=session` on all four would make these
`flush()` calls mean something.
--
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]
