Rename back RuntimeError to AirflowException

ephraimbuddy · ephraimbuddy · commit 91d4cc0d9401 · 2025-08-15T16:37:14.000+01:00
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -1657,6 +1657,11 @@ repos:
           ^airflow-core/src/airflow/api_fastapi/core_api/datamodels/hitl\.py$|
           ^airflow-core/src/airflow/api_fastapi/core_api/datamodels/variables\.py$|
           ^airflow-core/src/airflow/api_fastapi/core_api/routes/public/dag_run\.py$|
+          ^airflow-core/src/airflow/api_fastapi/core_api/routes/public/extra_links\.py$|
+          ^airflow-core/src/airflow/api_fastapi/core_api/routes/public/task_instances\.py$|
+          ^airflow-core/src/airflow/api_fastapi/core_api/routes/public/log\.py$|
+          ^airflow-core/src/airflow/api_fastapi/core_api/routes/public/tasks\.py$|
+          ^airflow-core/src/airflow/api_fastapi/core_api/routes/public/xcom\.py$|
           ^airflow-core/src/airflow/api_fastapi/core_api/routes/ui/grid\.py$|
           ^airflow-core/src/airflow/api_fastapi/core_api/routes/ui/structure\.py$|
           ^airflow-core/src/airflow/api_fastapi/core_api/services/ui/connections\.py$|
@@ -1701,6 +1706,7 @@ repos:
           ^airflow-core/src/airflow/serialization/dag\.py$|
           ^airflow-core/src/airflow/serialization/enums\.py$|
           ^airflow-core/src/airflow/serialization/serialized_objects\.py$|
+          ^airflow-core/src/airflow/ti_deps/dep_context\.py$|
           ^airflow-core/src/airflow/task/task_runner/bash_task_runner\.py$|
           ^airflow-core/src/airflow/task/task_runner/standard_task_runner\.py$|
           ^airflow-core/src/airflow/utils/dag_cycle_tester\.py$|
diff --git a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/extra_links.py b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/extra_links.py
@@ -28,8 +28,8 @@
 from airflow.api_fastapi.core_api.datamodels.extra_links import ExtraLinkCollectionResponse
 from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc
 from airflow.api_fastapi.core_api.security import DagAccessEntity, requires_access_dag
-from airflow.exceptions import TaskNotFound
 from airflow.models import DagRun
+from airflow.sdk.exceptions import TaskNotFound
 
 if TYPE_CHECKING:
     from airflow.models.mappedoperator import MappedOperator
diff --git a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/log.py b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/log.py
@@ -35,9 +35,9 @@
 from airflow.api_fastapi.core_api.datamodels.log import ExternalLogUrlResponse, TaskInstancesLogResponse
 from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc
 from airflow.api_fastapi.core_api.security import DagAccessEntity, requires_access_dag
-from airflow.exceptions import TaskNotFound
 from airflow.models import TaskInstance, Trigger
 from airflow.models.taskinstancehistory import TaskInstanceHistory
+from airflow.sdk.exceptions import TaskNotFound
 from airflow.utils.log.log_reader import TaskLogReader
 
 task_instances_log_router = AirflowRouter(
diff --git a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/task_instances.py b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/task_instances.py
@@ -75,10 +75,10 @@
     _patch_ti_validate_request,
 )
 from airflow.api_fastapi.logging.decorators import action_logging
-from airflow.exceptions import TaskNotFound
 from airflow.models import Base, DagRun
 from airflow.models.taskinstance import TaskInstance as TI, clear_task_instances
 from airflow.models.taskinstancehistory import TaskInstanceHistory as TIH
+from airflow.sdk.exceptions import TaskNotFound
 from airflow.ti_deps.dep_context import DepContext
 from airflow.ti_deps.dependencies_deps import SCHEDULER_QUEUED_DEPS
 from airflow.utils.db import get_query_count
diff --git a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/tasks.py b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/tasks.py
@@ -29,7 +29,7 @@
 from airflow.api_fastapi.core_api.datamodels.tasks import TaskCollectionResponse, TaskResponse
 from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc
 from airflow.api_fastapi.core_api.security import requires_access_dag
-from airflow.exceptions import TaskNotFound
+from airflow.sdk.exceptions import TaskNotFound
 
 tasks_router = AirflowRouter(tags=["Task"], prefix="/dags/{dag_id}/tasks")
 
