Improve DAG and task missing error handling in callbacks (#56725)

When callbacks are executed but the DAG or task is missing from the
`DagBag`, users now receive clear error messages explaining what went
wrong instead of cryptic `KeyError` exceptions. This typically occurs
due to race conditions where DAGs are removed or fail to parse between
scheduling and callback execution.

Fixes #56701

Author: kaxil

airflow-core/src/airflow/dag_processing/processor.py

@@ -36,6 +36,7 @@
 )
 from airflow.configuration import conf
 from airflow.dag_processing.dagbag import DagBag
+from airflow.exceptions import TaskNotFound
 from airflow.sdk.execution_time.comms import (
     ConnectionResult,
     DeleteVariable,
@@ -71,7 +72,10 @@
 
     from airflow.api_fastapi.execution_api.app import InProcessExecutionAPI
     from airflow.sdk.api.client import Client
+    from airflow.sdk.bases.operator import BaseOperator
     from airflow.sdk.definitions.context import Context
+    from airflow.sdk.definitions.dag import DAG
+    from airflow.sdk.definitions.mappedoperator import MappedOperator
     from airflow.typing_compat import Self
 
 
@@ -234,6 +238,39 @@ def _serialize_dags(
     return serialized_dags, serialization_import_errors
 
 
+def _get_dag_with_task(
+    dagbag: DagBag, dag_id: str, task_id: str | None = None
+) -> tuple[DAG, BaseOperator | MappedOperator | None]:
+    """
+    Retrieve a DAG and optionally a task from the DagBag.
+
+    :param dagbag: DagBag to retrieve from
+    :param dag_id: DAG ID to retrieve
+    :param task_id: Optional task ID to retrieve from the DAG
+    :return: tuple of (dag, task) where task is None if not requested
+    :raises ValueError: If DAG or task is not found
+    """
+    if dag_id not in dagbag.dags:
+        raise ValueError(
+            f"DAG '{dag_id}' not found in DagBag. "
+            f"This typically indicates a race condition where the DAG was removed or failed to parse."
+        )
+
+    dag = dagbag.dags[dag_id]
+
+    if task_id is not None:
+        try:
+            task = dag.get_task(task_id)
+            return dag, task
+        except TaskNotFound:
+            raise ValueError(
+                f"Task '{task_id}' not found in DAG '{dag_id}'. "
+                f"This typically indicates a race condition where the task was removed or the DAG structure changed."
+            ) from None
+
+    return dag, None
+
+
 def _execute_callbacks(
     dagbag: DagBag, callback_requests: list[CallbackRequest], log: FilteringBoundLogger
 ) -> None:
@@ -250,8 +287,7 @@ def _execute_callbacks(
 def _execute_dag_callbacks(dagbag: DagBag, request: DagCallbackRequest, log: FilteringBoundLogger) -> None:
     from airflow.sdk.api.datamodels._generated import TIRunContext
 
-    dag = dagbag.dags[request.dag_id]
-
+    dag, _ = _get_dag_with_task(dagbag, request.dag_id)
     callbacks = dag.on_failure_callback if request.is_failure_callback else dag.on_success_callback
     if not callbacks:
         log.warning("Callback requested, but dag didn't have any", dag_id=request.dag_id)
@@ -303,8 +339,10 @@ def _execute_task_callbacks(dagbag: DagBag, request: TaskCallbackRequest, log: F
         )
         return
 
-    dag = dagbag.dags[request.ti.dag_id]
-    task = dag.get_task(request.ti.task_id)
+    dag, task = _get_dag_with_task(dagbag, request.ti.dag_id, request.ti.task_id)
+
+    if TYPE_CHECKING:
+        assert task is not None
 
     if request.task_callback_type is TaskInstanceState.UP_FOR_RETRY:
         callbacks = task.on_retry_callback
@@ -356,8 +394,10 @@ def get_callback_representation(callback):
 
 def _execute_email_callbacks(dagbag: DagBag, request: EmailRequest, log: FilteringBoundLogger) -> None:
     """Execute email notification for task failure/retry."""
-    dag = dagbag.dags[request.ti.dag_id]
-    task = dag.get_task(request.ti.task_id)
+    dag, task = _get_dag_with_task(dagbag, request.ti.dag_id, request.ti.task_id)
+
+    if TYPE_CHECKING:
+        assert task is not None
 
     if not task.email:
         log.warning(

airflow-core/tests/unit/dag_processing/test_processor.py

@@ -901,6 +901,25 @@ def fake_collect_dags(self, *args, **kwargs):
         # Should log warning about no callback found
         log.warning.assert_called_once_with("Callback requested, but dag didn't have any", dag_id="test_dag")
 
+    def test_execute_dag_callbacks_missing_dag(self):
+        """Test _execute_dag_callbacks raises ValueError for missing DAG"""
+        dagbag = DagBag()
+
+        request = DagCallbackRequest(
+            filepath="test.py",
+            dag_id="missing_dag",
+            run_id="test_run",
+            bundle_name="testing",
+            bundle_version=None,
+            is_failure_callback=True,
+            msg="Test failure message",
+        )
+
+        log = structlog.get_logger()
+
+        with pytest.raises(ValueError, match="DAG 'missing_dag' not found in DagBag"):
+            _execute_dag_callbacks(dagbag, request, log)
+
     @pytest.mark.parametrize(
         "xcom_operation,expected_message_type,expected_message,mock_response",
         [
@@ -1277,6 +1296,58 @@ def fake_collect_dags(self, *args, **kwargs):
 
         assert call_count == 2
 
+    @pytest.mark.parametrize(
+        "dag_exists,task_exists,expected_error",
+        [
+            (False, False, "DAG 'missing_dag' not found in DagBag"),
+            (True, False, "Task 'missing_task' not found in DAG 'test_dag'"),
+        ],
+    )
+    def test_execute_task_callbacks_missing_dag_or_task(
+        self, spy_agency, dag_exists, task_exists, expected_error
+    ):
+        """Test _execute_task_callbacks raises ValueError for missing DAG or task"""
+        if dag_exists:
+            with DAG(dag_id="test_dag") as dag:
+                BaseOperator(task_id="existing_task")
+
+            def fake_collect_dags(self, *args, **kwargs):
+                self.dags[dag.dag_id] = dag
+
+            spy_agency.spy_on(DagBag.collect_dags, call_fake=fake_collect_dags, owner=DagBag)
+
+            dagbag = DagBag()
+            dagbag.collect_dags()
+            dag_id = "test_dag"
+            task_id = "missing_task"
+        else:
+            dagbag = DagBag()
+            dag_id = "missing_dag"
+            task_id = "test_task"
+
+        ti_data = TIDataModel(
+            id=uuid.uuid4(),
+            dag_id=dag_id,
+            task_id=task_id,
+            run_id="test_run",
+            try_number=1,
+            dag_version_id=uuid.uuid4(),
+        )
+
+        request = TaskCallbackRequest(
+            filepath="test.py",
+            msg="Task failed",
+            ti=ti_data,
+            bundle_name="testing",
+            bundle_version=None,
+            task_callback_type=TaskInstanceState.FAILED,
+        )
+
+        log = structlog.get_logger()
+
+        with pytest.raises(ValueError, match=expected_error):
+            _execute_task_callbacks(dagbag, request, log)
+
 
 class TestExecuteEmailCallbacks:
     """Test the email callback execution functionality."""
@@ -1517,6 +1588,75 @@ def test_execute_email_callbacks_email_disabled_for_type(self, mock_send_email):
         info_call = log.info.call_args[0][0]
         assert "Email not sent - task configured with email_on_" in info_call
 
+    @pytest.mark.parametrize(
+        "dag_exists,task_exists,expected_error",
+        [
+            (False, False, "DAG 'missing_dag' not found in DagBag"),
+            (True, False, "Task 'missing_task' not found in DAG 'test_dag'"),
+        ],
+    )
+    def test_execute_email_callbacks_missing_dag_or_task(
+        self, spy_agency, dag_exists, task_exists, expected_error
+    ):
+        """Test _execute_email_callbacks raises ValueError for missing DAG or task"""
+        if dag_exists:
+            with DAG(dag_id="test_dag") as dag:
+                BaseOperator(task_id="existing_task", email="[email protected]")
+
+            def fake_collect_dags(self, *args, **kwargs):
+                self.dags[dag.dag_id] = dag
+
+            spy_agency.spy_on(DagBag.collect_dags, call_fake=fake_collect_dags, owner=DagBag)
+
+            dagbag = DagBag()
+            dagbag.collect_dags()
+            dag_id = "test_dag"
+            task_id = "missing_task"
+        else:
+            dagbag = DagBag()
+            dag_id = "missing_dag"
+            task_id = "test_task"
+
+        ti_data = TIDataModel(
+            id=uuid.uuid4(),
+            dag_id=dag_id,
+            task_id=task_id,
+            run_id="test_run",
+            try_number=1,
+            dag_version_id=uuid.uuid4(),
+        )
+
+        current_time = timezone.utcnow()
+        request = EmailRequest(
+            filepath="test.py",
+            bundle_name="testing",
+            bundle_version=None,
+            ti=ti_data,
+            context_from_server=TIRunContext(
+                dag_run=DRDataModel(
+                    dag_id=dag_id,
+                    run_id="test_run",
+                    logical_date=current_time,
+                    data_interval_start=current_time,
+                    data_interval_end=current_time,
+                    run_after=current_time,
+                    start_date=current_time,
+                    end_date=None,
+                    run_type="manual",
+                    state="running",
+                    consumed_asset_events=[],
+                ),
+                max_tries=2,
+            ),
+            email_type="failure",
+            msg="Task failed",
+        )
+
+        log = structlog.get_logger()
+
+        with pytest.raises(ValueError, match=expected_error):
+            _execute_email_callbacks(dagbag, request, log)
+
 
 class TestDagProcessingMessageTypes:
     def test_message_types_in_dag_processor(self):