kaxil commented on code in PR #66608:
URL: https://github.com/apache/airflow/pull/66608#discussion_r3366549043


##########
airflow-core/src/airflow/jobs/triggerer_job_runner.py:
##########
@@ -742,10 +742,23 @@ def _create_workload(
         session: Session,
     ) -> workloads.RunTrigger | None:
         if trigger.task_instance is None:
+            # Only fetch DagRun data for callback triggers (not all non-task 
triggers).
+            dag_run_data = (
+                self._fetch_callback_dag_run_data(trigger, session=session) if 
trigger.callback else None
+            )
+            if trigger.callback and dag_run_data is None:

Review Comment:
   On 3.2.x the `TriggererCallback` branch of `handle_miss` stored context 
inside `data["kwargs"]["context"]` and set no top-level `dag_id`/`run_id` 
(verified on `v3-2-stable`). After a rolling upgrade, 
`_fetch_callback_dag_run_data` reads `callback_data.get("dag_id")`, gets None, 
and returns None, so this gate skips the trigger on every loop and 
`CallbackTrigger.run()` (with its `self._callback_context or 
kwargs.pop("context")` compat fallback) is never reached, because the workload 
is never built. The deadline alert then stalls silently, with no FAILED state. 
`test_run_with_stored_context_in_kwargs_32x_compat` calls `run()` directly, so 
it passes while this path stays broken. Building the workload with 
`dag_run_data=None` when `dag_id`/`run_id` are absent (letting the compat 
fallback run), and only skipping/retrying on an actual DagRun lookup miss, 
would keep those rows alive.



##########
task-sdk/src/airflow/sdk/execution_time/callback_supervisor.py:
##########
@@ -209,14 +218,36 @@ def _target():
                         _log.debug(
                             "Added bundle path to sys.path", 
bundle_name=bundle_info.name, path=bundle_path
                         )
+                    # DAG processor loads bundle files with a mangled module 
name
+                    # (unusual_prefix_{hash}_{stem}) to avoid collisions. The 
callback path
+                    # was serialized using that mangled name. Register the 
module under that
+                    # name so import_string can find it in the subprocess.
+                    if callback_path and 
callback_path.startswith("unusual_prefix_"):
+                        _register_unusual_prefix_module(callback_path, 
bundle.path, _log)
                 except Exception:
                     _log.warning(
                         "Failed to initialize DAG bundle for callback",
                         bundle_name=bundle_info.name,
                         exc_info=True,
                     )
 
-            success, error_msg = execute_callback(callback_path, 
callback_kwargs, _log)
+            # When DagRun identifiers are provided, fetch the DagRun via 
SUPERVISOR_COMMS
+            # and build a context dict to pass to the callback function.
+            effective_kwargs = dict(callback_kwargs)
+            if dag_id and run_id:
+                context = 
_fetch_and_build_context(task_runner.SUPERVISOR_COMMS, dag_id, run_id, _log)
+                if context is None:
+                    _log.error(
+                        "Cannot execute callback without context — failing so 
it can be retried",
+                        dag_id=dag_id,
+                        run_id=run_id,
+                    )
+                    sys.exit(1)

Review Comment:
   The comment here says "failing so it can be retried," but the executor path 
has no retry: a non-zero exit becomes a terminal `CallbackState.FAILED` in 
`scheduler_job_runner.py:1309-1311` (the state is set to FAILED with no 
requeue). So a transient context-fetch failure (API blip, network partition, 
token expiry) permanently loses the deadline notification, a regression from 
the pre-PR design where context lived on the DB row and didn't depend on a 
runtime fetch. The triggerer path effectively retries (it skips and 
re-evaluates on the next loop) while the executor path fails terminally on the 
same failure class. Fetching context before the callback is claimed, or mapping 
this exit to a retryable result rather than terminal FAILED, would close the 
gap.



##########
airflow-core/src/airflow/models/callback.py:
##########
@@ -228,20 +228,24 @@ def __repr__(self):
         return f"{self.data['path']}({self.data['kwargs'] or ''}) on a 
triggerer"
 
     def queue(self, *, session: Session) -> None:
+        from airflow.models.dag import DagModel
         from airflow.models.trigger import Trigger
         from airflow.triggers.callback import CallbackTrigger
 
         team_name: str | None = None
         if self.bundle_name and conf.getboolean("core", "multi_team"):
             team_name = DagBundleModel.get_team_name(self.bundle_name, 
session=session)
+        elif dag_id := self.data.get("dag_id"):
+            team_name = DagModel.get_team_name(dag_id)

Review Comment:
   The `if` branch above passes `session=session` to 
`DagBundleModel.get_team_name`, but this `elif` calls 
`DagModel.get_team_name(dag_id)` without it. `DagModel.get_team_name` is `def 
get_team_name(dag_id, *, session: Session = NEW_SESSION)`, so this opens a 
fresh session mid-`queue()` instead of joining the active one. Worth passing 
`session=session` for consistency.



##########
airflow-core/src/airflow/triggers/callback.py:
##########
@@ -17,21 +17,60 @@
 
 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()
+                # Walk the bundle tree to handle nested DAG files (e.g. 
dags/team_a/my_dag.py)
+                for file_path in Path(bundle.path).rglob(f"{stem}.py"):

Review Comment:
   `get_unique_dag_module_name` sanitizes `.` and `-` to `_` in the stem 
(`re.sub(r"[.-]", "_", Path(file_path).stem)` in `utils/file.py`), so a DAG 
file `alerts-prod.py` yields stem `alerts_prod` in the mangled name. `parts[3]` 
recovers that sanitized stem, and `rglob(f"{stem}.py")` = 
`rglob("alerts_prod.py")` never matches the real `alerts-prod.py` on disk, so 
the hash check is unreachable and a bundle-local callback in a dashed/dotted 
filename silently ends in FAILED. Same pattern in 
`_register_unusual_prefix_module` (`callback_supervisor.py:385`). Globbing 
`*.py` and comparing `get_unique_dag_module_name(str(f)) == mod_name` directly 
handles it.



-- 
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]

Reply via email to