From 67af7d7b644135b578bef61e9ec9bcdc31a74df9 Mon Sep 17 00:00:00 2001 From: Amogh Desai Date: Sat, 9 Aug 2025 10:24:37 +0530 Subject: [PATCH] Unify error handling when connection is not found in aws hook --- .../providers/amazon/aws/hooks/base_aws.py | 21 +++++-------------- 1 file changed, 5 insertions(+), 16 deletions(-) diff --git a/providers/amazon/src/airflow/providers/amazon/aws/hooks/base_aws.py b/providers/amazon/src/airflow/providers/amazon/aws/hooks/base_aws.py index 832f3d92e310a..256630855fb42 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/hooks/base_aws.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/hooks/base_aws.py @@ -61,7 +61,6 @@ from airflow.providers.amazon.aws.utils.identifiers import generate_uuid from airflow.providers.amazon.aws.utils.suppress import return_on_error from airflow.providers.amazon.version_compat import BaseHook -from airflow.providers.common.compat.version_compat import AIRFLOW_V_3_0_PLUS from airflow.providers_manager import ProvidersManager from airflow.utils.helpers import exactly_one from airflow.utils.log.logging_mixin import LoggingMixin @@ -77,9 +76,6 @@ BaseAwsConnection = TypeVar("BaseAwsConnection", bound=Union[BaseClient, ServiceResource]) # noqa: UP007 -if AIRFLOW_V_3_0_PLUS: - from airflow.sdk.exceptions import AirflowRuntimeError, ErrorType - if TYPE_CHECKING: from aiobotocore.session import AioSession from botocore.client import ClientMeta @@ -621,19 +617,12 @@ def conn_config(self) -> AwsConnectionWrapper: if self.aws_conn_id: try: connection = self.get_connection(self.aws_conn_id) - except Exception as e: - not_found_exc_via_core = isinstance(e, AirflowNotFoundException) - not_found_exc_via_task_sdk = ( - AIRFLOW_V_3_0_PLUS - and isinstance(e, AirflowRuntimeError) - and e.error.error == ErrorType.CONNECTION_NOT_FOUND + except AirflowNotFoundException: + self.log.warning( + "Unable to find AWS Connection ID '%s', switching to empty.", self.aws_conn_id ) - if not_found_exc_via_core or not_found_exc_via_task_sdk: - self.log.warning( - "Unable to find AWS Connection ID '%s', switching to empty.", self.aws_conn_id - ) - else: - raise + except Exception: + raise return AwsConnectionWrapper( conn=connection,