diff --git a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/xcom.py b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/xcom.py
@@ -49,10 +49,10 @@
 from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc
 from airflow.api_fastapi.core_api.security import ReadableXComFilterDep, requires_access_dag
 from airflow.api_fastapi.logging.decorators import action_logging
-from airflow.exceptions import TaskNotFound
 from airflow.models import DagRun as DR
 from airflow.models.dag import DagModel
 from airflow.models.xcom import XComModel
+from airflow.sdk.exceptions import TaskNotFound
 
 xcom_router = AirflowRouter(
     tags=["XCom"], prefix="/dags/{dag_id}/dagRuns/{dag_run_id}/taskInstances/{task_id}/xcomEntries"
diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py
@@ -57,7 +57,6 @@
     TITerminalStatePayload,
 )
 from airflow.api_fastapi.execution_api.deps import JWTBearerTIPathDep
-from airflow.exceptions import TaskNotFound
 from airflow.models.asset import AssetActive
 from airflow.models.dagrun import DagRun as DR
 from airflow.models.taskinstance import TaskInstance as TI, _stop_remaining_tasks
@@ -66,6 +65,7 @@
 from airflow.models.xcom import XComModel
 from airflow.sdk.definitions._internal.expandinput import NotFullyPopulated
 from airflow.sdk.definitions.asset import Asset, AssetUniqueKey
+from airflow.sdk.exceptions import TaskNotFound
 from airflow.utils.state import DagRunState, TaskInstanceState, TerminalTIState
 
 if TYPE_CHECKING:
diff --git a/airflow-core/src/airflow/exceptions.py b/airflow-core/src/airflow/exceptions.py
@@ -36,10 +36,6 @@ class AirflowBadRequest(AirflowException):
     status_code = HTTPStatus.BAD_REQUEST
 
 
-class TaskNotFound(AirflowException):
-    """Raise when a Task is not available in the system."""
-
-
 class AirflowConfigException(AirflowException):
     """Raise when there is configuration problem."""
 
diff --git a/airflow-core/src/airflow/ti_deps/dep_context.py b/airflow-core/src/airflow/ti_deps/dep_context.py
@@ -22,7 +22,7 @@
 
 import attr
 
-from airflow.exceptions import TaskNotFound
+from airflow.sdk.exceptions import TaskNotFound
 from airflow.utils.state import State
 
 if TYPE_CHECKING:
diff --git a/task-sdk/src/airflow/sdk/api/client.py b/task-sdk/src/airflow/sdk/api/client.py
@@ -66,7 +66,7 @@
     XComSequenceIndexResponse,
     XComSequenceSliceResponse,
 )
