feat: adding weighted random host selection strategy

Author: JuanLeee

aws_advanced_python_wrapper/connection_provider.py

@@ -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": RandomHostSelector()}
 
     def accepts_host_info(self, host_info: HostInfo, props: Properties) -> bool:
         return True

aws_advanced_python_wrapper/host_selector.py

@@ -188,3 +188,61 @@ def _update_cache_properties_for_round_robin_cluster_info(self, round_robin_clus
 
     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:
+            raise AwsWrapperError(Messages.get_formatted("HostSelector.NoHostsMatchingRole", role))
+
+        self._update_host_weight_map_from_string()
+
+        default_weight: int = RoundRobinHostSelector._DEFAULT_WEIGHT
+        if props is not None:
+            props_weight = WrapperProperties.WEIGHTED_RANDOM_DEFAULT_WEIGHT.get_int(props)
+            if props_weight < WeightedRandomHostSelector._DEFAULT_WEIGHT:
+                raise AwsWrapperError(Messages.get("RoundRobinHostSelector.RoundRobinInvalidDefaultWeight"))
+            elif props_weight:
+                default_weight = props_weight
+
+        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)
+                    if match:
+                        host_name = match.group("host")
+                        host_weight = match.group("weight")
+                    else:
+                        raise AwsWrapperError(Messages.get("WeightedRandomHostSelector.WeightedRandomInvalidHostWeightPairs"))
+
+                    if len(host_name) == 0 or len(host_weight) == 0:
+                        raise AwsWrapperError(Messages.get("WeightedRandomHostSelector.WeightedRandomInvalidHostWeightPairs"))
+                    try:
+                        weight: int = int(host_weight)
+
+                        if weight < self._DEFAULT_WEIGHT:
+                            raise AwsWrapperError(Messages.get("WeightedRandomHostSelector.WeightedRandomInvalidHostWeightPairs"))
+
+                        self._host_weight_map[host_name] = weight
+                    except ValueError:
+                        raise AwsWrapperError(Messages.get("WeightedRandomHostSelector.WeightedRandomInvalidHostWeightPairs"))

aws_advanced_python_wrapper/resources/aws_advanced_python_wrapper_messages.properties

@@ -280,6 +280,8 @@ RoundRobinHostSelector.ClusterInfoNone=[RoundRobinHostSelector] The round robin
 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.
 
+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.
+
 SlidingExpirationCache.CleaningUp=[SlidingExpirationCache] Cleaning up...
 
 SqlAlchemyPooledConnectionProvider.PoolNone=[SqlAlchemyPooledConnectionProvider] Attempted to find or create a pool for '{}' but the result of the attempt evaluated to None.

aws_advanced_python_wrapper/sql_alchemy_connection_provider.py

@@ -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": RandomHostSelector()}
     _rds_utils: ClassVar[RdsUtils] = RdsUtils()
     _database_pools: ClassVar[SlidingExpirationCache[PoolKey, QueuePool]] = SlidingExpirationCache(
         should_dispose_func=lambda queue_pool: queue_pool.checkedout() == 0,

aws_advanced_python_wrapper/utils/properties.py

@@ -271,6 +271,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",

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_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`        |

docs/using-the-python-driver/using-plugins/UsingTheReadWriteSplittingPlugin.md

@@ -104,6 +104,7 @@ conn = AwsWrapperConnection.connect(psycopg.Connection.connect, **params)
 | `random`                        | The random strategy is the default connection strategy. When switching to a reader connection, the reader instance will be chosen randomly from the available database instances.                                                                                                                                                                                                                                                                                                                                                                                 | N/A                                                                                                        |
 | `least_connections`             | 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`                   | 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.                                                                                                                                                                                                                                             | See the rows `round_robin_host_weight_pairs` and `round_robin_default_weight` for configuration parameters |
+| `weighted_random`               | 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.                                                                                                                                                                                                                                                                                    | See the rows `round_robin_host_weight_pairs` and `round_robin_default_weight` for configuration parameters |
 | `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. | `""` (empty string)                                                                                        |
 | `round_robin_default_weight`    | This parameter value must be an integer value. This parameter represents the default weight for any hosts that have not been configured with the `round_robin_host_weight_pairs` parameter. For example, if a connection were already established and host weights were set with `round_robin_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                                                                                                          |
 | `fastest_response`              | The fastest_response strategy will select reader instances based on which database instance has the fastest response time. Note that this strategy requires that the `fastest_response_strategy` and `read_write_splitting` plugins are both enabled. See [`Fastest Response Strategy Plugin`](./UsingTheFastestResponseStrategyPlugin.md)                                                                                                                                                                                                                        | N/A                                                                                                        |

tests/test_weight_random_host_selector.py

@@ -0,0 +1,34 @@
+import pytest
+
+from aws_advanced_python_wrapper.utils.properties import Properties
+from aws_advanced_python_wrapper.host_availability import HostAvailability
+from aws_advanced_python_wrapper.host_selector import WeightedRandomHostSelector
+from aws_advanced_python_wrapper.hostinfo import HostInfo, HostRole
+
+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 = WeightedRandomHostSelector()
+    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),
+        HostInfo(host="some_available_host", role=HOST_ROLE, availability=HostAvailability.AVAILABLE)
+    )
+
+    host_selector = WeightedRandomHostSelector()
+    actual_host = host_selector.get_host(hosts, HOST_ROLE, Properties())
+
+    assert HostAvailability.AVAILABLE == actual_host.get_availability()