fix: properly handle nested errors

karenc-bq · karenc-bq · commit c8e1c8d6f3f3 · 2026-02-08T15:11:27.000-08:00
diff --git a/.github/workflows/integration_tests.yml b/.github/workflows/integration_tests.yml
@@ -5,6 +5,7 @@ on:
   push:
     branches:
       - main
+      - fix/properly-handle-network-exceptions
 
 permissions:
   id-token: write   # This is required for requesting the JWT
@@ -18,7 +19,7 @@ jobs:
       fail-fast: false
       matrix:
         python-version: [ "3.11", "3.12", "3.13"  ]
-        environment: [ "mysql", "pg" ]
+        environment: [ "pg" ]
 
     steps:
       - name: 'Clone repository'
@@ -55,6 +56,7 @@ jobs:
           RDS_DB_REGION: ${{ secrets.AWS_DEFAULT_REGION }}
           AURORA_MYSQL_DB_ENGINE_VERSION: lts
           AURORA_PG_ENGINE_VERSION: lts
+          FILTER: test_writer_failover_in_idle_connections
 
       - name: 'Archive LTS results'
         if: always()
@@ -71,7 +73,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: [ "3.11", "3.12", "3.13"  ]
+        python-version: [ "3.11", "3.12", "3.13" ]
         environment: ["mysql", "pg"]
 
     steps:
@@ -109,6 +111,7 @@ jobs:
           RDS_DB_REGION: ${{ secrets.AWS_DEFAULT_REGION }}
           AURORA_MYSQL_DB_ENGINE_VERSION: latest
           AURORA_PG_ENGINE_VERSION: latest
+          FILTER: TestAuroraFailover
 
       - name: 'Archive Latest results'
         if: always()
diff --git a/aws_advanced_python_wrapper/aws_secrets_manager_plugin.py b/aws_advanced_python_wrapper/aws_secrets_manager_plugin.py
@@ -32,7 +32,7 @@
     from aws_advanced_python_wrapper.pep249 import Connection
     from aws_advanced_python_wrapper.plugin_service import PluginService
 
-from aws_advanced_python_wrapper.errors import AwsWrapperError
+from aws_advanced_python_wrapper.errors import AwsConnectError, AwsWrapperError
 from aws_advanced_python_wrapper.pep249_methods import DbApiMethod
 from aws_advanced_python_wrapper.plugin import Plugin, PluginFactory
 from aws_advanced_python_wrapper.utils.log import Logger
@@ -113,9 +113,12 @@ def _connect(self, host_info: HostInfo, props: Properties, connect_func: Callabl
             return connect_func()
 
         except Exception as e:
+            if self._plugin_service.is_network_exception(error=e):
+                raise AwsConnectError(Messages.get_formatted("AwsSecretsManagerPlugin.ConnectException", e)) from e
+
             if not self._plugin_service.is_login_exception(error=e) or secret_fetched:
                 raise AwsWrapperError(
-                    Messages.get_formatted("AwsSecretsManagerPlugin.ConnectException", e)) from e
+                    Messages.get_formatted("AwsSecretsManagerPlugin.ConnectException", e), e) from e
 
             secret_fetched = self._update_secret(host_info, props, token_expiration_ns=token_expiration_ns, force_refetch=True)
 
