feat: adding weighted random host selection strategy

JuanLeee · JuanLeee · commit 7e4f6d44f1b1 · 2025-07-03T03:51:31.000-07:00
diff --git a/aws_advanced_python_wrapper/connection_provider.py b/aws_advanced_python_wrapper/connection_provider.py
@@ -25,9 +25,9 @@
 from threading import Lock
 
 from aws_advanced_python_wrapper.errors import AwsWrapperError
-from aws_advanced_python_wrapper.host_selector import (HostSelector,
-                                                       RandomHostSelector,
-                                                       RoundRobinHostSelector)
+from aws_advanced_python_wrapper.host_selector import (
+    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
@@ -96,8 +96,9 @@ def connect(
 
 
 class DriverConnectionProvider(ConnectionProvider):
-    _accepted_strategies: Dict[str, HostSelector] = \
-        {"random": RandomHostSelector(), "round_robin": RoundRobinHostSelector()}
+    _accepted_strategies: Dict[str, HostSelector] = {"random": RandomHostSelector(),
+                                                     "round_robin": RoundRobinHostSelector(),
+                                                     "weighted_random": WeightedRandomHostSelector()}
 
     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
@@ -18,6 +18,7 @@
 from re import search
 from typing import TYPE_CHECKING, Dict, List, Optional, Protocol, Tuple
 
+from aws_advanced_python_wrapper.utils.log import Logger
 from .host_availability import HostAvailability
 
 if TYPE_CHECKING:
@@ -29,6 +30,8 @@
 from .utils.messages import Messages
 from .utils.properties import Properties, WrapperProperties
 
+logger = Logger(__name__)
+
 
 class HostSelector(Protocol):
     """
@@ -168,23 +171,92 @@ def _update_cache_properties_for_round_robin_cluster_info(self, round_robin_clus
 
                 for pair in host_weight_pairs:
                     match = search(RoundRobinHostSelector._HOST_WEIGHT_PAIRS_PATTERN, pair)
+                    message = "RoundRobinHostSelector.RoundRobinInvalidHostWeightPairs"
                     if match:
                         host_name = match.group("host")
                         host_weight = match.group("weight")
                     else:
-                        raise AwsWrapperError(Messages.get("RoundRobinHostSelector.RoundRobinInvalidHostWeightPairs"))
+                        logger.error(message, pair)
+                        raise AwsWrapperError(Messages.get_formatted(message, pair))
 
                     if len(host_name) == 0 or len(host_weight) == 0:
-                        raise AwsWrapperError(Messages.get("RoundRobinHostSelector.RoundRobinInvalidHostWeightPairs"))
+                        logger.error(message, pair)
+                        raise AwsWrapperError(Messages.get_formatted(message, pair))
                     try:
                         weight: int = int(host_weight)
 
                         if weight < RoundRobinHostSelector._DEFAULT_WEIGHT:
-                            raise AwsWrapperError(Messages.get("RoundRobinHostSelector.RoundRobinInvalidHostWeightPairs"))
+                            logger.error(message, pair)
+                            raise AwsWrapperError(Messages.get_formatted(message, pair))
 
                         round_robin_cluster_info.cluster_weights_dict[host_name] = weight
                     except ValueError:
-                        raise AwsWrapperError(Messages.get("RoundRobinHostSelector.RoundRobinInvalidHostWeightPairs"))
+                        logger.error(message, pair)
+                        raise AwsWrapperError(Messages.get_formatted(message, pair))
 
     def clear_cache(self):
         RoundRobinHostSelector._round_robin_cache.clear()
+
+
+class WeightedRandomHostSelector(HostSelector):
+    _DEFAULT_WEIGHT: int = 1
+    _HOST_WEIGHT_PAIRS_PATTERN = r"((?P<host>[^:/?#]*):(?P<weight>.*))"
+    _host_weight_map: Dict[str, int] = {}
+
+    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]
+        eligible_hosts.sort(key=lambda host: host.host, reverse=False)
+        if len(eligible_hosts) == 0:
+            message = "HostSelector.NoHostsMatchingRole"
+            logger.error(message, role)
+            raise AwsWrapperError(Messages.get_formatted("HostSelector.NoHostsMatchingRole", role))
+
+        self._update_host_weight_map_from_string(props)
+
+        default_weight: int = WeightedRandomHostSelector._DEFAULT_WEIGHT
+        if props is not None:
+            default_weight = WrapperProperties.WEIGHTED_RANDOM_DEFAULT_WEIGHT.get_int(props)
+            if default_weight < WeightedRandomHostSelector._DEFAULT_WEIGHT:
+                logger.error("WeightedRandomHostSelector.WeightedRandomInvalidDefaultWeight")
+                raise AwsWrapperError(Messages.get("WeightedRandomHostSelector.WeightedRandomInvalidDefaultWeight"))
+
+        selection_list: List[HostInfo] = []
+        for host in eligible_hosts:
+            if host.host in self._host_weight_map:
+                selection_list = selection_list + self._host_weight_map[host.host] * [host]
+            else:
+                selection_list = selection_list + default_weight * [host]
+
+        return random.choice(selection_list)
+
+    def _update_host_weight_map_from_string(self, props: Optional[Properties] = None) -> None:
+        if props is not None:
+            host_weights: Optional[str] = WrapperProperties.WEIGHTED_RANDOM_HOST_WEIGHT_PAIRS.get(props)
+            if host_weights is not None and len(host_weights) != 0:
+                host_weight_pairs: List[str] = host_weights.split(",")
+
+                for pair in host_weight_pairs:
+                    match = search(WeightedRandomHostSelector._HOST_WEIGHT_PAIRS_PATTERN, pair)
+                    message = "WeightedRandomHostSelector.WeightedRandomInvalidHostWeightPairs"
+                    if match:
+                        host_name = match.group("host")
+                        host_weight = match.group("weight")
+                    else:
+                        logger.error(message, pair)
+                        raise AwsWrapperError(Messages.get_formatted(message, pair))
+
+                    if len(host_name) == 0 or len(host_weight) == 0:
+                        logger.error(message, pair)
+                        raise AwsWrapperError(Messages.get_formatted(message, pair))
+                    try:
+                        weight: int = int(host_weight)
+
+                        if weight < WeightedRandomHostSelector._DEFAULT_WEIGHT:
+                            logger.error(message, pair)
+                            raise AwsWrapperError(Messages.get_formatted(message, pair))
+
+                        self._host_weight_map[host_name] = weight
+                    except ValueError:
+                        logger.error(message, pair)
+                        raise AwsWrapperError(Messages.get_formatted(message, pair))
diff --git a/aws_advanced_python_wrapper/resources/aws_advanced_python_wrapper_messages.properties b/aws_advanced_python_wrapper/resources/aws_advanced_python_wrapper_messages.properties
@@ -279,7 +279,10 @@ ReadWriteSplittingPlugin.UnsupportedHostInfoSelectorStrategy=[ReadWriteSplitting
 
 RoundRobinHostSelector.ClusterInfoNone=[RoundRobinHostSelector] The round robin cluster information cache should have an entry for the current cluster, but no entry was found.
 RoundRobinHostSelector.RoundRobinInvalidDefaultWeight=[RoundRobinHostSelector] The provided default weight value is not valid. Weight values must be an integer greater than or equal to 1.
-RoundRobinHostSelector.RoundRobinInvalidHostWeightPairs= [RoundRobinHostSelector] The provided host weight pairs have not been configured correctly. Please ensure the provided host weight pairs is a comma separated list of pairs, each pair in the format of <host>:<weight>. Weight values must be an integer greater than or equal to the default weight value of 1.
+RoundRobinHostSelector.RoundRobinInvalidHostWeightPairs= [RoundRobinHostSelector] The provided host weight pairs have not been configured correctly. Please ensure the provided host weight pairs is a comma separated list of pairs, each pair in the format of <host>:<weight>. Weight values must be an integer greater than or equal to the default weight value of 1. Weight pair: '{}'
+
+WeightedRandomHostSelector.WeightedRandomInvalidHostWeightPairs= [WeightedRandomHostSelector] The provided host weight pairs have not been configured correctly. Please ensure the provided host weight pairs is a comma separated list of pairs, each pair in the format of <host>:<weight>. Weight values must be an integer greater than or equal to the default weight value of 1. Weight pair: '{}'
+WeightedRandomHostSelector.WeightedRandomInvalidDefaultWeight=[WeightedRandomHostSelector] The provided default weight value is not valid. Weight values must be an integer greater than or equal to 1.
 
 SlidingExpirationCache.CleaningUp=[SlidingExpirationCache] Cleaning up...
 
diff --git a/aws_advanced_python_wrapper/sql_alchemy_connection_provider.py b/aws_advanced_python_wrapper/sql_alchemy_connection_provider.py
@@ -25,9 +25,9 @@
 
 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)
+from aws_advanced_python_wrapper.host_selector import (
+    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,
@@ -46,7 +46,9 @@ class SqlAlchemyPooledConnectionProvider(ConnectionProvider, CanReleaseResources
     """
     _POOL_EXPIRATION_CHECK_NS: ClassVar[int] = 30 * 60_000_000_000  # 30 minutes
     _LEAST_CONNECTIONS: ClassVar[str] = "least_connections"
-    _accepted_strategies: Dict[str, HostSelector] = {"random": RandomHostSelector(), "round_robin": RoundRobinHostSelector()}
+    _accepted_strategies: Dict[str, HostSelector] = {"random": RandomHostSelector(),
+                                                     "round_robin": RoundRobinHostSelector(),
+                                                     "weighted_random": WeightedRandomHostSelector()}
     _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/aws_advanced_python_wrapper/utils/properties.py b/aws_advanced_python_wrapper/utils/properties.py
@@ -276,6 +276,15 @@ class WrapperProperties:
     ROUND_ROBIN_HOST_WEIGHT_PAIRS = WrapperProperty("round_robin_host_weight_pairs",
                                                     "Comma separated list of database host-weight pairs in the format of `<host>:<weight>`.",
                                                     "")
+
+    WEIGHTED_RANDOM_DEFAULT_WEIGHT = WrapperProperty("weighted_random_default_weight", "The default weight for any hosts that have not been " +
+                                                     "configured with the `weighted_random_host_weight_pairs` parameter.",
+                                                     1)
+
+    WEIGHTED_RANDOM_HOST_WEIGHT_PAIRS = WrapperProperty("weighted_random_host_weight_pairs",
+                                                        "Comma separated list of database host-weight pairs in the format of `<host>:<weight>`.",
+                                                        "")
+
     # Federated Auth Plugin
     IDP_ENDPOINT = WrapperProperty("idp_endpoint",
                                    "The hosting URL of the Identity Provider",
diff --git a/docs/using-the-python-driver/ReaderSelectionStrategies.md b/docs/using-the-python-driver/ReaderSelectionStrategies.md
@@ -4,6 +4,9 @@ To balance connections to reader instances more evenly, different selection stra
 | Reader Selection Strategy | Configuration Parameter                               | Description                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                       | Default Value |
 |---------------------------|-------------------------------------------------------|-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|---------------|
 | `random`                  | This strategy does not have configuration parameters. | The random strategy is the default selection strategy. When switching to a reader connection, the reader instance will be chosen randomly from the available database instances.                                                                                                                                                                                                                                                                                                                                                                                  | N/A           |
+| `weighted_random`         | See the following rows for configuration parameters.  | The weighted random strategy will be chosen randomly from the available database instances. A slight addition to the random strategy is the weighted random strategy, where more connections will be passed to reader instances based on user specified connection properties.                                                                                                                                                                                                                                                                                    | N/A           |
+|                           | `weighted_random_host_weight_pairs`                   | This parameter value must be a `string` type comma separated list of database host-weight pairs in the format `<host>:<weight>`. The host represents the database instance name, and the weight represents the likeliness of the connection to be directed to the host. Larger the number, the more likely the connection is to be directed to the host. <br><br> **Note:** The `<weight>` value in the string must be an integer greater than or equal to 1.                                                                                                     | `null`        |
+|                           | `weighted_random_default_weight`                      | This parameter value must be an integer value in the form of a `string`. This parameter represents the default weight for any hosts that have not been configured with the `weighted_random_host_weight_pairs` parameter. For example, if a connection were already established and host weights were set with `weighted_random_host_weight_pairs` but a new reader host was added to the database, the new reader host would use the default weight. <br><br> **Note:** This value must be an integer greater than or equal to 1.                                | `1`           |
 | `least_connections`       | This strategy does not have configuration parameters. | The least connections strategy will select reader instances based on which database instance has the least number of currently active connections. Note that this strategy is only available when internal connection pools are enabled - if you set the connection property without enabling internal pools, an exception will be thrown.                                                                                                                                                                                                                        | N/A           |
 | `round_robin`             | See the following rows for configuration parameters.  | The round robin strategy will select a reader instance by taking turns with all available database instances in a cycle. A slight addition to the round robin strategy is the weighted round robin strategy, where more connections will be passed to reader instances based on user specified connection properties.                                                                                                                                                                                                                                             | N/A           |
 |                           | `round_robin_host_weight_pairs`                       | This parameter value must be a `string` type comma separated list of database host-weight pairs in the format `<host>:<weight>`. The host represents the database instance name, and the weight represents how many connections should be directed to the host in one cycle through all available hosts. For example, the value `instance-1:1,instance-2:4` means that for every connection to `instance-1`, there will be four connections to `instance-2`. <br><br> **Note:** The `<weight>` value in the string must be an integer greater than or equal to 1. | `null`        |
diff --git a/docs/using-the-python-driver/using-plugins/UsingTheReadWriteSplittingPlugin.md b/docs/using-the-python-driver/using-plugins/UsingTheReadWriteSplittingPlugin.md
diff --git a/tests/test_weight_random_host_selector.py b/tests/test_weight_random_host_selector.py