-from airflow.sdk.exceptions import ErrorType
+from airflow.sdk.exceptions import AirflowException, ErrorType
 from airflow.sdk.execution_time.comms import (
     CreateHITLDetailPayload,
     DRCount,
@@ -422,7 +422,7 @@ def head(self, dag_id: str, run_id: str, task_id: str, key: str) -> int:
         if not (content_range := resp.headers["Content-Range"]) or not content_range.startswith(
             "map_indexes "
         ):
-            raise RuntimeError(f"Unable to parse Content-Range header from HEAD {resp.request.url}")
+            raise AirflowException(f"Unable to parse Content-Range header from HEAD {resp.request.url}")
         return int(content_range[len("map_indexes ") :])
 
     def get(
diff --git a/task-sdk/src/airflow/sdk/bases/operator.py b/task-sdk/src/airflow/sdk/bases/operator.py
@@ -60,7 +60,7 @@
 from airflow.sdk.definitions.edges import EdgeModifier
 from airflow.sdk.definitions.mappedoperator import OperatorPartial, validate_mapping_kwargs
 from airflow.sdk.definitions.param import ParamsDict
-from airflow.sdk.exceptions import RemovedInAirflow4Warning
+from airflow.sdk.exceptions import AirflowException, RemovedInAirflow4Warning
 from airflow.task.priority_strategy import (
     PriorityWeightStrategy,
     airflow_priority_weight_strategies,
@@ -161,7 +161,7 @@ def parse_retries(retries: Any) -> int | None:
     try:
         parsed_retries = int(retries)
     except (TypeError, ValueError):
-        raise RuntimeError(f"'retries' type must be int, not {type(retries).__name__}")
+        raise AirflowException(f"'retries' type must be int, not {type(retries).__name__}")
     return parsed_retries
 
 
@@ -390,7 +390,7 @@ def wrapper(self, *args, **kwargs):
                 if not cls.test_mode and sentinel is not self:
                     message = f"{self.__class__.__name__}.{func.__name__} cannot be called outside of the Task Runner!"
                     if not self.allow_nested_operators:
-                        raise RuntimeError(message)
+                        raise AirflowException(message)
                     self.log.warning(message)
 
                     # Now that we've logged, set sentinel so that `super()` calls don't log again
@@ -1276,7 +1276,7 @@ def dag(self) -> DAG:
         """Returns the Operator's DAG if set, otherwise raises an error."""
         if dag := self._dag:
             return dag
-        raise RuntimeError(f"Operator {self} has not been assigned to a DAG yet")
+        raise AirflowException(f"Operator {self} has not been assigned to a DAG yet")
 
     @dag.setter
     def dag(self, dag: DAG | None) -> None:
@@ -1446,7 +1446,7 @@ def get_serialized_fields(cls):
 
             # make sure the following "fake" task is not added to current active
             # dag in context, otherwise, it will result in
-            # `RuntimeError: dictionary changed size during iteration`
+            # `AirflowException: dictionary changed size during iteration`
             # Exception in SerializedDAG.serialize_dag() call.
             DagContext.push(None)
             cls.__serialized_fields = frozenset(
diff --git a/task-sdk/src/airflow/sdk/bases/sensor.py b/task-sdk/src/airflow/sdk/bases/sensor.py
@@ -29,6 +29,7 @@
 from airflow.sdk import timezone
 from airflow.sdk.bases.operator import BaseOperator
 from airflow.sdk.exceptions import (
+    AirflowException,
     AirflowFailException,
     AirflowRescheduleException,
     AirflowSensorTimeout,
@@ -144,31 +145,33 @@ def _coerce_poke_interval(poke_interval: float | timedelta) -> timedelta:
             return poke_interval
         if isinstance(poke_interval, (int, float)) and poke_interval >= 0:
             return timedelta(seconds=poke_interval)
-        raise RuntimeError("Operator arg `poke_interval` must be timedelta object or a non-negative number")
+        raise AirflowException(
+            "Operator arg `poke_interval` must be timedelta object or a non-negative number"
+        )
 
     @staticmethod
     def _coerce_timeout(timeout: float | timedelta) -> timedelta:
         if isinstance(timeout, timedelta):
             return timeout
         if isinstance(timeout, (int, float)) and timeout >= 0:
             return timedelta(seconds=timeout)
-        raise RuntimeError("Operator arg `timeout` must be timedelta object or a non-negative number")
+        raise AirflowException("Operator arg `timeout` must be timedelta object or a non-negative number")
 
     @staticmethod
     def _coerce_max_wait(max_wait: float | timedelta | None) -> timedelta | None:
         if max_wait is None or isinstance(max_wait, timedelta):
             return max_wait
         if isinstance(max_wait, (int, float)) and max_wait >= 0:
             return timedelta(seconds=max_wait)
-        raise RuntimeError("Operator arg `max_wait` must be timedelta object or a non-negative number")
+        raise AirflowException("Operator arg `max_wait` must be timedelta object or a non-negative number")
 
     def _validate_input_values(self) -> None:
         if not isinstance(self.poke_interval, (int, float)) or self.poke_interval < 0:
-            raise RuntimeError("The poke_interval must be a non-negative number")
+            raise AirflowException("The poke_interval must be a non-negative number")
         if not isinstance(self.timeout, (int, float)) or self.timeout < 0:
-            raise RuntimeError("The timeout must be a non-negative number")
+            raise AirflowException("The timeout must be a non-negative number")
         if self.mode not in self.valid_modes:
-            raise RuntimeError(
+            raise AirflowException(
                 f"The mode must be one of {self.valid_modes},'{self.dag.dag_id if self.has_dag() else ''} "
                 f".{self.task_id}'; received '{self.mode}'."
             )
diff --git a/task-sdk/src/airflow/sdk/definitions/connection.py b/task-sdk/src/airflow/sdk/definitions/connection.py
@@ -24,7 +24,7 @@
 
 import attrs
 
-from airflow.sdk.exceptions import AirflowNotFoundException, AirflowRuntimeError, ErrorType
+from airflow.sdk.exceptions import AirflowException, AirflowNotFoundException, AirflowRuntimeError, ErrorType
 
 log = logging.getLogger(__name__)
 
@@ -130,7 +130,7 @@ def get_hook(self, *, hook_params=None):
         hook = ProvidersManager().hooks.get(self.conn_type, None)
 
         if hook is None:
-            raise RuntimeError(f'Unknown hook type "{self.conn_type}"')
+            raise AirflowException(f'Unknown hook type "{self.conn_type}"')
         try:
             hook_class = import_string(hook.hook_class_name)
         except ImportError:
diff --git a/task-sdk/src/airflow/sdk/definitions/decorators/setup_teardown.py b/task-sdk/src/airflow/sdk/definitions/decorators/setup_teardown.py
@@ -23,6 +23,7 @@
 from airflow.sdk.bases.operator import BaseOperator
 from airflow.sdk.definitions._internal.setup_teardown import SetupTeardownContext
 from airflow.sdk.definitions.decorators.task_group import _TaskGroupFactory
+from airflow.sdk.exceptions import AirflowException
 
 if TYPE_CHECKING:
     from airflow.sdk.bases.decorator import _TaskDecorator
@@ -51,7 +52,7 @@ def initialize_context(...):
     if isinstance(func, types.FunctionType):
         func = python_task(func)
     if isinstance(func, _TaskGroupFactory):
-        raise RuntimeError("Task groups cannot be marked as setup or teardown.")
+        raise AirflowException("Task groups cannot be marked as setup or teardown.")
     func = cast("_TaskDecorator", func)
     func.is_setup = True
     return func
@@ -76,7 +77,7 @@ def teardown(func: Callable) -> Callable:
         if isinstance(func, types.FunctionType):
             func = python_task(func)
         if isinstance(func, _TaskGroupFactory):
-            raise RuntimeError("Task groups cannot be marked as setup or teardown.")
+            raise AirflowException("Task groups cannot be marked as setup or teardown.")
         func = cast("_TaskDecorator", func)
 
         func.is_teardown = True
@@ -101,9 +102,9 @@ def __enter__(self):
             if isinstance(task, BaseOperator):
                 operators.append(task)
                 if not task.is_setup and not task.is_teardown:
-                    raise RuntimeError("Only setup/teardown tasks can be used as context managers.")
+                    raise AirflowException("Only setup/teardown tasks can be used as context managers.")
             elif not task.operator.is_setup and not task.operator.is_teardown:
-                raise RuntimeError("Only setup/teardown tasks can be used as context managers.")
+                raise AirflowException("Only setup/teardown tasks can be used as context managers.")
         if not operators:
             # means we have XComArgs
             operators = [task.operator for task in self.tasks]
diff --git a/task-sdk/src/airflow/sdk/definitions/param.py b/task-sdk/src/airflow/sdk/definitions/param.py
@@ -25,7 +25,7 @@
 
 from airflow.sdk.definitions._internal.mixins import ResolveMixin
 from airflow.sdk.definitions._internal.types import NOTSET, ArgNotSet
-from airflow.sdk.exceptions import ParamValidationError
+from airflow.sdk.exceptions import AirflowException, ParamValidationError
 
 if TYPE_CHECKING:
     from airflow.sdk.definitions.context import Context
@@ -295,7 +295,7 @@ def resolve(self, context: Context) -> Any:
             return self._default
         with contextlib.suppress(KeyError):
             return context["params"][self._name]
-        raise RuntimeError(f"No value could be resolved for parameter {self._name}")
+        raise AirflowException(f"No value could be resolved for parameter {self._name}")
 
     def serialize(self) -> dict:
         """Serialize the DagParam object into a dictionary."""
diff --git a/task-sdk/src/airflow/sdk/definitions/taskgroup.py b/task-sdk/src/airflow/sdk/definitions/taskgroup.py
@@ -36,6 +36,7 @@
 from airflow.sdk.definitions._internal.node import DAGNode, validate_group_key
 from airflow.sdk.exceptions import (
     AirflowDagCycleException,
+    AirflowException,
     DuplicateTaskIdFound,
     TaskAlreadyInTaskGroup,
 )
@@ -259,7 +260,7 @@ def add(self, task: DAGNode) -> DAGNode:
                     )
                 task.dag = self.dag
             if task.children:
-                raise RuntimeError("Cannot add a non-empty TaskGroup")
+                raise AirflowException("Cannot add a non-empty TaskGroup")
 
         self.children[key] = task
         return task
@@ -316,7 +317,7 @@ def update_relative(
             # Handles setting relationship between a TaskGroup and a task
             for task in other.roots:
                 if not isinstance(task, DAGNode):
-                    raise RuntimeError(
+                    raise AirflowException(
                         "Relationships can only be set between TaskGroup "
                         f"or operators; received {task.__class__.__name__}"
                     )
diff --git a/task-sdk/src/airflow/sdk/definitions/xcom_arg.py b/task-sdk/src/airflow/sdk/definitions/xcom_arg.py
@@ -30,7 +30,7 @@
 from airflow.sdk.definitions._internal.mixins import DependencyMixin, ResolveMixin
 from airflow.sdk.definitions._internal.setup_teardown import SetupTeardownContext
 from airflow.sdk.definitions._internal.types import NOTSET, ArgNotSet
-from airflow.sdk.exceptions import XComNotFound
+from airflow.sdk.exceptions import AirflowException, XComNotFound
 from airflow.sdk.execution_time.lazy_sequence import LazyXComSequence
 from airflow.sdk.execution_time.xcom import BaseXCom
 from airflow.utils.trigger_rule import TriggerRule
@@ -182,7 +182,7 @@ def resolve(self, context: Mapping[str, Any]) -> Any:
 
     def __enter__(self):
         if not self.operator.is_setup and not self.operator.is_teardown:
-            raise RuntimeError("Only setup/teardown tasks can be used as context managers.")
+            raise AirflowException("Only setup/teardown tasks can be used as context managers.")
         SetupTeardownContext.push_setup_teardown_task(self.operator)
         return SetupTeardownContext
 
diff --git a/task-sdk/src/airflow/sdk/execution_time/task_runner.py b/task-sdk/src/airflow/sdk/execution_time/task_runner.py
@@ -57,6 +57,7 @@
 from airflow.sdk.definitions.mappedoperator import MappedOperator
 from airflow.sdk.definitions.param import process_params
 from airflow.sdk.exceptions import (
+    AirflowException,
     AirflowInactiveAssetInInletOrOutletException,
     AirflowRuntimeError,
     AirflowTaskTimeout,
@@ -967,7 +968,7 @@ def _on_term(signum, frame):
         )
         state = TaskInstanceState.FAILED
         error = e
-    except (AirflowTaskTimeout, RuntimeError, AirflowRuntimeError) as e:
+    except (AirflowTaskTimeout, AirflowException, AirflowRuntimeError) as e:
         # We should allow retries if the task has defined it.
         log.exception("Task failed with exception")
         msg, state = _handle_current_task_failed(ti)
diff --git a/task-sdk/tests/task_sdk/bases/test_operator.py b/task-sdk/tests/task_sdk/bases/test_operator.py
@@ -765,7 +765,7 @@ class Branch(Mixin, sql.BranchSQLOperator):
             pass
 
         # The following throws an exception if metaclass breaks MRO:
-        #   airflow.exceptions.RuntimeError: Invalid arguments were passed to Branch (task_id: test). Invalid arguments were:
+        #   airflow.sdk.exceptions.AirflowException: Invalid arguments were passed to Branch (task_id: test). Invalid arguments were:
         #   **kwargs: {'sql': 'sql', 'follow_task_ids_if_true': ['x'], 'follow_task_ids_if_false': ['y']}
         op = Branch(
             task_id="test",
diff --git a/task-sdk/tests/task_sdk/bases/test_sensor.py b/task-sdk/tests/task_sdk/bases/test_sensor.py
diff --git a/task-sdk/tests/task_sdk/definitions/decorators/test_setup_teardown.py b/task-sdk/tests/task_sdk/definitions/decorators/test_setup_teardown.py
diff --git a/task-sdk/tests/task_sdk/definitions/test_connections.py b/task-sdk/tests/task_sdk/definitions/test_connections.py
diff --git a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py
