Skip to content

feat: adding weighted random host selection strategy #907

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Jul 3, 2025
Merged

feat: adding weighted random host selection strategy #907

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 6 additions & 5 deletions aws_advanced_python_wrapper/connection_provider.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
80 changes: 76 additions & 4 deletions aws_advanced_python_wrapper/host_selector.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -29,6 +30,8 @@
from .utils.messages import Messages
from .utils.properties import Properties, WrapperProperties

logger = Logger(__name__)


class HostSelector(Protocol):
"""
Expand Down Expand Up @@ -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))
Comment on lines +211 to +213
Copy link
Contributor

@karenc-bq karenc-bq Jul 3, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
message = "HostSelector.NoHostsMatchingRole"
logger.error(message, role)
raise AwsWrapperError(Messages.get_formatted("HostSelector.NoHostsMatchingRole", role))
message = Messages.get_formatted("HostSelector.NoHostsMatchingRole", role)
logger.error(message, role)
raise AwsWrapperError(message)

Copy link
Collaborator Author

@JuanLeee JuanLeee Jul 3, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

logger.error already gets the message formatted:

def error(self, msg, *args, **kwargs):
    if not self.logger.isEnabledFor(logging.ERROR):
        return

    if args is not None and len(args) > 0:
        self.logger.error(Messages.get_formatted(msg, *args), **kwargs)
    else:
        try:
            self.logger.error(Messages.get(msg), **kwargs)
        except NotInResourceBundleError:
            self.logger.error(msg, **kwargs)


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))
Comment on lines +246 to +247
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same comment as above


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))
5 changes: 4 additions & 1 deletion aws_advanced_python_wrapper/resources/aws_advanced_python_wrapper_messages.properties
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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...

Expand Down
10 changes: 6 additions & 4 deletions aws_advanced_python_wrapper/sql_alchemy_connection_provider.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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,
Expand Down
9 changes: 9 additions & 0 deletions aws_advanced_python_wrapper/utils/properties.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
3 changes: 3 additions & 0 deletions docs/using-the-python-driver/ReaderSelectionStrategies.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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` |
Expand Down
Loading