@@ -126,8 +129,8 @@ def _connect(self, host_info: HostInfo, props: Properties, connect_func: Callabl
                 except Exception as unhandled_error:
                     raise AwsWrapperError(
                         Messages.get_formatted("AwsSecretsManagerPlugin.UnhandledException",
-                                               unhandled_error)) from unhandled_error
-            raise AwsWrapperError(Messages.get_formatted("AwsSecretsManagerPlugin.FailedLogin", e)) from e
+                                               unhandled_error), unhandled_error) from unhandled_error
+            raise AwsWrapperError(Messages.get_formatted("AwsSecretsManagerPlugin.FailedLogin", e), e) from e
 
     def _update_secret(self, host_info: HostInfo, props: Properties, token_expiration_ns: int, force_refetch: bool = False) -> bool:
         """
@@ -154,19 +157,19 @@ def _update_secret(self, host_info: HostInfo, props: Properties, token_expiratio
                 except (ClientError, AttributeError) as e:
                     logger.debug("AwsSecretsManagerPlugin.FailedToFetchDbCredentials", e)
                     raise AwsWrapperError(
-                        Messages.get_formatted("AwsSecretsManagerPlugin.FailedToFetchDbCredentials", e)) from e
+                        Messages.get_formatted("AwsSecretsManagerPlugin.FailedToFetchDbCredentials", e), e) from e
                 except JSONDecodeError as e:
                     logger.debug("AwsSecretsManagerPlugin.JsonDecodeError", e)
                     raise AwsWrapperError(
-                        Messages.get_formatted("AwsSecretsManagerPlugin.JsonDecodeError", e))
-                except EndpointConnectionError:
+                        Messages.get_formatted("AwsSecretsManagerPlugin.JsonDecodeError", e), e) from e
+                except EndpointConnectionError as e:
                     logger.debug("AwsSecretsManagerPlugin.EndpointOverrideInvalidConnection", endpoint)
                     raise AwsWrapperError(
-                        Messages.get_formatted("AwsSecretsManagerPlugin.EndpointOverrideInvalidConnection", endpoint))
-                except ValueError:
+                        Messages.get_formatted("AwsSecretsManagerPlugin.EndpointOverrideInvalidConnection", endpoint), e) from e
+                except ValueError as e:
                     logger.debug("AwsSecretsManagerPlugin.EndpointOverrideMisconfigured", endpoint)
                     raise AwsWrapperError(
-                        Messages.get_formatted("AwsSecretsManagerPlugin.EndpointOverrideMisconfigured", endpoint))
+                        Messages.get_formatted("AwsSecretsManagerPlugin.EndpointOverrideMisconfigured", endpoint), e) from e
 
             return fetched
         except Exception as ex:
diff --git a/aws_advanced_python_wrapper/custom_endpoint_plugin.py b/aws_advanced_python_wrapper/custom_endpoint_plugin.py
@@ -330,8 +330,8 @@ def _wait_for_info(self, monitor: CustomEndpointMonitor):
             while not has_info and perf_counter_ns() < wait_for_info_timeout_ns:
                 sleep(0.1)
                 has_info = monitor.has_custom_endpoint_info()
-        except InterruptedError:
-            raise AwsWrapperError(Messages.get_formatted("CustomEndpointPlugin.InterruptedThread", hostname))
+        except InterruptedError as e:
+            raise AwsWrapperError(Messages.get_formatted("CustomEndpointPlugin.InterruptedThread", hostname), e) from e
 
         if not has_info:
             raise AwsWrapperError(
diff --git a/aws_advanced_python_wrapper/database_dialect.py b/aws_advanced_python_wrapper/database_dialect.py
@@ -75,8 +75,8 @@ class DialectCode(Enum):
     def from_string(value: str) -> DialectCode:
         try:
             return DialectCode(value)
-        except ValueError:
-            raise AwsWrapperError(Messages.get_formatted("DialectCode.InvalidStringValue", value))
+        except ValueError as e:
+            raise AwsWrapperError(Messages.get_formatted("DialectCode.InvalidStringValue", value), e) from e
 
 
 class TargetDriverType(Enum):
diff --git a/aws_advanced_python_wrapper/errors.py b/aws_advanced_python_wrapper/errors.py
@@ -11,12 +11,22 @@
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
+from typing import Optional
 
 from .pep249 import Error
 
 
 class AwsWrapperError(Error):
     __module__ = "aws_advanced_python_wrapper"
+    driver_error: Optional[Exception]
+
+    def __init__(self, message: str = "", original_error: Optional[Exception] = None):
+        super().__init__(message)
+        # If wrapping another AwsWrapperError, preserve the original driver exception
+        if isinstance(original_error, AwsWrapperError) and original_error.driver_error is not None:
+            self.driver_error = original_error.driver_error
+        else:
+            self.driver_error = original_error
 
 
 class UnsupportedOperationError(AwsWrapperError):
@@ -45,3 +55,7 @@ class FailoverSuccessError(FailoverError):
 
 class ReadWriteSplittingError(AwsWrapperError):
     __module__ = "aws_advanced_python_wrapper"
+
+
+class AwsConnectError(AwsWrapperError):
+    __module__ = "aws_advanced_python_wrapper"
diff --git a/aws_advanced_python_wrapper/failover_plugin.py b/aws_advanced_python_wrapper/failover_plugin.py
@@ -175,7 +175,7 @@ def execute(self, target: type, method_name: str, execute_func: Callable, *args:
 
                 self._pick_new_connection()
                 self._last_exception = ex
-            raise AwsWrapperError(Messages.get_formatted("FailoverPlugin.DetectedException", str(ex))) from ex
+            raise AwsWrapperError(Messages.get_formatted("FailoverPlugin.DetectedException", str(ex)), ex) from ex
 
     def notify_host_list_changed(self, changes: Dict[str, Set[HostEvent]]):
         if not self._enable_failover_setting:
diff --git a/aws_advanced_python_wrapper/failover_v2_plugin.py b/aws_advanced_python_wrapper/failover_v2_plugin.py
@@ -204,7 +204,7 @@ def _deal_with_original_exception(self, original_exception: Exception) -> None:
             self._pick_new_connection()
             self._last_exception_dealt_with = original_exception
 
-        raise AwsWrapperError(Messages.get_formatted("FailoverPlugin.DetectedException", str(original_exception))) \
+        raise AwsWrapperError(Messages.get_formatted("FailoverPlugin.DetectedException", str(original_exception)), original_exception) \
             from original_exception
 
     def _failover(self) -> None:
diff --git a/aws_advanced_python_wrapper/federated_plugin.py b/aws_advanced_python_wrapper/federated_plugin.py
@@ -39,7 +39,7 @@
 
 import requests
 
-from aws_advanced_python_wrapper.errors import AwsWrapperError
+from aws_advanced_python_wrapper.errors import AwsConnectError, AwsWrapperError
 from aws_advanced_python_wrapper.plugin import Plugin, PluginFactory
 from aws_advanced_python_wrapper.utils.log import Logger
 from aws_advanced_python_wrapper.utils.messages import Messages
@@ -114,8 +114,11 @@ def _connect(self, host_info: HostInfo, props: Properties, connect_func: Callabl
         try:
             return connect_func()
         except Exception as e:
+            if self._plugin_service.is_network_exception(e):
+                raise AwsConnectError(Messages.get_formatted("FederatedAuthPlugin.ConnectException", e)) from e
+
             if token_info is None or token_info.is_expired() or not self._plugin_service.is_login_exception(e):
-                raise e
+                raise AwsWrapperError(Messages.get_formatted("FederatedAuthPlugin.ConnectException", e), e) from e
 
             self._update_authentication_token(token_host_info, props, user, region, cache_key)
 
@@ -124,7 +127,7 @@ def _connect(self, host_info: HostInfo, props: Properties, connect_func: Callabl
             except Exception as e:
                 error_message = "FederatedAuthPlugin.UnhandledException"
                 logger.debug(error_message, e)
-                raise AwsWrapperError(Messages.get_formatted(error_message, e)) from e
+                raise AwsWrapperError(Messages.get_formatted(error_message, e), e) from e
 
     def force_connect(
             self,
@@ -211,7 +214,7 @@ def get_saml_assertion(self, props: Properties):
         except IOError as e:
             error_message = "FederatedAuthPlugin.UnhandledException"
             logger.debug(error_message, e)
-            raise AwsWrapperError(Messages.get_formatted(error_message, e))
+            raise AwsWrapperError(Messages.get_formatted(error_message, e), e) from e
 
     def _get_sign_in_page_body(self, url: str, props: Properties) -> str:
         logger.debug("AdfsCredentialsProviderFactory.SignOnPageUrl", url)
diff --git a/aws_advanced_python_wrapper/host_list_provider.py b/aws_advanced_python_wrapper/host_list_provider.py
@@ -641,7 +641,7 @@ def _query_for_topology(self, conn: Connection) -> Optional[Tuple[HostInfo, ...]
                 cursor.execute(self._dialect.topology_query)
                 return self._process_query_results(cursor)
         except ProgrammingError as e:
-            raise AwsWrapperError(Messages.get("RdsHostListProvider.InvalidQuery")) from e
+            raise AwsWrapperError(Messages.get("RdsHostListProvider.InvalidQuery"), e) from e
 
     def _process_query_results(self, cursor: Cursor) -> Tuple[HostInfo, ...]:
         """
