Retain AirflowNotFoundException

Author: ephraimbuddy

airflow-core/src/airflow/exceptions.py

@@ -27,17 +27,13 @@
 if TYPE_CHECKING:
     from airflow.models import DagRun
 
-from airflow.sdk.exceptions import AirflowException
+from airflow.sdk.exceptions import AirflowException, AirflowNotFoundException
 
 
 class TaskNotFound(AirflowException):
     """Raise when a Task is not available in the system."""
 
 
-class AirflowNotFoundException(AirflowException):
-    """Raise when the requested object/resource is not available in the system."""
-
-
 class AirflowBadRequest(AirflowException):
     """Raise when the application or server cannot handle the request."""
 

providers/amazon/src/airflow/providers/amazon/aws/hooks/athena_sql.py

@@ -119,7 +119,7 @@ def conn_config(self) -> AwsConnectionWrapper:
                     raise ValueError(
                         f"Encountered non-JSON in `extra` field for connection {self.aws_conn_id!r}."
                     )
-            except (AirflowNotFoundException, RuntimeError):
+            except AirflowNotFoundException:
                 connection = athena_conn
                 connection.conn_type = "aws"
                 self.log.warning(

providers/amazon/src/airflow/providers/amazon/aws/hooks/base_aws.py

@@ -620,7 +620,7 @@ def conn_config(self) -> AwsConnectionWrapper:
         if self.aws_conn_id:
             try:
                 connection = self.get_connection(self.aws_conn_id)
-            except (AirflowNotFoundException, RuntimeError):
+            except AirflowNotFoundException:
                 self.log.warning(
                     "Unable to find AWS Connection ID '%s', switching to empty.", self.aws_conn_id
                 )

providers/amazon/src/airflow/providers/amazon/aws/hooks/emr.py

@@ -112,7 +112,7 @@ def create_job_flow(self, job_flow_overrides: dict[str, Any]) -> dict[str, Any]:
         if self.emr_conn_id:
             try:
                 emr_conn = self.get_connection(self.emr_conn_id)
-            except (AirflowNotFoundException, RuntimeError):
+            except AirflowNotFoundException:
                 warnings.warn(
                     f"Unable to find {self.hook_name} Connection ID {self.emr_conn_id!r}, "
                     "using an empty initial configuration. If you want to get rid of this warning "

providers/apache/spark/src/airflow/providers/apache/spark/hooks/spark_sql.py

@@ -104,7 +104,7 @@ def __init__(
 
         try:
             conn = self.get_connection(conn_id)
-        except (AirflowNotFoundException, RuntimeError):
+        except AirflowNotFoundException:
             conn = None
         if conn:
             options = conn.extra_dejson

providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/hooks/kubernetes.py

@@ -208,7 +208,7 @@ def get_connection(cls, conn_id: str) -> Connection:
         """
         try:
             return super().get_connection(conn_id)  # type: ignore[return-value]
-        except (AirflowNotFoundException, RuntimeError):
+        except AirflowNotFoundException:
             if conn_id == cls.default_conn_name:
                 return Connection(conn_id=cls.default_conn_name)
             raise

providers/cncf/kubernetes/tests/unit/cncf/kubernetes/hooks/test_kubernetes.py

@@ -448,9 +448,7 @@ def test_missing_default_connection_is_ok(self, remove_default_conn, sdk_connect
 
         # meanwhile, asking for non-default should still fail if it doesn't exist
         hook = KubernetesHook("some_conn")
-        with pytest.raises(
-            (AirflowNotFoundException, RuntimeError), match="The conn_id `some_conn` isn't defined"
-        ):
+        with pytest.raises(AirflowNotFoundException, match="The conn_id `some_conn` isn't defined"):
             hook.conn_extras
 
     @patch("kubernetes.config.kube_config.KubeConfigLoader")

providers/edge3/src/airflow/providers/edge3/example_dags/integration_test.py

@@ -129,7 +129,7 @@ def connection():
         try:
             conn = BaseHook.get_connection("integration_test")
             print(f"Got connection {conn}")
-        except (AirflowNotFoundException, RuntimeError):
+        except AirflowNotFoundException:
             print("Connection not found... but also OK.")
 
     @task_group(prefix_group_id=False)

providers/edge3/src/airflow/providers/edge3/example_dags/win_test.py

@@ -303,7 +303,7 @@ def connection():
             try:
                 conn = BaseHook.get_connection("integration_test")
                 print(f"Got connection {conn}")
-            except (AirflowNotFoundException, RuntimeError):
+            except AirflowNotFoundException:
                 print("Connection 'integration_test' not found... but also OK.")
 
         command = CmdOperator(task_id="command", command="echo Hello World")

providers/google/src/airflow/providers/google/cloud/log/gcs_task_handler.py

@@ -93,7 +93,7 @@ def hook(self) -> GCSHook | None:
         if conn_id:
             try:
                 return GCSHook(gcp_conn_id=conn_id)
-            except (AirflowNotFoundException, RuntimeError):
+            except AirflowNotFoundException:
                 pass
         return None