From 72a347c938c8588fa9382d4c2e81cf6926a2af99 Mon Sep 17 00:00:00 2001 From: Sean Ghaeli Date: Thu, 9 Apr 2026 22:26:06 +0000 Subject: [PATCH 1/4] Add Jinja template rendering and richer context for async deadline callbacks Async deadline callbacks (TriggererCallbacks) now receive a full context with dag_id, run_id, logical_date, ds, ts, conf, and other standard template variables. Plain function callbacks get their kwargs rendered with Jinja2 before execution. Notifier classes are skipped since they self-render via __await__. Replaces the minimal "simple context" from PR #55241 with a richer deadline context that enables useful templating like: AsyncCallback(my_func, kwargs={"msg": "DAG {{ dag_id }} missed at {{ ds }}"}) --- airflow-core/src/airflow/models/deadline.py | 42 ++++- airflow-core/src/airflow/triggers/callback.py | 39 +++- .../tests/unit/models/test_deadline.py | 19 +- .../tests/unit/triggers/test_callback.py | 176 +++++++++++++++++- 4 files changed, 259 insertions(+), 17 deletions(-) diff --git a/airflow-core/src/airflow/models/deadline.py b/airflow-core/src/airflow/models/deadline.py index 6a94b21223fc0..3d032cacb2662 100644 --- a/airflow-core/src/airflow/models/deadline.py +++ b/airflow-core/src/airflow/models/deadline.py @@ -216,28 +216,50 @@ def prune_deadlines(cls, *, session: Session, conditions: dict[Mapped, Any]) -> def handle_miss(self, session: Session): """Handle a missed deadline by queueing the callback.""" - def get_simple_context(): + def _build_deadline_context(): from airflow.api_fastapi.core_api.datamodels.dag_run import DAGRunResponse from airflow.models import DagRun - # TODO: Use the TaskAPI from within Triggerer to fetch full context instead of sending this context - # from the scheduler - - # Fetch the DagRun from the database again to avoid errors when self.dagrun's relationship fields - # are not in the current session. + # Fetch the DagRun from the database again to avoid errors when self.dagrun's + # relationship fields are not in the current session. dagrun = session.get(DagRun, self.dagrun_id) + logical_date = dagrun.logical_date - return { + context: dict[str, Any] = { + # Full DAGRunResponse as a JSON-serializable dict "dag_run": DAGRunResponse.model_validate(dagrun).model_dump(mode="json"), - "deadline": {"id": self.id, "deadline_time": self.deadline_time}, + # Top-level convenience keys for Jinja templates (match standard context naming) + "dag_id": dagrun.dag_id, + "run_id": dagrun.run_id, + "logical_date": logical_date, + "data_interval_start": dagrun.data_interval_start, + "data_interval_end": dagrun.data_interval_end, + "run_type": dagrun.run_type, + "conf": dagrun.conf or {}, + # Deadline-specific information + "deadline": { + "id": self.id, + "deadline_time": self.deadline_time, + "alert_name": self.deadline_alert.name if self.deadline_alert else None, + }, } + # Derived date/time template variables + if logical_date is not None: + context["ds"] = logical_date.strftime("%Y-%m-%d") + context["ds_nodash"] = logical_date.strftime("%Y%m%d") + context["ts"] = logical_date.isoformat() + context["ts_nodash"] = logical_date.strftime("%Y%m%dT%H%M%S") + context["ts_nodash_with_tz"] = logical_date.isoformat().replace("-", "").replace(":", "") + + return context + 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() + "context": _build_deadline_context() } self.callback.queue() @@ -248,7 +270,7 @@ def get_simple_context(): 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() + "context": _build_deadline_context() } self.callback.data["deadline_id"] = str(self.id) self.callback.data["dag_run_id"] = str(self.dagrun.id) diff --git a/airflow-core/src/airflow/triggers/callback.py b/airflow-core/src/airflow/triggers/callback.py index 9c2470c77eae6..45d2cd6c8ac3e 100644 --- a/airflow-core/src/airflow/triggers/callback.py +++ b/airflow-core/src/airflow/triggers/callback.py @@ -32,6 +32,38 @@ PAYLOAD_BODY_KEY = "body" +def _is_notifier_class(callback: Any) -> bool: + """Check if the callback is a BaseNotifier subclass (not an instance).""" + try: + from airflow.sdk.bases.notifier import BaseNotifier + + return isinstance(callback, type) and issubclass(callback, BaseNotifier) + except ImportError: + return False + + +def _render_callback_kwargs(kwargs: dict[str, Any], context: dict) -> dict[str, Any]: + """ + Render Jinja2 templates in callback kwargs using the provided context. + + Uses ``Templater.render_template`` to recursively render all string values + in the kwargs dict. Non-string values (int, float, datetime, …) pass + through unchanged. + """ + from typing import TYPE_CHECKING, cast + + from airflow.sdk.definitions._internal.templater import SandboxedEnvironment, Templater + + if TYPE_CHECKING: + from airflow.sdk.definitions.context import Context + + templater = Templater() + templater.template_fields = () + templater.template_ext = () + jinja_env = SandboxedEnvironment(cache_size=0) + return templater.render_template(kwargs, cast("Context", context), jinja_env) + + class CallbackTrigger(BaseTrigger): """Trigger that executes a callback function asynchronously.""" @@ -52,9 +84,14 @@ async def run(self) -> AsyncIterator[TriggerEvent]: try: yield TriggerEvent({PAYLOAD_STATUS_KEY: CallbackState.RUNNING}) callback = import_string(self.callback_path) - # TODO: get full context and run template rendering. Right now, a simple context is included in `callback_kwargs` context = self.callback_kwargs.pop("context", None) + # Render Jinja templates in kwargs for plain function callbacks. + # Notifiers handle their own template rendering in __await__ via + # render_template_fields(), so we skip rendering here for them. + if context is not None and not _is_notifier_class(callback): + self.callback_kwargs = _render_callback_kwargs(self.callback_kwargs, context) + if _accepts_context(callback) and context is not None: result = await callback(**self.callback_kwargs, context=context) else: diff --git a/airflow-core/tests/unit/models/test_deadline.py b/airflow-core/tests/unit/models/test_deadline.py index 94c6977ae0c16..0cf3b00e0c564 100644 --- a/airflow-core/tests/unit/models/test_deadline.py +++ b/airflow-core/tests/unit/models/test_deadline.py @@ -234,9 +234,26 @@ def test_handle_miss(self, dagrun, session): context = callback_kwargs.pop("context") assert callback_kwargs == TEST_CALLBACK_KWARGS + # Verify enriched context — dag_run and deadline info + assert context["dag_run"] == DAGRunResponse.model_validate(dagrun).model_dump(mode="json") assert context["deadline"]["id"] == 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") + assert context["deadline"]["alert_name"] is None # no deadline_alert in this test + + # Verify top-level convenience keys + assert context["dag_id"] == dagrun.dag_id + assert context["run_id"] == dagrun.run_id + assert context["logical_date"] == dagrun.logical_date + assert context["data_interval_start"] == dagrun.data_interval_start + assert context["data_interval_end"] == dagrun.data_interval_end + assert context["run_type"] == dagrun.run_type + assert context["conf"] == (dagrun.conf or {}) + + # Verify derived template variables + assert context["ds"] == dagrun.logical_date.strftime("%Y-%m-%d") + assert context["ds_nodash"] == dagrun.logical_date.strftime("%Y%m%d") + assert context["ts"] == dagrun.logical_date.isoformat() + assert context["ts_nodash"] == dagrun.logical_date.strftime("%Y%m%dT%H%M%S") @pytest.mark.db_test diff --git a/airflow-core/tests/unit/triggers/test_callback.py b/airflow-core/tests/unit/triggers/test_callback.py index 99eca603323bb..77797a210a3b4 100644 --- a/airflow-core/tests/unit/triggers/test_callback.py +++ b/airflow-core/tests/unit/triggers/test_callback.py @@ -23,16 +23,32 @@ from airflow.models.callback import CallbackState from airflow.sdk import BaseNotifier -from airflow.triggers.callback import PAYLOAD_BODY_KEY, PAYLOAD_STATUS_KEY, CallbackTrigger +from airflow.triggers.callback import ( + PAYLOAD_BODY_KEY, + PAYLOAD_STATUS_KEY, + CallbackTrigger, + _is_notifier_class, + _render_callback_kwargs, +) TEST_MESSAGE = "test_message" TEST_CALLBACK_PATH = "classpath.test_callback" -TEST_CALLBACK_KWARGS = {"message": TEST_MESSAGE, "context": {"dag_run": "test"}} +TEST_CONTEXT = { + "dag_run": {"dag_id": "test_dag"}, + "dag_id": "test_dag", + "run_id": "test_run", + "ds": "2024-01-01", + "ts": "2024-01-01T00:00:00+00:00", + "deadline": {"id": "abc-123", "deadline_time": "2024-01-01T01:00:00+00:00", "alert_name": "my_alert"}, +} +TEST_CALLBACK_KWARGS = {"message": TEST_MESSAGE, "context": TEST_CONTEXT} class ExampleAsyncNotifier(BaseNotifier): """Example of a properly implemented async notifier.""" + template_fields = ("message",) + def __init__(self, message, **kwargs): super().__init__(**kwargs) self.message = message @@ -93,7 +109,7 @@ async def test_run_success_with_async_function(self, trigger, mock_import_string success_event = await anext(trigger_gen) mock_import_string.assert_called_once_with(TEST_CALLBACK_PATH) # AsyncMock accepts **kwargs, so _accepts_context returns True and context is passed through - mock_callback.assert_called_once_with(**TEST_CALLBACK_KWARGS) + mock_callback.assert_called_once_with(message=TEST_MESSAGE, context=TEST_CONTEXT) assert success_event.payload[PAYLOAD_STATUS_KEY] == CallbackState.SUCCESS assert success_event.payload[PAYLOAD_BODY_KEY] == callback_return_value @@ -112,7 +128,7 @@ async def test_run_success_with_notifier(self, trigger, mock_import_string): assert success_event.payload[PAYLOAD_STATUS_KEY] == CallbackState.SUCCESS assert ( success_event.payload[PAYLOAD_BODY_KEY] - == f"Async notification: {TEST_MESSAGE}, context: {{'dag_run': 'test'}}" + == f"Async notification: {TEST_MESSAGE}, context: {TEST_CONTEXT}" ) @pytest.mark.asyncio @@ -129,6 +145,156 @@ async def test_run_failure(self, trigger, mock_import_string): failure_event = await anext(trigger_gen) mock_import_string.assert_called_once_with(TEST_CALLBACK_PATH) # AsyncMock accepts **kwargs, so _accepts_context returns True and context is passed through - mock_callback.assert_called_once_with(**TEST_CALLBACK_KWARGS) + mock_callback.assert_called_once_with(message=TEST_MESSAGE, context=TEST_CONTEXT) assert failure_event.payload[PAYLOAD_STATUS_KEY] == CallbackState.FAILED assert all(s in failure_event.payload[PAYLOAD_BODY_KEY] for s in ["raise", "RuntimeError", exc_msg]) + + +class TestTemplateRendering: + """Tests for Jinja2 template rendering in callback kwargs.""" + + @pytest.mark.asyncio + async def test_run_renders_jinja_templates_in_function_kwargs(self): + """Plain async function callbacks get their kwargs rendered.""" + context = {"dag_id": "my_dag", "ds": "2024-06-15"} + trigger = CallbackTrigger( + callback_path="classpath.test", + callback_kwargs={ + "message": "DAG {{ dag_id }} missed deadline at {{ ds }}", + "context": context, + }, + ) + mock_callback = mock.AsyncMock(return_value="ok") + with mock.patch("airflow.triggers.callback.import_string", return_value=mock_callback): + events = [event async for event in trigger.run()] + + assert events[-1].payload[PAYLOAD_STATUS_KEY] == CallbackState.SUCCESS + mock_callback.assert_called_once_with( + message="DAG my_dag missed deadline at 2024-06-15", + context=context, + ) + + @pytest.mark.asyncio + async def test_run_does_not_double_render_notifier_kwargs(self): + """Notifier classes should NOT have kwargs pre-rendered — they handle it themselves.""" + context = {"dag_id": "my_dag", "ds": "2024-06-15"} + trigger = CallbackTrigger( + callback_path="classpath.test", + callback_kwargs={ + "message": "DAG {{ dag_id }}", + "context": context, + }, + ) + with mock.patch("airflow.triggers.callback.import_string", return_value=ExampleAsyncNotifier): + events = [event async for event in trigger.run()] + + assert events[-1].payload[PAYLOAD_STATUS_KEY] == CallbackState.SUCCESS + # The notifier's __await__ renders template_fields, so the final output + # should show the rendered message (rendered by the notifier, not pre-rendered). + assert "DAG my_dag" in events[-1].payload[PAYLOAD_BODY_KEY] + + @pytest.mark.asyncio + async def test_run_renders_nested_kwargs(self): + """Template rendering works recursively on nested dicts and lists.""" + context = {"dag_id": "etl_pipeline"} + trigger = CallbackTrigger( + callback_path="classpath.test", + callback_kwargs={ + "recipients": ["{{ dag_id }}-team@example.com"], + "metadata": {"dag": "{{ dag_id }}"}, + "context": context, + }, + ) + mock_callback = mock.AsyncMock(return_value="ok") + with mock.patch("airflow.triggers.callback.import_string", return_value=mock_callback): + events = [event async for event in trigger.run()] + + assert events[-1].payload[PAYLOAD_STATUS_KEY] == CallbackState.SUCCESS + mock_callback.assert_called_once_with( + recipients=["etl_pipeline-team@example.com"], + metadata={"dag": "etl_pipeline"}, + context=context, + ) + + @pytest.mark.asyncio + async def test_run_skips_rendering_when_no_context(self): + """Without context, kwargs pass through unrendered.""" + trigger = CallbackTrigger( + callback_path="classpath.test", + callback_kwargs={"message": "{{ dag_id }}"}, + ) + mock_callback = mock.AsyncMock(return_value="ok") + with mock.patch("airflow.triggers.callback.import_string", return_value=mock_callback): + events = [event async for event in trigger.run()] + + assert events[-1].payload[PAYLOAD_STATUS_KEY] == CallbackState.SUCCESS + mock_callback.assert_called_once_with(message="{{ dag_id }}") + + @pytest.mark.asyncio + async def test_notifier_template_fields_rendered_with_context(self): + """Notifier template_fields are rendered using the provided context.""" + context = {"dag_id": "my_dag", "ds": "2024-06-15"} + trigger = CallbackTrigger( + callback_path="classpath.test", + callback_kwargs={ + "message": "Alert for {{ dag_id }} on {{ ds }}", + "context": context, + }, + ) + with mock.patch("airflow.triggers.callback.import_string", return_value=ExampleAsyncNotifier): + events = [event async for event in trigger.run()] + + assert events[-1].payload[PAYLOAD_STATUS_KEY] == CallbackState.SUCCESS + # The notifier's __await__ renders template_fields (self.message), so the + # notification body contains the rendered message. + assert "Alert for my_dag on 2024-06-15" in events[-1].payload[PAYLOAD_BODY_KEY] + + +class TestHelpers: + """Tests for module-level helper functions.""" + + def test_is_notifier_class_with_notifier(self): + assert _is_notifier_class(ExampleAsyncNotifier) is True + + def test_is_notifier_class_with_function(self): + async def my_func(): + pass + + assert _is_notifier_class(my_func) is False + + def test_is_notifier_class_with_non_notifier_class(self): + class MyClass: + pass + + assert _is_notifier_class(MyClass) is False + + def test_is_notifier_class_with_notifier_instance(self): + """Instances are not classes — should return False.""" + instance = ExampleAsyncNotifier(message="hi") + assert _is_notifier_class(instance) is False + + def test_render_callback_kwargs_renders_strings(self): + result = _render_callback_kwargs( + {"message": "Hello {{ name }}", "count": 5}, + {"name": "World"}, + ) + assert result == {"message": "Hello World", "count": 5} + + def test_render_callback_kwargs_handles_nested_structures(self): + result = _render_callback_kwargs( + {"items": ["{{ x }}", "{{ y }}"], "meta": {"key": "{{ x }}"}}, + {"x": "a", "y": "b"}, + ) + assert result == {"items": ["a", "b"], "meta": {"key": "a"}} + + def test_render_callback_kwargs_missing_key_renders_empty(self): + result = _render_callback_kwargs( + {"message": "Hello {{ nonexistent }}"}, + {"name": "World"}, + ) + assert result == {"message": "Hello "} + + def test_render_callback_kwargs_no_templates_is_noop(self): + kwargs = {"message": "plain text", "count": 42} + result = _render_callback_kwargs(kwargs, {"dag_id": "test"}) + assert result == kwargs From 131ccca507f15c67263578f41d4f29956dcd0794 Mon Sep 17 00:00:00 2001 From: Sean Ghaeli Date: Thu, 9 Apr 2026 22:53:19 +0000 Subject: [PATCH 2/4] Fix review findings: add missing test assertion, clean up imports - Add ts_nodash_with_tz assertion to test_handle_miss - Move cast and TYPE_CHECKING imports to module level in callback.py - Remove dead TYPE_CHECKING block inside _render_callback_kwargs --- airflow-core/src/airflow/triggers/callback.py | 10 ++++------ airflow-core/tests/unit/models/test_deadline.py | 3 +++ 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/airflow-core/src/airflow/triggers/callback.py b/airflow-core/src/airflow/triggers/callback.py index 45d2cd6c8ac3e..3f53024b0d1fd 100644 --- a/airflow-core/src/airflow/triggers/callback.py +++ b/airflow-core/src/airflow/triggers/callback.py @@ -20,12 +20,15 @@ import logging import traceback from collections.abc import AsyncIterator -from typing import Any +from typing import TYPE_CHECKING, Any, cast from airflow._shared.module_loading import import_string, qualname from airflow.models.callback import CallbackState, _accepts_context from airflow.triggers.base import BaseTrigger, TriggerEvent +if TYPE_CHECKING: + from airflow.sdk.definitions.context import Context + log = logging.getLogger(__name__) PAYLOAD_STATUS_KEY = "state" @@ -50,13 +53,8 @@ def _render_callback_kwargs(kwargs: dict[str, Any], context: dict) -> dict[str, in the kwargs dict. Non-string values (int, float, datetime, …) pass through unchanged. """ - from typing import TYPE_CHECKING, cast - from airflow.sdk.definitions._internal.templater import SandboxedEnvironment, Templater - if TYPE_CHECKING: - from airflow.sdk.definitions.context import Context - templater = Templater() templater.template_fields = () templater.template_ext = () diff --git a/airflow-core/tests/unit/models/test_deadline.py b/airflow-core/tests/unit/models/test_deadline.py index 0cf3b00e0c564..d60a5a194641a 100644 --- a/airflow-core/tests/unit/models/test_deadline.py +++ b/airflow-core/tests/unit/models/test_deadline.py @@ -254,6 +254,9 @@ def test_handle_miss(self, dagrun, session): assert context["ds_nodash"] == dagrun.logical_date.strftime("%Y%m%d") assert context["ts"] == dagrun.logical_date.isoformat() assert context["ts_nodash"] == dagrun.logical_date.strftime("%Y%m%dT%H%M%S") + assert context["ts_nodash_with_tz"] == dagrun.logical_date.isoformat().replace("-", "").replace( + ":", "" + ) @pytest.mark.db_test From 04d524eb5f39fd5d1ea779fe451c5e59403562de Mon Sep 17 00:00:00 2001 From: Sean Ghaeli Date: Fri, 10 Apr 2026 00:51:56 +0000 Subject: [PATCH 3/4] Fix check-sdk-imports: remove airflow.sdk imports from core trigger Replace direct airflow.sdk imports with alternatives: - BaseNotifier check: duck-typing via hasattr instead of isinstance - Templater: reuse CallbackTrigger (inherits Templater via BaseTrigger) - SandboxedEnvironment: import from jinja2.sandbox directly - Context type: removed (no longer needed) --- airflow-core/src/airflow/triggers/callback.py | 33 ++++++++++--------- 1 file changed, 18 insertions(+), 15 deletions(-) diff --git a/airflow-core/src/airflow/triggers/callback.py b/airflow-core/src/airflow/triggers/callback.py index 3f53024b0d1fd..830d374d38d2b 100644 --- a/airflow-core/src/airflow/triggers/callback.py +++ b/airflow-core/src/airflow/triggers/callback.py @@ -17,18 +17,16 @@ from __future__ import annotations +import inspect import logging import traceback from collections.abc import AsyncIterator -from typing import TYPE_CHECKING, Any, cast +from typing import Any from airflow._shared.module_loading import import_string, qualname from airflow.models.callback import CallbackState, _accepts_context from airflow.triggers.base import BaseTrigger, TriggerEvent -if TYPE_CHECKING: - from airflow.sdk.definitions.context import Context - log = logging.getLogger(__name__) PAYLOAD_STATUS_KEY = "state" @@ -36,13 +34,18 @@ def _is_notifier_class(callback: Any) -> bool: - """Check if the callback is a BaseNotifier subclass (not an instance).""" - try: - from airflow.sdk.bases.notifier import BaseNotifier + """ + Check if the callback is a BaseNotifier subclass (not an instance). - return isinstance(callback, type) and issubclass(callback, BaseNotifier) - except ImportError: - return False + Uses duck-typing (checks for ``async_notify`` and ``template_fields``) + to avoid importing ``airflow.sdk`` in core. + """ + return ( + inspect.isclass(callback) + and hasattr(callback, "async_notify") + and hasattr(callback, "template_fields") + and hasattr(callback, "__await__") + ) def _render_callback_kwargs(kwargs: dict[str, Any], context: dict) -> dict[str, Any]: @@ -53,13 +56,13 @@ def _render_callback_kwargs(kwargs: dict[str, Any], context: dict) -> dict[str, in the kwargs dict. Non-string values (int, float, datetime, …) pass through unchanged. """ - from airflow.sdk.definitions._internal.templater import SandboxedEnvironment, Templater + # Use CallbackTrigger (which inherits Templater via BaseTrigger) to access + # render_template without importing airflow.sdk directly in core. + from jinja2.sandbox import SandboxedEnvironment - templater = Templater() - templater.template_fields = () - templater.template_ext = () + trigger = CallbackTrigger(callback_path="", callback_kwargs={}) jinja_env = SandboxedEnvironment(cache_size=0) - return templater.render_template(kwargs, cast("Context", context), jinja_env) + return trigger.render_template(kwargs, context, jinja_env) class CallbackTrigger(BaseTrigger): From 9630ad9a20c8290af45993153ef102d38af20ba5 Mon Sep 17 00:00:00 2001 From: Sean Ghaeli Date: Fri, 10 Apr 2026 07:14:02 +0000 Subject: [PATCH 4/4] Fix mypy: cast context arg to satisfy Templater.render_template type --- airflow-core/src/airflow/triggers/callback.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/airflow-core/src/airflow/triggers/callback.py b/airflow-core/src/airflow/triggers/callback.py index 830d374d38d2b..a32819971bf06 100644 --- a/airflow-core/src/airflow/triggers/callback.py +++ b/airflow-core/src/airflow/triggers/callback.py @@ -21,7 +21,7 @@ import logging import traceback from collections.abc import AsyncIterator -from typing import Any +from typing import Any, cast from airflow._shared.module_loading import import_string, qualname from airflow.models.callback import CallbackState, _accepts_context @@ -62,7 +62,7 @@ def _render_callback_kwargs(kwargs: dict[str, Any], context: dict) -> dict[str, trigger = CallbackTrigger(callback_path="", callback_kwargs={}) jinja_env = SandboxedEnvironment(cache_size=0) - return trigger.render_template(kwargs, context, jinja_env) + return trigger.render_template(kwargs, cast("Any", context), jinja_env) class CallbackTrigger(BaseTrigger):