Add Jinja template rendering and richer context for async deadline callbacks#64984
Add Jinja template rendering and richer context for async deadline callbacks#64984seanghaeli wants to merge 4 commits intoapache:mainfrom
Conversation
…llbacks 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 apache#55241 with a richer deadline context that enables useful templating like: AsyncCallback(my_func, kwargs={"msg": "DAG {{ dag_id }} missed at {{ ds }}"})
- 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
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)
|
@ramitkataria - I know you had some plans here a while back, does this work with what you have in mind? |
There was a problem hiding this comment.
Pull request overview
This PR enhances async deadline callback execution by (1) sending a richer, more “standard Airflow-like” context to TriggererCallbacks and (2) adding Jinja template rendering for kwargs passed to plain async function callbacks (while skipping Notifier classes that self-render).
Changes:
- Enrich deadline callback
contextpayload with DAG run metadata and common template keys (dag_id,run_id,ds,ts, etc.). - Render Jinja templates in callback kwargs for plain async function callbacks in
CallbackTrigger. - Add unit tests covering richer context expectations and template-rendering behavior (including Notifier skip logic).
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 3 comments.
| File | Description |
|---|---|
| airflow-core/src/airflow/triggers/callback.py | Adds notifier detection and Jinja rendering for function-callback kwargs. |
| airflow-core/src/airflow/models/deadline.py | Builds and attaches an enriched deadline context before queuing callbacks. |
| airflow-core/tests/unit/triggers/test_callback.py | Adds tests for template rendering behavior and helper functions. |
| airflow-core/tests/unit/models/test_deadline.py | Extends assertions to validate the enriched deadline context fields. |
| 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. | ||
| """ | ||
| # Use CallbackTrigger (which inherits Templater via BaseTrigger) to access | ||
| # render_template without importing airflow.sdk directly in core. | ||
| from jinja2.sandbox import SandboxedEnvironment | ||
|
|
||
| trigger = CallbackTrigger(callback_path="", callback_kwargs={}) | ||
| jinja_env = SandboxedEnvironment(cache_size=0) | ||
| return trigger.render_template(kwargs, cast("Any", context), jinja_env) |
There was a problem hiding this comment.
_render_callback_kwargs() builds a raw jinja2.sandbox.SandboxedEnvironment, which bypasses Airflow’s templating environment (custom sandbox behavior, filters like ds/ts, extensions, etc.). This can make template rendering for plain function callbacks behave differently than Notifier rendering (which uses Templater.get_template_env()). Consider using CallbackTrigger.get_template_env() (or importing Airflow’s SDK SandboxedEnvironment from airflow.sdk.definitions._internal.templater) instead of the raw Jinja environment, and drop the string-based cast("Any", ...) in favor of a proper type (or no cast).
| # Deadline-specific information | ||
| "deadline": { | ||
| "id": self.id, | ||
| "deadline_time": self.deadline_time, | ||
| "alert_name": self.deadline_alert.name if self.deadline_alert else None, | ||
| }, |
There was a problem hiding this comment.
alert_name is derived via self.deadline_alert.name, but Deadline.deadline_alert is not eager-loaded in the scheduler’s deadline query (it currently selectinloads only callback and dagrun). Since handle_miss() is called in a loop, this will trigger a per-deadline lazy-load query (N+1). Either eager-load Deadline.deadline_alert in the scheduler query, or avoid relationship access here by fetching the name in bulk/alongside the deadline rows.
| # 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 |
There was a problem hiding this comment.
This test validates the enriched context shape, but it mocks deadline_orm.callback.queue(), so it doesn’t exercise the new context through TriggererCallback.queue() → Trigger.from_object() → airflow.sdk.serde.serialize(). Since the enriched context now includes additional types (UUIDs/datetimes/nested dicts), a regression test should ensure the callback can be queued successfully and the trigger kwargs can be serialized/deserialized without error.
|
For context, the reason I previously added the context workaround to add "simple context" was because the triggerrer was not ready for fetching context. Now that #55068 is merged, we can handle context using the same approach that the triggerrer does for AsyncCallbacks and the way it works in the executor for SyncCallbacks. So I think the way I would want to do implement this would involve first reverting #55241 and making changes in triggerer and executor to support context in non-task workloads like callbacks. With the current approach, we are not using the |
|
@seanghaeli This PR has been converted to draft because it does not yet meet our Pull Request quality criteria. Issues found:
What to do next:
Converting a PR to draft is not a rejection — it is an invitation to bring the PR up to the project's standards so that maintainer review time is spent productively. There is no rush — take your time and work at your own pace. We appreciate your contribution and are happy to wait for updates. If you have questions, feel free to ask on the Airflow Slack. Note: This comment was drafted by an AI-assisted triage tool and may contain mistakes. Once you have addressed the points above, an Apache Airflow maintainer — a real person — will take the next look at your PR. We use this two-stage triage process so that our maintainers' limited time is spent where it matters most: the conversation with you. |
|
Quick follow-up to the triage comment above — one clarification on the "Unresolved review comments" item: Once you believe a thread has been addressed — whether by pushing a fix, or by replying in-thread with an explanation of why the suggestion doesn't apply — please mark the thread as resolved yourself by clicking the "Resolve conversation" button at the bottom of each thread. Reviewers don't auto-close their own threads, so an addressed-but-unresolved thread reads as "still waiting on the author" and keeps the PR from moving forward. The author doing the resolve-click is the expected convention on this project. Note: This comment was drafted by an AI-assisted triage tool and may contain mistakes. Once you have addressed the points above, an Apache Airflow maintainer — a real person — will take the next look at your PR. We use this two-stage triage process so that our maintainers' limited time is spent where it matters most: the conversation with you. |
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.