fix: populate opened connection queue with url (#1094)

Author: karenc-bq

.github/workflows/integration_tests.yml

@@ -71,8 +71,8 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: [ "3.11", "3.12", "3.13"  ]
-        environment: ["mysql", "pg"]
+        python-version: [ "3.11", "3.12", "3.13" ]
+        environment: [ "mysql", "pg" ]
 
     steps:
       - name: 'Clone repository'

aws_advanced_python_wrapper/aurora_connection_tracker_plugin.py

@@ -124,7 +124,7 @@ def populate_opened_connection_set(self, host_info: HostInfo, conn: Connection):
         aliases: FrozenSet[str] = host_info.as_aliases()
 
         if self._rds_utils.is_rds_instance(host_info.host):
-            self._track_connection(host_info.as_alias(), conn)
+            self._track_connection(host_info.url, conn)
             return
 
         instance_endpoint: Optional[str] = next(
@@ -133,6 +133,7 @@ def populate_opened_connection_set(self, host_info: HostInfo, conn: Connection):
             logger.debug("OpenedConnectionTracker.UnableToPopulateOpenedConnectionSet")
             return
 
+        instance_endpoint = instance_endpoint if instance_endpoint.endswith("/") else instance_endpoint + "/"
         self._track_connection(instance_endpoint, conn)
 
     def invalidate_all_connections(self, host_info: Optional[HostInfo] = None, host: Optional[FrozenSet[str]] = None):
@@ -144,7 +145,7 @@ def invalidate_all_connections(self, host_info: Optional[HostInfo] = None, host:
         """
 
         if host_info:
-            self.invalidate_all_connections(host=frozenset([host_info.as_alias()]))
+            self.invalidate_all_connections(host=frozenset([host_info.url]))
             self.invalidate_all_connections(host=host_info.as_aliases())
             return
 
@@ -160,6 +161,8 @@ def invalidate_all_connections(self, host_info: Optional[HostInfo] = None, host:
         if not instance_endpoint:
             return
 
+        instance_endpoint = instance_endpoint if instance_endpoint.endswith("/") else instance_endpoint + "/"
+
         with self._lock:
             connection_set: Optional[WeakSet] = self._opened_connections.get(instance_endpoint)
             connections_list = list(connection_set) if connection_set is not None else None
@@ -173,14 +176,16 @@ def remove_connection_tracking(self, host_info: HostInfo, connection: Connection
             return
 
         if self._rds_utils.is_rds_instance(host_info.host):
-            host = host_info.as_alias()
+            host = host_info.url
         else:
             host = next((alias for alias in host_info.as_aliases()
                          if self._rds_utils.is_rds_instance(self._rds_utils.remove_port(alias))), "")
 
         if not host:
             return
 
+        host = host if host.endswith("/") else host + "/"
+
         with self._lock:
             connection_set = self._opened_connections.get(host)
             if connection_set:

aws_advanced_python_wrapper/failover_v2_plugin.py

@@ -149,7 +149,7 @@ def connect(
                 break
 
         conn = None
-        if (host_with_availability is None or host_with_availability.availability != HostAvailability.UNAVAILABLE):
+        if host_with_availability is None or host_with_availability.availability != HostAvailability.UNAVAILABLE:
             try:
                 conn = self._stale_dns_helper.get_verified_connection(
                     is_initial_connection, self._host_list_provider_service,

aws_advanced_python_wrapper/utils/pg_exception_handler.py

@@ -28,11 +28,13 @@ class PgExceptionHandler(ExceptionHandler):
     _CONNECTION_FAILED = "connection failed"
     _CONSUMING_INPUT_FAILED = "consuming input failed"
     _CONNECTION_SOCKET_CLOSED = "connection socket closed"
+    _CONNECTION_CLOSED = "the connection is closed"
 
     _NETWORK_ERROR_MESSAGES: List[str] = [
         _CONNECTION_FAILED,
         _CONSUMING_INPUT_FAILED,
-        _CONNECTION_SOCKET_CLOSED
+        _CONNECTION_SOCKET_CLOSED,
+        _CONNECTION_CLOSED,
     ]
     _ACCESS_ERROR_MESSAGES: List[str] = [
         _PASSWORD_AUTHENTICATION_FAILED_MSG,

aws_advanced_python_wrapper/wrapper.py

@@ -22,8 +22,7 @@
 
 from aws_advanced_python_wrapper.driver_dialect_manager import \
     DriverDialectManager
-from aws_advanced_python_wrapper.errors import (AwsWrapperError,
-                                                FailoverSuccessError)
+from aws_advanced_python_wrapper.errors import AwsWrapperError
 from aws_advanced_python_wrapper.pep249 import Connection, Cursor, Error
 from aws_advanced_python_wrapper.pep249_methods import DbApiMethod
 from aws_advanced_python_wrapper.plugin import CanReleaseResources
@@ -285,12 +284,8 @@ def execute(
             *args: Any,
             **kwargs: Any
     ) -> AwsWrapperCursor:
-        try:
-            return self._plugin_manager.execute(self.target_cursor, DbApiMethod.CURSOR_EXECUTE,
-                                                lambda: self.target_cursor.execute(*args, **kwargs), *args, **kwargs)
-        except FailoverSuccessError as e:
-            self._target_cursor = self.connection.target_connection.cursor()
-            raise e
+        return self._plugin_manager.execute(self.target_cursor, DbApiMethod.CURSOR_EXECUTE,
+                                            lambda: self.target_cursor.execute(*args, **kwargs), *args, **kwargs)
 
     def executemany(
             self,