chore: catch segfault on connect

Author: karenc-bq

aws_advanced_python_wrapper/default_plugin.py

@@ -14,6 +14,7 @@
 
 from __future__ import annotations
 
+import signal
 from typing import TYPE_CHECKING, List, Optional
 
 if TYPE_CHECKING:
@@ -28,7 +29,7 @@
 import copy
 from typing import Any, Callable, Set
 
-from aws_advanced_python_wrapper.errors import AwsWrapperError
+from aws_advanced_python_wrapper.errors import AwsWrapperError, SegfaultOnConnectError
 from aws_advanced_python_wrapper.host_availability import HostAvailability
 from aws_advanced_python_wrapper.hostinfo import HostInfo, HostRole
 from aws_advanced_python_wrapper.plugin import Plugin
@@ -40,10 +41,17 @@
 class DefaultPlugin(Plugin):
     _SUBSCRIBED_METHODS: Set[str] = {"*"}
     _CLOSE_METHOD = "Connection.close"
+    _in_connect = False
 
     def __init__(self, plugin_service: PluginService, connection_provider_manager: ConnectionProviderManager):
         self._plugin_service: PluginService = plugin_service
         self._connection_provider_manager = connection_provider_manager
+        signal.signal(signal.SIGSEGV, self._segfault_handler)
+
+    def _segfault_handler(self, signum, frame):
+        if DefaultPlugin._in_connect:
+            raise SegfaultOnConnectError(Messages.get("DefaultPlugin.SegFault"))
+        signal.default_int_handler(signum, frame)
 
     def connect(
             self,
@@ -71,9 +79,11 @@ def _connect(
 
         conn: Connection
         try:
+            DefaultPlugin._in_connect = True
             database_dialect = self._plugin_service.database_dialect
             conn = conn_provider.connect(target_func, driver_dialect, database_dialect, host_info, props)
         finally:
+            DefaultPlugin._in_connect = False
             context.close_context()
 
         self._plugin_service.set_availability(host_info.all_aliases, HostAvailability.AVAILABLE)

aws_advanced_python_wrapper/errors.py

@@ -45,3 +45,7 @@ class FailoverSuccessError(FailoverError):
 
 class ReadWriteSplittingError(AwsWrapperError):
     __module__ = "aws_advanced_python_wrapper"
+
+
+class SegfaultOnConnectError(AwsWrapperError):
+    __module__ = "aws_advanced_python_wrapper"

aws_advanced_python_wrapper/resources/aws_advanced_python_wrapper_messages.properties

@@ -99,6 +99,7 @@ CustomEndpointPlugin.WaitingForCustomEndpointInfo=[CustomEndpointPlugin] Custom
 
 DefaultPlugin.EmptyHosts=[DefaultPlugin] The default connection plugin received an empty host list from the plugin service.
 DefaultPlugin.UnknownHosts=[DefaultPlugin] A HostInfo with the role of HostRole.UNKNOWN was requested via get_host_info_by_strategy. The requested role must be either HostRole.WRITER or HostRole.READER.
+DefaultPlugin.SegFault=[DefaultPlugin] Segmentation fault thrown during connect
 
 DefaultTelemetryFactory.InvalidMetricsBackend=[DefaultTelemetryFactory] {} is not a valid metrics backend. Available options: OTLP, NONE.
 DefaultTelemetryFactory.InvalidTracingBackend=[DefaultTelemetryFactory] {} is not a valid tracing backend. Available options: OTLP, XRAY, NONE.

aws_advanced_python_wrapper/utils/mysql_exception_handler.py

@@ -16,7 +16,7 @@
 
 from mysql.connector import DatabaseError, InterfaceError
 
-from aws_advanced_python_wrapper.errors import QueryTimeoutError
+from aws_advanced_python_wrapper.errors import QueryTimeoutError, SegfaultOnConnectError
 from aws_advanced_python_wrapper.exception_handling import ExceptionHandler
 
 
@@ -37,6 +37,9 @@ class MySQLExceptionHandler(ExceptionHandler):
     ]
 
     def is_network_exception(self, error: Optional[Exception] = None, sql_state: Optional[str] = None) -> bool:
+        if isinstance(error, SegfaultOnConnectError):
+            return True
+
         if isinstance(error, QueryTimeoutError):
             return True
 

aws_advanced_python_wrapper/utils/pg_exception_handler.py

@@ -18,7 +18,7 @@
                             InvalidAuthorizationSpecification, InvalidPassword,
                             OperationalError)
 
-from aws_advanced_python_wrapper.errors import QueryTimeoutError
+from aws_advanced_python_wrapper.errors import QueryTimeoutError, SegfaultOnConnectError
 from aws_advanced_python_wrapper.exception_handling import ExceptionHandler
 
 
@@ -42,6 +42,9 @@ class PgExceptionHandler(ExceptionHandler):
     _ACCESS_ERROR_CODES: List[str]
 
     def is_network_exception(self, error: Optional[Exception] = None, sql_state: Optional[str] = None) -> bool:
+        if isinstance(error, SegfaultOnConnectError):
+            return True
+
         if isinstance(error, QueryTimeoutError) or isinstance(error, ConnectionTimeout):
             return True
         if sql_state is None: