feat: limitless plugin implementation

Author: JuanLeee

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

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)

aws_advanced_python_wrapper/limitless_connection_plugin.py

@@ -363,7 +363,7 @@ def establish_connection(self, context: LimitlessConnectionContext) -> None:
 
             self._retry_connection_with_least_loaded_routers(context)
             return
-
+        selected_host_info = None
         if selected_host_info is None:
             self._retry_connection_with_least_loaded_routers(context)
             return
@@ -417,7 +417,7 @@ 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())
+                    HostRole.WRITER, "highest_weight", context.get_limitless_routers())
                 logger.debug("LimitlessRouterServiceImpl.SelectedHostForRetry",
                              "None" if selected_host_info is None else selected_host_info.host)
                 if selected_host_info is None:

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,

tests/unit/test_highest_weight_host_selector.py

@@ -0,0 +1,66 @@
+#  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.
+
+import pytest
+
+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
+
+
+@pytest.mark.parametrize("execution_number", range(50))
+def test_get_host_given_unavailable_host(execution_number):
+    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
+
+
+@pytest.mark.parametrize("execution_number", range(50))
+def test_get_host_given_multiple_unavailable_hosts(execution_number):
+    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()
+
+
+@pytest.mark.parametrize("execution_number", range(50))
+def test_get_host_given_different_weights(execution_number):
+
+    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