feat: limitless plugin implementation

JuanLeee · JuanLeee · commit ff1cee640c2c · 2025-07-15T10:03:11.000-07:00
diff --git a/aws_advanced_python_wrapper/connection_provider.py b/aws_advanced_python_wrapper/connection_provider.py
@@ -26,8 +26,8 @@
 
 from aws_advanced_python_wrapper.errors import AwsWrapperError
 from aws_advanced_python_wrapper.host_selector import (
-    HostSelector, RandomHostSelector, RoundRobinHostSelector,
-    WeightedRandomHostSelector)
+    HighestWeightHostSelector, HostSelector, RandomHostSelector,
+    RoundRobinHostSelector, WeightedRandomHostSelector)
 from aws_advanced_python_wrapper.plugin import CanReleaseResources
 from aws_advanced_python_wrapper.utils.log import Logger
 from aws_advanced_python_wrapper.utils.messages import Messages
@@ -98,7 +98,8 @@ def connect(
 class DriverConnectionProvider(ConnectionProvider):
     _accepted_strategies: Dict[str, HostSelector] = {"random": RandomHostSelector(),
                                                      "round_robin": RoundRobinHostSelector(),
-                                                     "weighted_random": WeightedRandomHostSelector()}
+                                                     "weighted_random": WeightedRandomHostSelector(),
+                                                     "highest_weight": HighestWeightHostSelector()}
 
     def accepts_host_info(self, host_info: HostInfo, props: Properties) -> bool:
         return True
diff --git a/aws_advanced_python_wrapper/host_selector.py b/aws_advanced_python_wrapper/host_selector.py
@@ -260,3 +260,15 @@ def _update_host_weight_map_from_string(self, props: Optional[Properties] = None
                     except ValueError:
                         logger.error(message, pair)
                         raise AwsWrapperError(Messages.get_formatted(message, pair))
+
+
+class HighestWeightHostSelector(HostSelector):
+
+    def get_host(self, hosts: Tuple[HostInfo, ...], role: HostRole, props: Optional[Properties] = None) -> HostInfo:
+        eligible_hosts: List[HostInfo] = [host for host in hosts if
+                                          host.role == role and host.get_availability() == HostAvailability.AVAILABLE]
+
+        if len(eligible_hosts) == 0:
+            raise AwsWrapperError(Messages.get_formatted("HostSelector.NoHostsMatchingRole", role))
+
+        return max(eligible_hosts, key=lambda host: host.weight)
diff --git a/aws_advanced_python_wrapper/limitless_connection_plugin.py b/aws_advanced_python_wrapper/limitless_connection_plugin.py
@@ -35,6 +35,7 @@
     SlidingExpirationCacheWithCleanupThread
 from aws_advanced_python_wrapper.utils.telemetry.telemetry import (
     TelemetryContext, TelemetryFactory, TelemetryTraceLevel)
+from aws_advanced_python_wrapper.utils.utils import LogUtils
 
 if TYPE_CHECKING:
     from aws_advanced_python_wrapper.driver_dialect import DriverDialect
@@ -50,6 +51,10 @@ class LimitlessConnectionPlugin(Plugin):
     def __init__(self, plugin_service: PluginService, props: Properties):
         self._plugin_service = plugin_service
         self._properties = props
+        self._limitless_router_service = LimitlessRouterService(
+            self._plugin_service,
+            LimitlessQueryHelper(self._plugin_service)
+        )
 
     @property
     def subscribed_methods(self) -> Set[str]:
@@ -68,34 +73,27 @@ def connect(
 
         dialect: DatabaseDialect = self._plugin_service.database_dialect
         if not isinstance(dialect, AuroraLimitlessDialect):
-            connection = connect_func()
             refreshed_dialect = self._plugin_service.database_dialect
-
             if not isinstance(refreshed_dialect, AuroraLimitlessDialect):
                 raise UnsupportedOperationError(
                     Messages.get_formatted("LimitlessConnectionPlugin.UnsupportedDialectOrDatabase",
                                            type(refreshed_dialect).__name__))
 
-        limitless_router_service = LimitlessRouterService(
-            self._plugin_service,
-            LimitlessQueryHelper(self._plugin_service)
-        )
-
         if is_initial_connection:
-            limitless_router_service.start_monitoring(host_info, props)
+            self._limitless_router_service.start_monitoring(host_info, props)
 
-        context: LimitlessConnectionContext = LimitlessConnectionContext(
+        self._context: LimitlessConnectionContext = LimitlessConnectionContext(
             host_info,
             props,
             connection,
             connect_func,
             [],
             self
         )
-        limitless_router_service.establish_connection(context)
-        connection = context.get_connection()
+        self._limitless_router_service.establish_connection(self._context)
+        connection = self._context.get_connection()
         if connection is not None and not self._plugin_service.driver_dialect.is_closed(connection):
-            return context.get_connection()
+            return connection
 
         raise AwsWrapperError(Messages.get_formatted("LimitlessConnectionPlugin.FailedToConnectToHost", host_info.host))
 
@@ -172,6 +170,7 @@ def run(self):
                                                                    lambda _: new_limitless_routers,
                                                                    WrapperProperties.LIMITLESS_MONITOR_DISPOSAL_TIME_MS.get(
                                                                        self._properties) * 1_000_000)
+                    logger.debug(LogUtils.log_topology(tuple(new_limitless_routers), "[limitlessRouterMonitor] Topology:"))
 
                 sleep(self._interval_ms / 1000)
 
@@ -331,18 +330,19 @@ def establish_connection(self, context: LimitlessConnectionContext) -> None:
             self._plugin_service.host_list_provider.get_cluster_id(), context.get_props()))
 
         if context.get_limitless_routers() is None or len(context.get_limitless_routers()) == 0:
