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 into
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 Converting to draft — this PR doesn't yet meet our Pull Request quality criteria.
See the linked criteria for how to fix each item, then mark the PR "Ready for review". This is not a rejection — just an invitation to bring the PR up to standard. No rush. 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. |
|
@seanghaeli This draft PR has been inactive for 13 days since the last triage comment and no response from the author. Closing to keep the queue clean. You are welcome to reopen this PR when you resume work, or to open a new one addressing the issues previously raised. There is no rush — take your time. 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. |
Renders Jinja2 templates in callback kwargs using the simple context
(dag_run, deadline info) that the scheduler passes to the triggerer.
This allows users to write callbacks like:
AsyncCallback(my_func, kwargs={"msg": "DAG {{ dag_id }} missed deadline"})
Key design decisions:
- Uses BaseTrigger's inherited render_template() (via Templater) for
consistent rendering behavior with the rest of Airflow
- Skips pre-rendering for Notifier classes since they handle their own
template rendering in __await__ via render_template_fields()
- Only renders when context is present (no-op for callbacks without context)
- Acknowledges this is interim: full context should be fetched via the
Execution API at execution time (tracked in PR apache#64984 TODO comments)
Addresses feedback from Ramit (PR apache#64984 comment) by documenting the
path forward while still providing template rendering value on top of
the existing simple context from PR apache#55241.
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.