kaxil commented on code in PR #66608:
URL: https://github.com/apache/airflow/pull/66608#discussion_r3350852732
##########
airflow-core/src/airflow/jobs/triggerer_job_runner.py:
##########
@@ -802,6 +807,45 @@ def _create_workload(
timeout_after=trigger.task_instance.trigger_timeout,
)
+ def _fetch_callback_dag_run_data(self, trigger: Trigger, *, session:
Session) -> dict | None:
+ """
+ Fetch DagRun data for deadline callback triggers.
+
+ When a callback trigger has dag_id/run_id stored in its associated
Callback.data,
+ fetch the DagRun and return serialized dag_run_data so the
TriggerRunner can build
+ a standard Context at runtime (same pattern as start_from_trigger).
+ """
+ from airflow.models.dagrun import DagRun
+
+ # The trigger's callback relationship stores the identifiers we need
+ if not trigger.callback:
+ return None
+
+ callback_data = trigger.callback.data
+ dag_id = callback_data.get("dag_id")
+ run_id = callback_data.get("run_id")
+ if not dag_id or not run_id:
+ return None
+
+ dagrun = session.scalar(select(DagRun).where(DagRun.dag_id == dag_id,
DagRun.run_id == run_id))
+ if not dagrun:
+ log.warning(
+ "Could not find DagRun for callback context",
+ dag_id=dag_id,
+ run_id=run_id,
+ trigger_id=trigger.id,
+ )
+ return None
Review Comment:
The executor path now fails with `sys.exit(1)` when context can't be
fetched, but the triggerer path here still returns None on a missing DagRun,
and `build_trigger_workloads` then sends `dag_run_data=None`, so the trigger
runs the callback with `self.context = None`. A deadline notification
templating `{{ dag_run }}` / `{{ deadline.deadline_time }}` renders empty while
the deadline is marked handled. Worth failing the trigger here too so it
retries, to match the executor side.
##########
shared/module_loading/src/airflow_shared/module_loading/__init__.py:
##########
@@ -131,6 +131,48 @@ def qualname(o: object | Callable, use_qualname: bool =
False, exclude_module: b
return name
+def load_mangled_dag_module(mod_name: str, file_path: str) -> bool:
+ """
+ Load a DAG bundle file into sys.modules under its mangled
unusual_prefix_{hash}_{stem} name.
+
+ The DAG processor assigns each DAG file a unique module name
+ (unusual_prefix_{sha1_of_filepath}_{stem}) to prevent import collisions
between bundles.
+ Callback paths are serialized using that mangled name, so both the
executor subprocess
+ and the triggerer must register the file under exactly that name before
calling
+ import_string().
+
+ Returns True if the module was registered, False if the file was not found
or load failed.
+ """
+ import importlib.machinery
Review Comment:
The callback.py imports did get hoisted, but `load_mangled_dag_module` still
has these four inline stdlib imports
(`importlib.machinery`/`importlib.util`/`sys`/`Path`), and `sys` is already
imported at module top on line 24. The executor-side
`_register_unusual_prefix_module` in callback_supervisor.py:373-375 also still
imports `Path` and `load_mangled_dag_module` inline. These can all move up.
##########
airflow-core/src/airflow/triggers/callback.py:
##########
@@ -41,6 +84,10 @@ def __init__(self, callback_path: str, callback_kwargs:
dict[str, Any] | None =
super().__init__()
self.callback_path = callback_path
self.callback_kwargs = callback_kwargs or {}
+ #: Set externally by TriggerRunner from workload dag_run_data before
run() is called,
+ #: the same pattern as task_instance is set for task-bound triggers.
Not serialized;
+ #: always None on a freshly deserialized trigger.
+ self.context: dict | None = None
Review Comment:
The reply said this was renamed to `_callback_context`, but at HEAD it's
still `self.context` here, in `run()` (line 106), and in the triggerer
assignment (`triggerer_job_runner.py:1329`); `_callback_context` isn't anywhere
in the tree. The docstring note landed, so just the rename is left, and the
generic-name shadowing risk from the original comment still stands.
##########
airflow-core/src/airflow/triggers/callback.py:
##########
@@ -51,14 +98,19 @@ def serialize(self) -> tuple[str, dict[str, Any]]:
async def run(self) -> AsyncIterator[TriggerEvent]:
try:
yield TriggerEvent({PAYLOAD_STATUS_KEY: CallbackState.RUNNING})
+ if self.callback_path.startswith("unusual_prefix_"):
+ await asyncio.to_thread(_ensure_bundle_module_registered,
self.callback_path)
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)
+
+ # self.context is set by TriggerRunner from workload.dag_run_data
before run() is called.
+ context = self.context
+
+ kwargs = dict(self.callback_kwargs)
Review Comment:
This drops the `callback_kwargs.pop("context", None)` that's on `main` (and
shipped in 3.2.0-3.2.2, where `deadline.py` stored the context inside
`callback.data["kwargs"]`). After upgrade, a deadline callback queued on 3.2.x
still has `context` in `callback_kwargs`: since `self.context` is None on a
deserialized trigger, the `else` branch runs `callback(**kwargs)`, which either
passes the stale stored context (callables accepting `context`/`**kwargs`) or
raises `TypeError: unexpected keyword argument 'context'` (callables that
don't). The PR description says old stored-context callbacks still work, so
this looks like a gap. Preferring `self.context` but falling back to a popped
`callback_kwargs["context"]`, and stripping `context` from the kwargs handed to
the callable, would honor that.
##########
airflow-core/src/airflow/triggers/callback.py:
##########
@@ -17,21 +17,64 @@
from __future__ import annotations
+import asyncio
import logging
+import sys
import traceback
from collections.abc import AsyncIterator
+from pathlib import Path
from typing import Any
-from airflow._shared.module_loading import accepts_context, import_string,
qualname
+from airflow._shared.module_loading import accepts_context, import_string,
load_mangled_dag_module, qualname
+from airflow.dag_processing.bundles.manager import DagBundlesManager
from airflow.models.callback import CallbackState
from airflow.triggers.base import BaseTrigger, TriggerEvent
+from airflow.utils.file import get_unique_dag_module_name
log = logging.getLogger(__name__)
PAYLOAD_STATUS_KEY = "state"
PAYLOAD_BODY_KEY = "body"
+def _ensure_bundle_module_registered(callback_path: str) -> None:
+ """
+ Register an unusual_prefix_{hash}_{stem} module so import_string can find
it.
+
+ The triggerer event loop doesn't run the DAG processor, so DAG files with
mangled
+ module names must be registered manually. Walks all configured bundles to
find the
+ file, then delegates the actual loading to _load_mangled_module in
callback_supervisor.
+ """
+ mod_name = callback_path.split(".")[0]
+ if mod_name in sys.modules:
+ return
+
+ # unusual_prefix_{hex40}_{stem} → {stem}.py
+ parts = mod_name.split("_", 3)
+ if len(parts) < 4:
+ return
+ stem = parts[3]
+
+ try:
+ for bundle in DagBundlesManager().get_all_dag_bundles():
+ try:
+ bundle.initialize()
+ file_path = Path(bundle.path) / f"{stem}.py"
Review Comment:
`get_unique_dag_module_name` hashes the full file path, but this
reconstructs the file as `bundle.path / {stem}.py` (bundle root only). For a
callback whose DAG lives in a subdirectory (`dags/team_a/my_dag.py`), the hash
check on line 67 never matches, the module is never registered, and
`import_string(self.callback_path)` fails with `CallbackState.FAILED`. Same
assumption in `_register_unusual_prefix_module` (callback_supervisor.py:386).
Walking the bundle tree (`Path(bundle.path).rglob(f"{stem}.py")`) and matching
on the recomputed hash would handle nested DAG files. The `unusual_prefix_`
stem-parsing is also duplicated across both helpers and could share one
function.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