-            logger.debug("LimitlessRouterServiceImpl.limitlessRouterCacheEmpty")
+            logger.debug("LimitlessRouterService.:LimitlessRouterCacheEmpty")
 
             wait_for_router_info = WrapperProperties.WAIT_FOR_ROUTER_INFO.get(context.get_props())
             if wait_for_router_info:
                 self._synchronously_get_limitless_routers_with_retry(context)
             else:
-                logger.debug("LimitlessRouterServiceImpl.UsingProvidedConnectUrl")
+                logger.debug("LimitlessRouterService.UsingProvidedConnectUrl")
                 if context.get_connection() is None or self._plugin_service.driver_dialect.is_closed(context.get_connection()):
                     context.set_connection(context.get_connect_func()())
+                    return
 
-        if context.get_host_info in context.get_limitless_routers():
-            logger.debug("LimitlessRouterServiceImpl.ConnectWithHost")
+        if context.get_host_info() in context.get_limitless_routers():
+            logger.debug("LimitlessRouterService.ConnectWithHost")
             if context.get_connection() is None:
                 try:
                     context.set_connection(context.get_connect_func()())
@@ -356,7 +356,7 @@ def establish_connection(self, context: LimitlessConnectionContext) -> None:
         try:
             selected_host_info = self._plugin_service.get_host_info_by_strategy(
                 HostRole.WRITER, "weighted_random", context.get_limitless_routers())
-            logger.debug("LimitlessRouterServiceImpl.SelectedHost", "None" if selected_host_info is None else selected_host_info.host)
+            logger.debug("LimitlessRouterService.SelectedHost", "None" if selected_host_info is None else selected_host_info.host)
         except Exception as e:
             if self._is_login_exception(e) or isinstance(e, UnsupportedOperationError):
                 raise e
@@ -375,7 +375,7 @@ def establish_connection(self, context: LimitlessConnectionContext) -> None:
                 raise e
 
             if selected_host_info is not None:
-                logger.debug("LimitlessRouterServiceImpl.FailedToConnectToHost", selected_host_info.host)
+                logger.debug("LimitlessRouterService.FailedToConnectToHost", selected_host_info.host)
                 selected_host_info.set_availability(HostAvailability.UNAVAILABLE)
 
             self._retry_connection_with_least_loaded_routers(context)
@@ -400,7 +400,7 @@ def _retry_connection_with_least_loaded_routers(self, context: LimitlessConnecti
                 if (context.get_limitless_routers() is None
                         or len(context.get_limitless_routers()) == 0
                         or not context.is_any_router_available()):
-                    logger.debug("LimitlessRouterServiceImpl.NoRoutersAvailableForRetry")
+                    logger.debug("LimitlessRouterService.NoRoutersAvailableForRetry")
 
                     if context.get_connection() is not None and not self._plugin_service.driver_dialect.is_closed(context.get_connection()):
                         return
@@ -417,14 +417,14 @@ def _retry_connection_with_least_loaded_routers(self, context: LimitlessConnecti
 
             try:
                 selected_host_info = self._plugin_service.get_host_info_by_strategy(
-                    HostRole.WRITER, "weighted_random", context.get_limitless_routers())
-                logger.debug("LimitlessRouterServiceImpl.SelectedHostForRetry",
+                    HostRole.WRITER, "highest_weight", context.get_limitless_routers())
+                logger.debug("LimitlessRouterService.SelectedHostForRetry",
                              "None" if selected_host_info is None else selected_host_info.host)
                 if selected_host_info is None:
                     continue
 
             except UnsupportedOperationError as e:
-                logger.error("LimitlessRouterServiceImpl.IncorrectConfiguration")
+                logger.error("LimitlessRouterService.IncorrectConfiguration")
                 raise e
             except AwsWrapperError:
                 continue
@@ -438,14 +438,14 @@ def _retry_connection_with_least_loaded_routers(self, context: LimitlessConnecti
                 if self._is_login_exception(e):
                     raise e
                 selected_host_info.set_availability(HostAvailability.UNAVAILABLE)
-                logger.debug("LimitlessRouterServiceImpl.FailedToConnectToHost", selected_host_info.host)
+                logger.debug("LimitlessRouterService.FailedToConnectToHost", selected_host_info.host)
 
         raise AwsWrapperError(Messages.get("LimitlessRouterService.MaxRetriesExceeded"))
 
     def _synchronously_get_limitless_routers_with_retry(self, context: LimitlessConnectionContext) -> None:
-        logger.debug("LimitlessRouterServiceImpl.SynchronouslyGetLimitlessRouters")
+        logger.debug("LimitlessRouterService.SynchronouslyGetLimitlessRouters")
         retry_count = -1
-        max_retries = WrapperProperties.MAX_RETRIES_MS.get_int(context.get_props())
+        max_retries = WrapperProperties.GET_ROUTER_MAX_RETRIES.get_int(context.get_props())
         retry_interval_ms = WrapperProperties.GET_ROUTER_RETRY_INTERVAL_MS.get_float(context.get_props())
         first_iteration = True
         while first_iteration or retry_count < max_retries:
diff --git a/aws_advanced_python_wrapper/sql_alchemy_connection_provider.py b/aws_advanced_python_wrapper/sql_alchemy_connection_provider.py
@@ -26,8 +26,8 @@
 from aws_advanced_python_wrapper.connection_provider import ConnectionProvider
 from aws_advanced_python_wrapper.errors import AwsWrapperError
 from aws_advanced_python_wrapper.host_selector import (
-    HostSelector, RandomHostSelector, RoundRobinHostSelector,
-    WeightedRandomHostSelector)
+    HighestWeightHostSelector, HostSelector, RandomHostSelector,
+    RoundRobinHostSelector, WeightedRandomHostSelector)
 from aws_advanced_python_wrapper.plugin import CanReleaseResources
 from aws_advanced_python_wrapper.utils.messages import Messages
 from aws_advanced_python_wrapper.utils.properties import (Properties,
@@ -48,7 +48,8 @@ class SqlAlchemyPooledConnectionProvider(ConnectionProvider, CanReleaseResources
     _LEAST_CONNECTIONS: ClassVar[str] = "least_connections"
     _accepted_strategies: Dict[str, HostSelector] = {"random": RandomHostSelector(),
                                                      "round_robin": RoundRobinHostSelector(),
-                                                     "weighted_random": WeightedRandomHostSelector()}
+                                                     "weighted_random": WeightedRandomHostSelector(),
+                                                     "highest_weight": HighestWeightHostSelector()}
     _rds_utils: ClassVar[RdsUtils] = RdsUtils()
     _database_pools: ClassVar[SlidingExpirationCache[PoolKey, QueuePool]] = SlidingExpirationCache(
         should_dispose_func=lambda queue_pool: queue_pool.checkedout() == 0,
diff --git a/tests/unit/test_highest_weight_host_selector.py b/tests/unit/test_highest_weight_host_selector.py
@@ -0,0 +1,61 @@
+#  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License").
+#  You may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#  http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  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 aws_advanced_python_wrapper.host_availability import HostAvailability
+from aws_advanced_python_wrapper.host_selector import HighestWeightHostSelector
+from aws_advanced_python_wrapper.hostinfo import HostInfo, HostRole
+from aws_advanced_python_wrapper.utils.properties import Properties
+
+HOST_ROLE = HostRole.READER
+
+
+def test_get_host_given_unavailable_host():
+    unavailable_host: HostInfo = HostInfo(host="some_unavailable_host", role=HOST_ROLE, availability=HostAvailability.UNAVAILABLE)
+    available_host: HostInfo = HostInfo(host="some_available_host", role=HOST_ROLE, availability=HostAvailability.AVAILABLE)
+
+    host_selector = HighestWeightHostSelector()
+    actual_host = host_selector.get_host((unavailable_host, available_host), HOST_ROLE, Properties())
+
+    assert available_host == actual_host
+
+
+def test_get_host_given_multiple_unavailable_hosts():
+    hosts = (
+        HostInfo(host="some_unavailable_host", role=HOST_ROLE, availability=HostAvailability.UNAVAILABLE),
+        HostInfo(host="some_unavailable_host", role=HOST_ROLE, availability=HostAvailability.UNAVAILABLE),
+        HostInfo(host="some_available_host", role=HOST_ROLE, availability=HostAvailability.AVAILABLE)
+    )
+
+    host_selector = HighestWeightHostSelector()
+    actual_host = host_selector.get_host(hosts, HOST_ROLE, Properties())
+
+    assert HostAvailability.AVAILABLE == actual_host.get_availability()
+
+
+def test_get_host_given_different_weights():
+
+    highest_weight_host = HostInfo(host="some_available_host", role=HOST_ROLE, availability=HostAvailability.AVAILABLE, weight=3)
+
+    hosts = (
+        HostInfo(host="some_unavailable_host", role=HOST_ROLE, availability=HostAvailability.UNAVAILABLE),
+        HostInfo(host="some_unavailable_host", role=HOST_ROLE, availability=HostAvailability.UNAVAILABLE),
+        HostInfo(host="some_available_host", role=HOST_ROLE, availability=HostAvailability.AVAILABLE, weight=1),
+        HostInfo(host="some_available_host", role=HOST_ROLE, availability=HostAvailability.AVAILABLE, weight=2),
+        highest_weight_host
+    )
+
+    host_selector = HighestWeightHostSelector()
+    actual_host = host_selector.get_host(hosts, HOST_ROLE, Properties())
+
+    assert actual_host == highest_weight_host
diff --git a/tests/unit/test_limitless_connection_plugin.py b/tests/unit/test_limitless_connection_plugin.py
diff --git a/tests/unit/test_limitless_connection_service.py b/tests/unit/test_limitless_connection_service.py