@@ -707,7 +707,7 @@ def _query_for_topology(self, conn: Connection) -> Optional[Tuple[HostInfo, ...]
                 cursor.execute(self._dialect.topology_query)
                 return self._process_multi_az_query_results(cursor, writer_id)
         except ProgrammingError as e:
-            raise AwsWrapperError(Messages.get("RdsHostListProvider.InvalidQuery")) from e
+            raise AwsWrapperError(Messages.get("RdsHostListProvider.InvalidQuery"), e) from e
 
     def _process_multi_az_query_results(self, cursor: Cursor, writer_id: str) -> Tuple[HostInfo, ...]:
         hosts_dict = {}
diff --git a/aws_advanced_python_wrapper/host_selector.py b/aws_advanced_python_wrapper/host_selector.py
@@ -257,9 +257,9 @@ def _update_host_weight_map_from_string(self, props: Optional[Properties] = None
                             raise AwsWrapperError(Messages.get_formatted(message, pair))
 
                         self._host_weight_map[host_name] = weight
-                    except ValueError:
+                    except ValueError as e:
                         logger.error(message, pair)
-                        raise AwsWrapperError(Messages.get_formatted(message, pair))
+                        raise AwsWrapperError(Messages.get_formatted(message, pair), e) from e
 
 
 class HighestWeightHostSelector(HostSelector):
diff --git a/aws_advanced_python_wrapper/iam_plugin.py b/aws_advanced_python_wrapper/iam_plugin.py
@@ -31,7 +31,7 @@
 from datetime import datetime, timedelta
 from typing import Callable, Dict, Set
 
-from aws_advanced_python_wrapper.errors import AwsWrapperError
+from aws_advanced_python_wrapper.errors import AwsConnectError, AwsWrapperError
 from aws_advanced_python_wrapper.pep249_methods import DbApiMethod
 from aws_advanced_python_wrapper.plugin import Plugin, PluginFactory
 from aws_advanced_python_wrapper.utils.log import Logger
@@ -125,10 +125,12 @@ def _connect(self, host_info: HostInfo, props: Properties, connect_func: Callabl
         except Exception as e:
             logger.debug("IamAuthPlugin.ConnectException", e)
 
+            if self._plugin_service.is_network_exception(error=e):
+                raise AwsConnectError(Messages.get_formatted("IamAuthPlugin.ConnectException", e)) from e
+
             is_cached_token = (token_info is not None and not token_info.is_expired())
             if not self._plugin_service.is_login_exception(error=e) or not is_cached_token:
-                raise AwsWrapperError(Messages.get_formatted("IamAuthPlugin.ConnectException", e)) from e
-
+                raise AwsWrapperError(Messages.get_formatted("IamAuthPlugin.ConnectException", e), e) from e
             # Login unsuccessful with cached token
             # Try to generate a new token and try to connect again
             token_expiry = datetime.now() + timedelta(seconds=token_expiration_sec)
@@ -145,7 +147,7 @@ def _connect(self, host_info: HostInfo, props: Properties, connect_func: Callabl
             try:
                 return connect_func()
             except Exception as e:
-                raise AwsWrapperError(Messages.get_formatted("IamAuthPlugin.UnhandledException", e)) from e
+                raise AwsWrapperError(Messages.get_formatted("IamAuthPlugin.UnhandledException", e), e) from e
 
     def force_connect(
             self,
diff --git a/aws_advanced_python_wrapper/limitless_plugin.py b/aws_advanced_python_wrapper/limitless_plugin.py
@@ -420,8 +420,8 @@ def _retry_connection_with_least_loaded_routers(self, context: LimitlessContext)
                             if self._is_login_exception(e):
                                 raise e
 
-                            raise AwsWrapperError(Messages.get_formatted("LimitlessRouterService.UnableToConnectNoRoutersAvailable"),
-                                                  context.get_host_info().host) from e
+                            raise AwsWrapperError(Messages.get_formatted("LimitlessRouterService.UnableToConnectNoRoutersAvailable",
+                                                  context.get_host_info().host), e) from e
 
             try:
                 selected_host_info = self._plugin_service.get_host_info_by_strategy(
@@ -466,7 +466,7 @@ def _synchronously_get_limitless_routers_with_retry(self, context: LimitlessCont
 
                 time.sleep(retry_interval_ms)
             except InterruptedError as e:
-                raise AwsWrapperError(Messages.get("LimitlessRouterService.InterruptedSynchronousGetRouter"), e)
+                raise AwsWrapperError(Messages.get("LimitlessRouterService.InterruptedSynchronousGetRouter"), e) from e
             except Exception as e:
                 if self._is_login_exception(e):
                     raise e
diff --git a/aws_advanced_python_wrapper/okta_plugin.py b/aws_advanced_python_wrapper/okta_plugin.py
@@ -36,7 +36,7 @@
 
 import requests
 
-from aws_advanced_python_wrapper.errors import AwsWrapperError
+from aws_advanced_python_wrapper.errors import AwsConnectError, AwsWrapperError
 from aws_advanced_python_wrapper.plugin import Plugin, PluginFactory
 from aws_advanced_python_wrapper.utils.log import Logger
 from aws_advanced_python_wrapper.utils.messages import Messages
@@ -110,8 +110,11 @@ def _connect(self, host_info: HostInfo, props: Properties, connect_func: Callabl
         try:
             return connect_func()
         except Exception as e:
+            if self._plugin_service.is_network_exception(e):
+                raise AwsConnectError(Messages.get_formatted("OktaAuthPlugin.ConnectException", e)) from e
+
             if token_info is None or token_info.is_expired() or not self._plugin_service.is_login_exception(e):
-                raise e
+                raise AwsWrapperError(Messages.get_formatted("OktaAuthPlugin.ConnectException", e), e) from e
 
             self._update_authentication_token(token_host_info, props, user, region, cache_key)
 
@@ -120,7 +123,7 @@ def _connect(self, host_info: HostInfo, props: Properties, connect_func: Callabl
             except Exception as e:
                 error_message = "OktaAuthPlugin.UnhandledException"
                 logger.debug(error_message, e)
-                raise AwsWrapperError(Messages.get_formatted(error_message, e)) from e
+                raise AwsWrapperError(Messages.get_formatted(error_message, e), e) from e
 
     def force_connect(
             self,
@@ -194,7 +197,7 @@ def _get_session_token(self, props: Properties) -> str:
         except IOError as e:
             error_message = "OktaAuthPlugin.UnhandledException"
             logger.debug(error_message, e)
-            raise AwsWrapperError(Messages.get_formatted(error_message, e))
+            raise AwsWrapperError(Messages.get_formatted(error_message, e), e) from e
 
     def _get_saml_url(self, props: Properties) -> str:
         idp_endpoint = WrapperProperties.IDP_ENDPOINT.get(props)
@@ -227,7 +230,7 @@ def get_saml_assertion(self, props: Properties):
         except IOError as e:
             error_message = "OktaAuthPlugin.UnhandledException"
             logger.debug(error_message, e)
-            raise AwsWrapperError(Messages.get_formatted(error_message, e))
+            raise AwsWrapperError(Messages.get_formatted(error_message, e), e) from e
 
 
 class OktaAuthPluginFactory(PluginFactory):
diff --git a/aws_advanced_python_wrapper/plugin_service.py b/aws_advanced_python_wrapper/plugin_service.py
@@ -945,7 +945,7 @@ def execute(self, target: object, method: DbApiMethod, target_driver_func: Calla
             raise AwsWrapperError(Messages.get_formatted("PluginManager.MethodInvokedAgainstOldConnection", target))
 
         if conn is None and method in [DbApiMethod.CONNECTION_CLOSE, DbApiMethod.CURSOR_CLOSE]:
-            return
+            return None
 
         context: TelemetryContext | None
         context = self._telemetry_factory.open_telemetry_context(method.method_name, TelemetryTraceLevel.TOP_LEVEL)
diff --git a/aws_advanced_python_wrapper/resources/aws_advanced_python_wrapper_messages.properties b/aws_advanced_python_wrapper/resources/aws_advanced_python_wrapper_messages.properties
@@ -173,9 +173,11 @@ FastestResponseStrategyPlugin.UnsupportedHostSelectorStrategy=[FastestResponseSt
 
 FederatedAuthPlugin.UnhandledException=[FederatedAuthPlugin] Unhandled exception: '{}'
 FederatedAuthPlugin.UseCachedToken=[FederatedAuthPlugin] Used cached authentication token = {}
+FederatedAuthPlugin.ConnectException=[FederatedAuthPlugin] Error occurred while opening a connection: {}
 
 OktaAuthPlugin.UnhandledException=[OktaAuthPlugin] Unhandled exception: '{}'
 OktaAuthPlugin.UseCachedToken=[OktaAuthPlugin] Used cached authentication token = {}
+OktaAuthPlugin.ConnectException=[OktaAuthPlugin] Error occurred while opening a connection: {}
 
 OktaCredentialsProviderFactory.SamlAssertionUrl=[OktaCredentialsProviderFactory] Retrieving SAML assertion from URL: '{}'
 
diff --git a/aws_advanced_python_wrapper/utils/mysql_exception_handler.py b/aws_advanced_python_wrapper/utils/mysql_exception_handler.py
diff --git a/aws_advanced_python_wrapper/utils/pg_exception_handler.py b/aws_advanced_python_wrapper/utils/pg_exception_handler.py
diff --git a/tests/unit/test_exception_handling.py b/tests/unit/test_exception_handling.py
diff --git a/tests/unit/test_federated_auth_plugin.py b/tests/unit/test_federated_auth_plugin.py
diff --git a/tests/unit/test_okta_plugin.py b/tests/unit/test_okta_plugin.py
