fix: bugs in host list provider

Author: karenc-bq

aws_advanced_python_wrapper/aurora_connection_tracker_plugin.py

@@ -49,14 +49,31 @@ def populate_opened_connection_set(self, host_info: HostInfo, conn: Connection):
         :param conn: currently opened connection.
         """
 
+        #     // Check if the connection was established using an instance endpoint
+        #     if (rdsUtils.isRdsInstance(hostSpec.getHost())) {
+        #       trackConnection(hostSpec.getHostAndPort(), conn);
+        #       return;
+        #     }
+        #
+        #     final String instanceEndpoint = aliases.stream()
+        #         .filter(x -> rdsUtils.isRdsInstance(rdsUtils.removePort(x)))
+        #         .max(String::compareToIgnoreCase)
+        #         .orElse(null);
+        #
+        #     if (instanceEndpoint == null) {
+        #       LOGGER.finest(
+        #           Messages.get("OpenedConnectionTracker.unableToPopulateOpenedConnectionQueue",
+        #             new Object[] {hostSpec.getHost()}));
+        #       return;
+        #     }
+
         aliases: FrozenSet[str] = host_info.as_aliases()
-        host: str = host_info.as_alias()
 
-        if self._rds_utils.is_rds_instance(host):
-            self._track_connection(host, conn)
+        if self._rds_utils.is_rds_instance(host_info.host):
+            self._track_connection(host_info.as_alias(), conn)
             return
 
-        instance_endpoint: Optional[str] = next((alias for alias in aliases if self._rds_utils.is_rds_instance(alias)),
+        instance_endpoint: Optional[str] = next((alias for alias in aliases if self._rds_utils.is_rds_instance(self._rds_utils.remove_port(alias))),
                                                 None)
         if not instance_endpoint:
             logger.debug("OpenedConnectionTracker.UnableToPopulateOpenedConnectionSet")

aws_advanced_python_wrapper/sql_alchemy_connection_provider.py

@@ -57,10 +57,12 @@ def __init__(
             self,
             pool_configurator: Optional[Callable] = None,
             pool_mapping: Optional[Callable] = None,
+            accept_url_func: Optional[Callable] = None,
             pool_expiration_check_ns: int = -1,
             pool_cleanup_interval_ns: int = -1):
         self._pool_configurator = pool_configurator
         self._pool_mapping = pool_mapping
+        self._accept_url_func = accept_url_func
 
         if pool_expiration_check_ns > -1:
             SqlAlchemyPooledConnectionProvider._POOL_EXPIRATION_CHECK_NS = pool_expiration_check_ns
@@ -80,6 +82,8 @@ def keys(self):
         return self._database_pools.keys()
 
     def accepts_host_info(self, host_info: HostInfo, props: Properties) -> bool:
+        if self._accept_url_func:
+            return self._accept_url_func(host_info, props)
         url_type = SqlAlchemyPooledConnectionProvider._rds_utils.identify_rds_type(host_info.host)
         return RdsUrlType.RDS_INSTANCE == url_type
 

aws_advanced_python_wrapper/utils/rdsutils.py

@@ -267,6 +267,13 @@ def _get_group(self, host: str, group: str):
     def _get_dns_group(self, host: str):
         return self._get_group(host, RdsUtils.DNS_GROUP)
 
+    def remove_port(self, url: str):
+        if not url or not url.strip():
+            return None
+        if ":" in url:
+            return url.split(":")[0]
+        return url
+
     @staticmethod
     def clear_cache():
         RdsUtils.CACHE_PATTERNS.clear()

tests/integration/container/test_read_write_splitting.py

@@ -515,7 +515,8 @@ def test_pooled_connection__cluster_url_failover(
     def test_pooled_connection__failover_failed(
             self, test_environment: TestEnvironment, test_driver: TestDriver,
             rds_utils, conn_utils, proxied_failover_props):
-        provider = SqlAlchemyPooledConnectionProvider(lambda _, __: {"pool_size": 1})
+        writer_host = test_environment.get_writer().get_host()
+        provider = SqlAlchemyPooledConnectionProvider(lambda _, __: {"pool_size": 1}, lambda host_info, props: writer_host in host_info.host)
         ConnectionProviderManager.set_connection_provider(provider)
 
         WrapperProperties.PLUGINS.set(proxied_failover_props, "read_write_splitting,failover,host_monitoring")