Repository navigation
Render Jinja templates in deadline callback kwargs - #75363
seanghaeli wants to merge 5 commits into
Conversation
5ad2074 to
631d0e6
Compare
631d0e6 to
097465f
Compare
ferruzzi
left a comment
There was a problem hiding this comment.
Looks like it needs a rebase, and left a non-blocking question. Otherwise LGTM
097465f to
8578a27
Compare
ramitkataria
left a comment
There was a problem hiding this comment.
Thanks for picking this up again. I like storing only ids in the scheduler and building the context where the callback runs. However, I would like to ideally see this implemented the same way as tasks like we have discussed earlier:
Target: a deadline callback gets a task's context without the task fields, rendered the same way a task is.
- One context builder. Move the Dag-run part of
RuntimeTaskInstance.get_template_context(task_runner.py) into a function inexecution_time/context.pythat takes a Dag run and an optional Dag. It returnsdag_run,run_id, thelogical_date/ds/ts/data_interval_*keys,partition_key,partition_date,triggering_asset_events,macros,varandconn, plusdagandparamswhen it has the Dag (process_paramsneeds to accept no task).get_template_contextcalls it and adds the task keys, keeping its cached-dict behaviour. Both callback paths use it in place ofbuild_context_from_dag_run, and a test asserts every key it returns is also in a task context. - Render the way tasks do. For plain-function callbacks, render every kwarg the way
PythonOperatorrendersop_kwargs(Templater.render_template, nested values included), after the kwargs are decoded. Don't pre-render Notifier kwargs: Notifiers render their owntemplate_fields, so each value is rendered once. Use the Dag'sget_template_env()soStrictUndefinedand user macros/filters apply. Keep this PR's fallback of passing a value unrendered with a warning when it fails to render. - Load the Dag where tasks get it. Executor: parse the Dag file in the callback subprocess the way
task_runner.parsedoes. Triggerer: send the serialized Dag on the workload, as_create_workloaddoes forstart_from_trigger. Serialized Dags drop user macros/filters, so async callbacks share that limit withstart_from_triggeruntil it is fixed for both; document it. - Callback-specific keys. Set
context["callback"] = CallbackMeta(source=CallbackSource.DEADLINE)as Dag callbacks do. Keepcontext["deadline"]withdeadline_timeas a datetime, and add the sync path'sxcom_pullto the context before rendering. Addprev_*_successonce the Execution API can look it up by Dag run instead of by TI id.
| data["dag_id"] = self.dagrun.dag_id | ||
| data["run_id"] = self.dagrun.run_id | ||
| data["deadline_id"] = str(self.id) | ||
| data["deadline_time"] = self.deadline_time.isoformat() |
There was a problem hiding this comment.
This turns context["deadline"]["deadline_time"] from a datetime into a string on both paths. Callbacks doing date math on it break, and {{ deadline.deadline_time | ds }} raises AttributeError and is passed unrendered. Storing the datetime won't survive the workload JSON to remote workers, so I'd parse it back to an aware datetime when the context is built, with a test on the type
Move the keys get_template_context takes from the Dag run (dag_run, run_id, the logical date and data interval keys, partition_key, partition_date, triggering_asset_events, macros, var and conn) into build_dag_run_context in execution_time/context.py, with dag and params when it is given the Dag. get_template_context calls it and keeps its cached dict, and process_params can be called without a task. Deadline callbacks use it next.
The scheduler now stores only the Dag run and deadline ids with the callback, and both callback paths build the context from the Dag run with build_dag_run_context: the executor callback subprocess fetches the Dag run over the Execution API, and the triggerer sends it on the workload, as for start_from_trigger tasks. The context also has callback=CallbackMeta(source=CallbackSource.DEADLINE), deadline with deadline_time parsed back to an aware datetime, and on the sync path xcom_pull. A callback whose Dag run is gone fails on both paths.
The executor callback subprocess parses the Dag file the way task_runner.parse does, and the triggerer sends the serialized Dag on the workload, as _create_workload does for start_from_trigger tasks, so the context has dag and params with the run's conf applied. A callback defined in the Dag file reuses the module parsed for the context instead of running it again. SerializedDAG.get_template_env lets notifiers render with a serialized Dag. Serialized Dags don't keep user macros or filters, which async callbacks share with start_from_trigger tasks; this is documented.
Every kwarg of a function callback, nested values included, is rendered with
Templater.render_template after the kwargs are decoded, using the Dag's
get_template_env(), so StrictUndefined and the Dag's macros and filters
apply. A value that fails to render is passed on unrendered with a warning.
Notifier kwargs are not rendered beforehand: notifiers render their own
template_fields, so each value is rendered once. On the sync path xcom_pull is
in the context before rendering, so {{ xcom_pull('task') }} works.
context["deadline"] is now a NotRequired key of Context, like callback, so it is in KNOWN_CONTEXT_KEYS, and build_deadline_callback_context returns a Context. Like other context keys, a @task argument named deadline can now only default to None, which the newsfragment mentions.
8578a27 to
e70697d
Compare
Deadline callbacks get a Dag run snapshot that the scheduler stores in their kwargs when the deadline is missed, and Airflow never renders their kwargs, so
{{ dag_run.run_id }}reaches the callback as raw text. This replaces #64984. The scheduler now stores only the ids, and the callback gets a task's context without the task fields, built where the callback runs and rendered the same way a task's is. There's one commit per review point:build_dag_run_contextinexecution_time/context.pyis the Dag run part ofget_template_context, which calls it and keeps its cached dict.process_paramstakes no task. A test checks every key it returns is in a task context.callback=CallbackMeta(source=CallbackSource.DEADLINE),deadlinewithdeadline_timeparsed back to an aware datetime, andxcom_pullon the sync path. If the Dag run is gone, the callback fails on both paths.task_runner.parse, and the triggerer sends the serialized Dag on the workload like_create_workloaddoes forstart_from_trigger.op_kwargs:Templater.render_template, nested values included, after decoding, with the Dag'sget_template_env(). A value that fails to render is passed through unrendered with a warning. Notifier kwargs aren't pre-rendered, since Notifiers render their owntemplate_fields.deadlineis aNotRequiredkey onContext, so it's inKNOWN_CONTEXT_KEYS, and the callback builder returnsContext.Serialized Dags don't keep
user_defined_macrosoruser_defined_filters, so async callbacks can't use them, the same asstart_from_triggertasks. That's documented in the deadline alerts docs and will be fixed for both in a follow-up.prev_*_successisn't in the context yet, since the Execution API only looks it up by task instance. The executor reads the Dag run with the single-use callback token from #71192, and this adds no new reads on top of it.context["dag_run"]is now an object,deadline_timea datetime, anddeadlinea context key, so there's a significant newsfragment.E2E in Breeze (scheduler, LocalExecutor, dag-processor, triggerer, postgres): one Dag with a sync and an async function callback and a sync and an async Notifier, kwargs using
{{ dag_run.run_id }},{{ deadline.deadline_time | ds }}, a nested dict and a user macro, triggered with conf{"braces": "{{ ds }}"}. The asyncmacrowarning is the serialized Dag limit above, andconf={{ ds }}shows each Notifier value is rendered once:Was generative AI tooling used to co-author this PR?
Generated-by: Claude Code (Opus 5.5) following the guidelines