This is an automated email from the ASF dual-hosted git repository.
eladkal pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 57cff7c409c Fix deadline callback data persistence (#69073)
57cff7c409c is described below
commit 57cff7c409c158e1c1d0458db11e4c0d00f79f4c
Author: Shivam Rastogi <[email protected]>
AuthorDate: Thu Jul 2 07:28:03 2026 -0700
Fix deadline callback data persistence (#69073)
* Fix deadline callback data persistence
* Clarify deadline callback data updates
---
airflow-core/src/airflow/models/deadline.py | 29 ++++++-----
airflow-core/tests/unit/models/test_deadline.py | 65 ++++++++++++++++++++++++-
2 files changed, 78 insertions(+), 16 deletions(-)
diff --git a/airflow-core/src/airflow/models/deadline.py
b/airflow-core/src/airflow/models/deadline.py
index f596c663014..67438ce428a 100644
--- a/airflow-core/src/airflow/models/deadline.py
+++ b/airflow-core/src/airflow/models/deadline.py
@@ -242,30 +242,29 @@ class Deadline(Base):
return {
"dag_run":
DAGRunResponse.model_validate(dagrun).model_dump(mode="json"),
- "deadline": {"id": self.id, "deadline_time":
self.deadline_time},
+ "deadline": {"id": str(self.id), "deadline_time":
self.deadline_time},
}
+ def callback_data_with_context():
+ data = self.callback.data.copy()
+ kwargs = dict(data.get("kwargs") or {})
+ kwargs["context"] = get_simple_context()
+ data["kwargs"] = kwargs
+ return data
+
if isinstance(self.callback, TriggererCallback):
- # Update the callback with context before queuing
- if "kwargs" not in self.callback.data:
- self.callback.data["kwargs"] = {}
- self.callback.data["kwargs"] = (self.callback.data.get("kwargs")
or {}) | {
- "context": get_simple_context()
- }
+ self.callback.data = callback_data_with_context()
self.callback.queue(session=session)
session.add(self.callback)
session.flush()
elif isinstance(self.callback, ExecutorCallback):
- if "kwargs" not in self.callback.data:
- self.callback.data["kwargs"] = {}
- self.callback.data["kwargs"] = (self.callback.data.get("kwargs")
or {}) | {
- "context": get_simple_context()
- }
- self.callback.data["deadline_id"] = str(self.id)
- self.callback.data["dag_run_id"] = str(self.dagrun.id)
- self.callback.data["dag_id"] = self.dagrun.dag_id
+ data = callback_data_with_context()
+ data["deadline_id"] = str(self.id)
+ data["dag_run_id"] = str(self.dagrun.id)
+ data["dag_id"] = self.dagrun.dag_id
+ self.callback.data = data
self.callback.state = CallbackState.PENDING
session.add(self.callback)
diff --git a/airflow-core/tests/unit/models/test_deadline.py
b/airflow-core/tests/unit/models/test_deadline.py
index 4fc1aada1cc..99f3d9b127b 100644
--- a/airflow-core/tests/unit/models/test_deadline.py
+++ b/airflow-core/tests/unit/models/test_deadline.py
@@ -252,10 +252,73 @@ class TestDeadline:
context = callback_kwargs.pop("context")
assert callback_kwargs == TEST_CALLBACK_KWARGS
- assert context["deadline"]["id"] == deadline_orm.id
+ assert context["deadline"]["id"] == str(deadline_orm.id)
assert context["deadline"]["deadline_time"].timestamp() ==
deadline_orm.deadline_time.timestamp()
assert context["dag_run"] ==
DAGRunResponse.model_validate(dagrun).model_dump(mode="json")
+ @pytest.mark.db_test
+ def test_handle_miss_persists_triggerer_callback_context(self, dagrun,
session):
+ deadline_orm = Deadline(
+ deadline_time=DEFAULT_DATE,
+ callback=AsyncCallback(TEST_CALLBACK_PATH, TEST_CALLBACK_KWARGS),
+ dagrun_id=dagrun.id,
+ dag_id=dagrun.dag_id,
+ deadline_alert_id=None,
+ )
+ session.add(deadline_orm)
+ session.flush()
+
+ callback_id = deadline_orm.callback.id
+ deadline_id = deadline_orm.id
+ deadline_time = deadline_orm.deadline_time
+ expected_dag_run =
DAGRunResponse.model_validate(dagrun).model_dump(mode="json")
+
+ deadline_orm.handle_miss(session)
+ session.commit()
+ session.expunge_all()
+
+ callback = session.scalar(select(Deadline).where(Deadline.id ==
deadline_id)).callback
+ assert callback.id == callback_id
+
+ callback_kwargs = callback.data["kwargs"]
+ context = callback_kwargs["context"]
+ assert {
+ key: value for key, value in callback_kwargs.items() if key !=
"context"
+ } == TEST_CALLBACK_KWARGS
+ assert context["deadline"]["id"] == str(deadline_id)
+ assert context["deadline"]["deadline_time"].timestamp() ==
deadline_time.timestamp()
+ assert context["dag_run"] == expected_dag_run
+
+ callback.trigger = None
+ session.commit()
+
+ @pytest.mark.db_test
+ def test_handle_miss_persists_executor_callback_routing_data(self, dagrun,
session):
+ deadline_orm = Deadline(
+ deadline_time=DEFAULT_DATE,
+ callback=SyncCallback(TEST_CALLBACK_PATH, TEST_CALLBACK_KWARGS),
+ dagrun_id=dagrun.id,
+ dag_id=dagrun.dag_id,
+ deadline_alert_id=None,
+ )
+ session.add(deadline_orm)
+ session.flush()
+
+ callback_id = deadline_orm.callback.id
+ deadline_id = deadline_orm.id
+ dagrun_id = dagrun.id
+ dag_id = dagrun.dag_id
+
+ deadline_orm.handle_miss(session)
+ session.commit()
+ session.expunge_all()
+
+ callback = session.scalar(select(Deadline).where(Deadline.id ==
deadline_id)).callback
+ assert callback.id == callback_id
+ assert callback.data["dag_run_id"] == str(dagrun_id)
+ assert callback.data["dag_id"] == dag_id
+ assert callback.data["deadline_id"] == str(deadline_id)
+
@pytest.mark.db_test
class TestCalculatedDeadlineDatabaseCalls:
