From 86db9c6036c5b39920f53929e90a59dded5d0e52 Mon Sep 17 00:00:00 2001 From: karezche <64801825+karenc-bq@users.noreply.github.com> Date: Tue, 9 Jun 2026 18:43:28 -0700 Subject: [PATCH 1/2] chore: update aurora initial connection strategy plugin implementation --- ...rora_initial_connection_strategy_plugin.py | 602 +++++++++++++----- .../cluster_topology_monitor.py | 10 +- .../host_list_provider.py | 8 +- aws_advanced_python_wrapper/plugin_service.py | 4 +- ...dvanced_python_wrapper_messages.properties | 10 +- .../utils/properties.py | 50 ++ .../utils/rds_utils.py | 8 + ...rora_initial_connection_strategy_plugin.py | 315 ++++++++- 8 files changed, 810 insertions(+), 197 deletions(-) diff --git a/aws_advanced_python_wrapper/aurora_initial_connection_strategy_plugin.py b/aws_advanced_python_wrapper/aurora_initial_connection_strategy_plugin.py index 6e7885dd6..59ba8a9aa 100644 --- a/aws_advanced_python_wrapper/aurora_initial_connection_strategy_plugin.py +++ b/aws_advanced_python_wrapper/aurora_initial_connection_strategy_plugin.py @@ -14,12 +14,14 @@ from __future__ import annotations +from enum import Enum from time import perf_counter_ns, sleep -from typing import TYPE_CHECKING, Callable, Optional, Set +from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Set, Tuple if TYPE_CHECKING: from aws_advanced_python_wrapper.driver_dialect import DriverDialect - from aws_advanced_python_wrapper.host_list_provider import HostListProviderService + from aws_advanced_python_wrapper.host_list_provider import \ + HostListProviderService from aws_advanced_python_wrapper.pep249 import Connection from aws_advanced_python_wrapper.plugin_service import PluginService @@ -27,162 +29,468 @@ from aws_advanced_python_wrapper.host_availability import HostAvailability from aws_advanced_python_wrapper.hostinfo import HostInfo, HostRole from aws_advanced_python_wrapper.plugin import Plugin, PluginFactory +from aws_advanced_python_wrapper.utils.log import Logger from aws_advanced_python_wrapper.utils.messages import Messages from aws_advanced_python_wrapper.utils.properties import (Properties, WrapperProperties) from aws_advanced_python_wrapper.utils.rds_url_type import RdsUrlType from aws_advanced_python_wrapper.utils.rds_utils import RdsUtils +logger = Logger(__name__) + + +class InstanceSubstitutionStrategy(Enum): + """Determines which host the plugin should connect to when opening a new connection.""" + SUBSTITUTE_WITH_WRITER = "writer" + SUBSTITUTE_WITH_READER = "reader" + SUBSTITUTE_WITH_ANY = "any" + DO_NOT_SUBSTITUTE = "none" + + @classmethod + def from_property_value(cls, value: Optional[str]) -> Optional[InstanceSubstitutionStrategy]: + if value is None: + return None + + strategy = _SUBSTITUTION_STRATEGY_BY_KEY.get(value.lower()) + if strategy is None: + raise AwsWrapperError(Messages.get_formatted( + "AuroraInitialConnectionStrategyPlugin.InvalidPropertyValue", + WrapperProperties.ENDPOINT_SUBSTITUTION_ROLE.name, + value, + ", ".join(item.value for item in cls))) + return strategy + + def to_target_role(self) -> Optional[HostRole]: + if self is InstanceSubstitutionStrategy.SUBSTITUTE_WITH_WRITER: + return HostRole.WRITER + if self is InstanceSubstitutionStrategy.SUBSTITUTE_WITH_READER: + return HostRole.READER + return None -class AuroraInitialConnectionStrategyPlugin(Plugin): - _SUBSCRIBED_METHODS: Set[str] = {"init_host_provider", "connect"} - _host_list_provider_service: Optional[HostListProviderService] = None +_SUBSTITUTION_STRATEGY_BY_KEY: Dict[str, InstanceSubstitutionStrategy] = { + item.value: item for item in InstanceSubstitutionStrategy +} - @property - def subscribed_methods(self) -> Set[str]: - return AuroraInitialConnectionStrategyPlugin._SUBSCRIBED_METHODS - def __init__(self, plugin_service: PluginService): - super() - self._plugin_service: PluginService = plugin_service - self._rds_utils = RdsUtils() +class RoleVerificationSetting(Enum): + """Determines what role, if any, an opened connection should be verified against.""" + WRITER = "writer" + READER = "reader" + NO_VERIFICATION = "none" - def connect(self, target_driver_func: Callable, driver_dialect: DriverDialect, host_info: HostInfo, props: Properties, - is_initial_connection: bool, connect_func: Callable) -> Connection: - url_type: RdsUrlType = self._rds_utils.identify_rds_type(host_info.host) - if not url_type.is_rds_cluster: - return connect_func() + @classmethod + def from_property_value(cls, value: Optional[str]) -> Optional[RoleVerificationSetting]: + if value is None: + return None - if url_type == RdsUrlType.RDS_WRITER_CLUSTER or url_type == RdsUrlType.RDS_GLOBAL_WRITER_CLUSTER: - writer_candidate_conn: Optional[Connection] = self._get_verified_writer_connection(props, is_initial_connection, connect_func) - if writer_candidate_conn is None: - return connect_func() - return writer_candidate_conn + setting = _VERIFICATION_SETTING_BY_KEY.get(value.lower()) + if setting is None: + raise AwsWrapperError(Messages.get_formatted( + "AuroraInitialConnectionStrategyPlugin.InvalidPropertyValue", + WrapperProperties.VERIFY_OPENED_CONNECTION_ROLE.name, + value, + ", ".join(item.value for item in cls))) + return setting - if url_type == RdsUrlType.RDS_READER_CLUSTER: - reader_candidate_conn: Optional[Connection] = self._get_verified_reader_connection(props, is_initial_connection, connect_func) - if reader_candidate_conn is None: - return connect_func() - return reader_candidate_conn - return connect_func() +_VERIFICATION_SETTING_BY_KEY: Dict[str, RoleVerificationSetting] = { + item.value: item for item in RoleVerificationSetting +} + + +class AuroraInitialConnectionStrategyPlugin(Plugin): + _SUBSCRIBED_METHODS: Set[str] = {"init_host_provider", "connect"} + + def __init__(self, plugin_service: PluginService, props: Properties): + self._plugin_service: PluginService = plugin_service + self._rds_utils = RdsUtils() + self._host_list_provider_service: Optional[HostListProviderService] = None + + self._retry_delay_ms: int = WrapperProperties.OPEN_CONNECTION_RETRY_INTERVAL_MS.get_int(props) + self._open_connection_retry_timeout_ns: int = \ + WrapperProperties.OPEN_CONNECTION_RETRY_TIMEOUT_MS.get_int(props) * 1_000_000 + self._wait_for_initial_topology_ms: int = max( + 0, WrapperProperties.WAIT_FOR_INITIAL_TOPOLOGY_MS.get_int(props)) + + verify_role_value = WrapperProperties.VERIFY_OPENED_CONNECTION_ROLE.get(props) + self._verify_role_prop_value: Optional[str] = \ + verify_role_value.lower() if verify_role_value is not None else None + + # INITIAL_CONNECTION_HOST_SELECTOR_STRATEGY overrides the deprecated + # READER_INITIAL_HOST_SELECTOR_STRATEGY when it is explicitly set. + if WrapperProperties.INITIAL_CONNECTION_HOST_SELECTOR_STRATEGY.name in props: + self._selection_strategy: Optional[str] = \ + WrapperProperties.INITIAL_CONNECTION_HOST_SELECTOR_STRATEGY.get(props) + else: + self._selection_strategy = WrapperProperties.READER_INITIAL_HOST_SELECTOR_STRATEGY.get(props) - def _get_verified_writer_connection(self, props: Properties, is_initial_connection: bool, connect_func: Callable) -> Connection | None: - retry_delay_ms: int = WrapperProperties.OPEN_CONNECTION_RETRY_INTERVAL_MS.get_int(props) - end_time_nano = perf_counter_ns() + (WrapperProperties.OPEN_CONNECTION_RETRY_TIMEOUT_MS.get_int(props) * 1000000) + @property + def subscribed_methods(self) -> Set[str]: + return AuroraInitialConnectionStrategyPlugin._SUBSCRIBED_METHODS - writer_candidate_conn: Optional[Connection] - writer_candidate: Optional[HostInfo] + def init_host_provider( + self, + props: Properties, + host_list_provider_service: HostListProviderService, + init_host_provider_func: Callable): + self._host_list_provider_service = host_list_provider_service + init_host_provider_func() - while perf_counter_ns() < end_time_nano: - writer_candidate_conn = None + def connect( + self, + target_driver_func: Callable, + driver_dialect: DriverDialect, + host_info: HostInfo, + props: Properties, + is_initial_connection: bool, + connect_func: Callable) -> Connection: + original_host = host_info.host + url_type: RdsUrlType = self._rds_utils.identify_rds_type(original_host) + substitution_strategy = self._get_instance_substitution_strategy( + props, url_type, is_initial_connection, original_host) + role_to_verify = self._get_role_to_verify(url_type, is_initial_connection, props, original_host) + end_time_ns = perf_counter_ns() + self._open_connection_retry_timeout_ns + + while perf_counter_ns() < end_time_ns: + candidate_conn: Optional[Connection] = None + candidate_host: Optional[HostInfo] = None try: - writer_candidate = self._get_writer() - if writer_candidate is None or self._rds_utils.is_rds_cluster_dns(writer_candidate.host): - # Writer is not found. Topology is outdated. - writer_candidate_conn = connect_func() - self._plugin_service.force_refresh_host_list(writer_candidate_conn) - writer_candidate = self._plugin_service.identify_connection(writer_candidate_conn) + candidate_host, candidate_conn = self._open_candidate_connection( + host_info, url_type, substitution_strategy, props, connect_func) - if writer_candidate is None or writer_candidate.role != HostRole.WRITER: - self._close_connection(writer_candidate_conn) - self._delay(retry_delay_ms) - continue + if candidate_conn is None: + # _open_candidate_connection always returns a connection on success; if none is + # present the attempt did not yield a usable connection, so retry until the + # timeout is reached. + continue - if is_initial_connection and self._host_list_provider_service is not None: - self._host_list_provider_service.initial_connection_host_info = writer_candidate + if role_to_verify is None: + # No verification required. + self._set_initial_connection_host_info(is_initial_connection, candidate_host) + return candidate_conn + + conn_role = self._plugin_service.get_host_role(candidate_conn) + if conn_role == role_to_verify: + # Verification succeeded. + self._set_initial_connection_host_info(is_initial_connection, candidate_host) + return candidate_conn + + # Verification failed. Retry, unless a reader was requested but the cluster has no readers. + self._plugin_service.force_refresh_host_list(candidate_conn) + if role_to_verify == HostRole.READER and self._has_hosts() and not self._has_readers(): + # A reader was requested but the cluster has no readers. + # Simulate the reader cluster endpoint logic and return the current (writer) connection. + if self._verify_role_prop_value == RoleVerificationSetting.READER.value: + logger.debug( + "AuroraInitialConnectionStrategyPlugin.VerifyReaderConfiguredButNoReadersExist", + WrapperProperties.VERIFY_OPENED_CONNECTION_ROLE.name) + self._set_initial_connection_host_info(is_initial_connection, candidate_host) + return candidate_conn + + logger.debug( + "AuroraInitialConnectionStrategyPlugin.IncorrectRole", candidate_host.host, role_to_verify) + self._close_connection(candidate_conn) + self._delay(self._retry_delay_ms) + except Exception as e: + self._close_connection(candidate_conn) + if self._plugin_service.is_login_exception(e): + raise - return writer_candidate_conn + if candidate_host is not None: + self._plugin_service.set_availability( + candidate_host.as_aliases(), HostAvailability.UNAVAILABLE) - writer_candidate_conn = self._plugin_service.connect(writer_candidate, props) + if self._plugin_service.is_network_exception(e): + # Retry connection. + continue - if self._plugin_service.get_host_role(writer_candidate_conn) != HostRole.WRITER: - self._plugin_service.force_refresh_host_list(writer_candidate_conn) - self._close_connection(writer_candidate_conn) - self._delay(retry_delay_ms) + if (self._plugin_service.is_read_only_connection_exception(e) + and (role_to_verify == HostRole.WRITER + or substitution_strategy is InstanceSubstitutionStrategy.SUBSTITUTE_WITH_WRITER)): + # Retry connection. continue - # Writer connection is valid and verified. - if is_initial_connection and self._host_list_provider_service is not None: - self._host_list_provider_service.initial_connection_host_info = writer_candidate - return writer_candidate_conn - except Exception as e: - self._close_connection(writer_candidate_conn) - raise e + raise + + raise AwsWrapperError(Messages.get_formatted( + "AuroraInitialConnectionStrategyPlugin.Timeout", + self._open_connection_retry_timeout_ns // 1_000_000, + WrapperProperties.VERIFY_OPENED_CONNECTION_ROLE.name)) + + def _open_candidate_connection( + self, + original_connect_host: HostInfo, + url_type: RdsUrlType, + substitution_strategy: InstanceSubstitutionStrategy, + props: Properties, + connect_func: Callable) -> Tuple[HostInfo, Optional[Connection]]: + """Opens a candidate connection, returning the host that was connected to and the connection. + + If no substitution is needed, the original endpoint is used. Otherwise, an instance host is + selected from the topology when available; if the topology isn't available yet, a connection is + opened via the initial endpoint (which also confirms the dialect and acts as a fallback) and, + when ``wait_for_initial_topology_ms > 0``, the topology fetch is awaited before re-attempting + instance selection. + """ + if substitution_strategy is InstanceSubstitutionStrategy.DO_NOT_SUBSTITUTE: + return original_connect_host, connect_func() + + candidate_host = self._get_candidate_host(original_connect_host, url_type, substitution_strategy) + if candidate_host is not None and self._rds_utils.is_rds_instance(candidate_host.host): + # Topology is already available; connect to the selected instance. + return candidate_host, self._plugin_service.connect(candidate_host, props, self) + + # Unable to find an instance URL host. Topology may not exist yet, or may be outdated. + # Connect via the initial endpoint. This connection also confirms the dialect (done by + # the default plugin on the initial connection), which is a prerequisite for fetching the + # topology, and it serves as a fallback connection if instance selection or connection fails. + candidate_conn = connect_func() + + if self._wait_for_initial_topology_ms <= 0: + # Feature disabled. Preserve the previous behavior. + self._plugin_service.force_refresh_host_list(candidate_conn) + return original_connect_host, candidate_conn + + return self._wait_for_topology_and_connect_to_instance( + original_connect_host, url_type, substitution_strategy, props, candidate_conn) + + def _wait_for_topology_and_connect_to_instance( + self, + original_connect_host: HostInfo, + url_type: RdsUrlType, + substitution_strategy: InstanceSubstitutionStrategy, + props: Properties, + fallback_conn: Connection) -> Tuple[HostInfo, Optional[Connection]]: + """Blocks until the topology for this cluster has been fetched, then re-attempts instance + selection and connection. This serializes concurrent/prefill connections (all waiting on the + same per-cluster topology monitor) so that the configured host selection strategy can + distribute them across instances instead of all relying on the initial endpoint resolved via + DNS. + + Returns the selected instance host and its connection if the topology was fetched and the + instance connection succeeded. Otherwise returns ``original_connect_host`` and the + already-opened initial-endpoint connection, which is kept as a fallback. + """ + logger.debug( + "AuroraInitialConnectionStrategyPlugin.WaitingForTopology", + self._wait_for_initial_topology_ms, original_connect_host.host) + + # Deviation from JDBC: force_monitoring_refresh_host_list takes seconds, and host list + # providers without monitor support raise instead of returning their host list. + timeout_sec = self._wait_for_initial_topology_ms / 1000 + try: + topology_fetched = self._plugin_service.force_monitoring_refresh_host_list(True, timeout_sec) + except Exception: + topology_fetched = False + + if not topology_fetched: + logger.debug( + "AuroraInitialConnectionStrategyPlugin.WaitForTopologyTimeout", + self._wait_for_initial_topology_ms, original_connect_host.host) + return original_connect_host, fallback_conn + + instance_host = self._get_candidate_host(original_connect_host, url_type, substitution_strategy) + if instance_host is None or not self._rds_utils.is_rds_instance(instance_host.host): + return original_connect_host, fallback_conn + + try: + instance_conn = self._plugin_service.connect(instance_host, props, self) + except Exception: + # Failed to connect to the selected instance; keep the initial-endpoint connection. + logger.debug( + "AuroraInitialConnectionStrategyPlugin.FailedToConnectToSelectedInstance", instance_host.host) + return original_connect_host, fallback_conn + + # Close the previous (fallback) connection once the instance connection is held. + self._close_connection(fallback_conn) + return instance_host, instance_conn + + def _get_instance_substitution_strategy( + self, + props: Properties, + url_type: RdsUrlType, + is_initial_connection: bool, + original_host: str) -> InstanceSubstitutionStrategy: + if is_initial_connection: + strategy = InstanceSubstitutionStrategy.from_property_value( + WrapperProperties.ENDPOINT_SUBSTITUTION_ROLE.get(props)) + if strategy is not None: + self._validate_substitution_strategy(strategy, url_type) + return strategy + + # This is not an initial connection, or ENDPOINT_SUBSTITUTION_ROLE was not set. + # Pick a strategy according to the default behavior. + if url_type == RdsUrlType.RDS_GLOBAL_WRITER_CLUSTER: + return InstanceSubstitutionStrategy.SUBSTITUTE_WITH_WRITER + + if url_type == RdsUrlType.RDS_WRITER_CLUSTER: + writer = self._get_writer() + if writer is None or not self._rds_utils.is_rds_instance(writer.host): + return InstanceSubstitutionStrategy.DO_NOT_SUBSTITUTE + + if self._rds_utils.is_same_region(writer.host, original_host): + return InstanceSubstitutionStrategy.SUBSTITUTE_WITH_WRITER + + # The cluster writer endpoint belongs to a different region than the current writer region. + # This means the cluster is an Aurora Global Database and the cluster writer endpoint is in a + # secondary region. In this case the cluster writer endpoint is inactive and doesn't represent + # the current writer. A user setting decides whether to substitute it with a writer instance URL. + inactive_strategy = InstanceSubstitutionStrategy.from_property_value( + WrapperProperties.INACTIVE_CLUSTER_WRITER_SUBSTITUTION_ROLE.get(props)) + return inactive_strategy if inactive_strategy is not None \ + else InstanceSubstitutionStrategy.SUBSTITUTE_WITH_WRITER - return None + if url_type == RdsUrlType.RDS_READER_CLUSTER: + return InstanceSubstitutionStrategy.SUBSTITUTE_WITH_READER - def _get_verified_reader_connection(self, props: Properties, is_initial_connection: bool, connect_func: Callable) -> Optional[Connection]: - retry_delay_ms: int = WrapperProperties.OPEN_CONNECTION_RETRY_INTERVAL_MS.get_int(props) - end_time_nano = perf_counter_ns() + (WrapperProperties.OPEN_CONNECTION_RETRY_TIMEOUT_MS.get_int(props) * 1000000) + return InstanceSubstitutionStrategy.DO_NOT_SUBSTITUTE - reader_candidate_conn: Optional[Connection] - reader_candidate: Optional[HostInfo] + def _validate_substitution_strategy( + self, setting: InstanceSubstitutionStrategy, url_type: RdsUrlType): + if setting is InstanceSubstitutionStrategy.DO_NOT_SUBSTITUTE: + return - while perf_counter_ns() < end_time_nano: - reader_candidate_conn = None - reader_candidate = None + if url_type == RdsUrlType.RDS_INSTANCE: + raise AwsWrapperError(Messages.get_formatted( + "AuroraInitialConnectionStrategyPlugin.InvalidSettingForInstanceEndpoint", + WrapperProperties.ENDPOINT_SUBSTITUTION_ROLE.name)) - try: - reader_candidate = self._get_reader(props) - if reader_candidate is None or self._rds_utils.is_rds_cluster_dns(reader_candidate.host): - # READER is not found. Topology is outdated. - reader_candidate_conn = connect_func() - self._plugin_service.force_refresh_host_list(reader_candidate_conn) - reader_candidate = self._plugin_service.identify_connection(reader_candidate_conn) - - if reader_candidate is None: - self._close_connection(reader_candidate_conn) - self._delay(retry_delay_ms) - continue - - if reader_candidate is not None and reader_candidate.role != HostRole.READER: - if self._has_no_readers(): - # Cluster has no readers. Simulate Aurora reader cluster endpoint logic and return the current writer connection. - if is_initial_connection and self._host_list_provider_service is not None: - self._host_list_provider_service.initial_connection_host_info = reader_candidate - return reader_candidate_conn - - self._close_connection(reader_candidate_conn) - self._delay(retry_delay_ms) - continue - - if is_initial_connection and self._host_list_provider_service is not None: - self._host_list_provider_service.initial_connection_host_info = reader_candidate - return reader_candidate_conn - - reader_candidate_conn = self._plugin_service.connect(reader_candidate, props) - - if self._plugin_service.get_host_role(reader_candidate_conn) != HostRole.READER: - # If the new connection resolves to a writer instance, the topology is outdated. - # Force refresh to update the topology. - self._plugin_service.force_refresh_host_list(reader_candidate_conn) - - if self._has_no_readers(): - # Cluster has no readers. Simulate Aurora reader cluster endpoint logic and return the current writer connection. - if is_initial_connection and self._host_list_provider_service is not None: - self._host_list_provider_service.initial_connection_host_info = reader_candidate - return reader_candidate_conn - - self._close_connection(reader_candidate_conn) - self._delay(retry_delay_ms) - continue + if not url_type.is_rds_cluster: + return + + # A custom cluster can only be of type "reader" or "any", so SUBSTITUTE_WITH_WRITER is not allowed. + if (setting is InstanceSubstitutionStrategy.SUBSTITUTE_WITH_WRITER + and url_type in (RdsUrlType.RDS_READER_CLUSTER, RdsUrlType.RDS_CUSTOM_CLUSTER)): + raise AwsWrapperError(Messages.get_formatted( + "AuroraInitialConnectionStrategyPlugin.InvalidSettingForEndpoint", + WrapperProperties.ENDPOINT_SUBSTITUTION_ROLE.name, "writer", "reader cluster or custom cluster")) + + if (setting is InstanceSubstitutionStrategy.SUBSTITUTE_WITH_READER + and url_type in (RdsUrlType.RDS_WRITER_CLUSTER, RdsUrlType.RDS_GLOBAL_WRITER_CLUSTER)): + raise AwsWrapperError(Messages.get_formatted( + "AuroraInitialConnectionStrategyPlugin.InvalidSettingForEndpoint", + WrapperProperties.ENDPOINT_SUBSTITUTION_ROLE.name, "reader", "writer cluster or global cluster")) + + if (setting is InstanceSubstitutionStrategy.SUBSTITUTE_WITH_ANY + and url_type != RdsUrlType.RDS_CUSTOM_CLUSTER): + raise AwsWrapperError(Messages.get_formatted( + "AuroraInitialConnectionStrategyPlugin.InvalidSettingForEndpoint", + WrapperProperties.ENDPOINT_SUBSTITUTION_ROLE.name, "any", + "writer cluster, reader cluster, or global cluster")) + + def _get_role_to_verify( + self, + url_type: RdsUrlType, + is_initial_connection: bool, + props: Properties, + original_host: str) -> Optional[HostRole]: + if not is_initial_connection: + return None + + setting = RoleVerificationSetting.from_property_value(self._verify_role_prop_value) + if setting is not None: + self._validate_verification_setting(setting, url_type) + + if setting is RoleVerificationSetting.NO_VERIFICATION: + return None + if setting is RoleVerificationSetting.WRITER: + return HostRole.WRITER + if setting is RoleVerificationSetting.READER: + return HostRole.READER + + # Role verification setting is not set. We still verify the correct role for a writer/reader cluster. + if url_type == RdsUrlType.RDS_GLOBAL_WRITER_CLUSTER: + return HostRole.WRITER + + if url_type == RdsUrlType.RDS_WRITER_CLUSTER: + writer = self._get_writer() + if (writer is not None and self._rds_utils.is_rds_instance(writer.host) + and self._rds_utils.is_same_region(writer.host, original_host)): + # The cluster writer endpoint belongs to the same region as the current writer; it's active. + return HostRole.WRITER + + # Writer is not found (topology cache may not be available yet) or the cluster writer endpoint + # belongs to a different region. In either case, assume the cluster writer endpoint may be + # inactive and use the corresponding setting. + inactive_strategy = InstanceSubstitutionStrategy.from_property_value( + WrapperProperties.VERIFY_INACTIVE_CLUSTER_WRITER_CONNECTION_ROLE.get(props)) + return inactive_strategy.to_target_role() if inactive_strategy is not None else HostRole.WRITER - # Reader connection is valid and verified. - if is_initial_connection and self._host_list_provider_service is not None: - self._host_list_provider_service.initial_connection_host_info = reader_candidate - return reader_candidate_conn - except Exception as e: - self._close_connection(reader_candidate_conn) - if not self._plugin_service.is_login_exception(e) and reader_candidate is not None: - self._plugin_service.set_availability(reader_candidate.as_aliases(), HostAvailability.UNAVAILABLE) + if url_type == RdsUrlType.RDS_READER_CLUSTER: + return HostRole.READER + + return None - raise e + def _validate_verification_setting(self, setting: RoleVerificationSetting, url_type: RdsUrlType): + if (setting is RoleVerificationSetting.READER + and url_type in (RdsUrlType.RDS_WRITER_CLUSTER, RdsUrlType.RDS_GLOBAL_WRITER_CLUSTER)): + raise AwsWrapperError(Messages.get_formatted( + "AuroraInitialConnectionStrategyPlugin.InvalidSettingForEndpoint", + WrapperProperties.VERIFY_OPENED_CONNECTION_ROLE.name, "reader", "writer cluster or global cluster")) + + # A custom cluster can only be of type "reader" or "any". + if (setting is RoleVerificationSetting.WRITER + and url_type in (RdsUrlType.RDS_READER_CLUSTER, RdsUrlType.RDS_CUSTOM_CLUSTER)): + raise AwsWrapperError(Messages.get_formatted( + "AuroraInitialConnectionStrategyPlugin.InvalidSettingForEndpoint", + WrapperProperties.VERIFY_OPENED_CONNECTION_ROLE.name, "writer", "reader cluster or custom cluster")) + + def _get_candidate_host( + self, + original_connect_host: HostInfo, + url_type: RdsUrlType, + substitution_strategy: InstanceSubstitutionStrategy) -> Optional[HostInfo]: + if substitution_strategy is InstanceSubstitutionStrategy.DO_NOT_SUBSTITUTE: + return original_connect_host + + if substitution_strategy is InstanceSubstitutionStrategy.SUBSTITUTE_WITH_WRITER: + return self._get_writer() + + # SUBSTITUTE_WITH_ANY has no specific target role, so to_target_role() returns None + target_role = substitution_strategy.to_target_role() + if (target_role is None + or self._selection_strategy is None + or not self._plugin_service.accepts_strategy(target_role, self._selection_strategy)): + raise AwsWrapperError(Messages.get_formatted( + "AuroraInitialConnectionStrategyPlugin.UnsupportedStrategy", self._selection_strategy)) + + try: + aws_region = self._rds_utils.get_rds_region(original_connect_host.host) \ + if url_type.has_region else None + if aws_region: + hosts_in_region: List[HostInfo] = [ + host for host in self._plugin_service.hosts + if (host_region := self._rds_utils.get_rds_region(host.host)) is not None + and aws_region.casefold() == host_region.casefold()] + return self._plugin_service.get_host_info_by_strategy( + target_role, self._selection_strategy, hosts_in_region) + + return self._plugin_service.get_host_info_by_strategy(target_role, self._selection_strategy) + except Exception: + # Unable to find a candidate host. + return None + + def _set_initial_connection_host_info( + self, is_initial_connection: bool, host_info: Optional[HostInfo]): + if (is_initial_connection + and self._host_list_provider_service is not None + and host_info is not None): + self._host_list_provider_service.initial_connection_host_info = host_info + def _get_writer(self) -> Optional[HostInfo]: + for host in self._plugin_service.all_hosts: + if host.role == HostRole.WRITER: + return host return None + def _has_hosts(self) -> bool: + return len(self._plugin_service.all_hosts) > 0 + + def _has_readers(self) -> bool: + return any(host.role == HostRole.READER for host in self._plugin_service.all_hosts) + def _close_connection(self, connection: Optional[Connection]): if connection is not None: try: @@ -194,56 +502,8 @@ def _close_connection(self, connection: Optional[Connection]): def _delay(self, delay_ms: int): sleep(delay_ms / 1000) - def _get_writer(self) -> Optional[HostInfo]: - for host in self._plugin_service.all_hosts: - if host.role == HostRole.WRITER: - return host - - return None - - def _get_reader(self, props: Properties) -> Optional[HostInfo]: - strategy = WrapperProperties.READER_INITIAL_HOST_SELECTOR_STRATEGY.get(props) - if (self._plugin_service is not None - and strategy is not None - and self._plugin_service.accepts_strategy(HostRole.READER, strategy)): - try: - original_host = self._plugin_service.current_host_info - url_type = self._rds_utils.identify_rds_type(original_host.host) if original_host else None - - if url_type and url_type.has_region: - aws_region = self._rds_utils.get_rds_region(original_host.host) - if aws_region: - hosts_in_region = [] - for h in self._plugin_service.all_hosts: - h_region = self._rds_utils.get_rds_region(h.host) - if h_region and aws_region.lower() == h_region.lower(): - hosts_in_region.append(h) - return self._plugin_service.get_host_info_by_strategy( - HostRole.READER, strategy, hosts_in_region) - - return self._plugin_service.get_host_info_by_strategy(HostRole.READER, strategy) - except Exception: - # Host isn't found. - return None - - raise AwsWrapperError(Messages.get_formatted("AuroraInitialConnectionStrategyPlugin.UnsupportedStrategy", strategy)) - - def init_host_provider(self, props: Properties, host_list_provider_service: HostListProviderService, init_host_provider_func: Callable): - self._host_list_provider_service = host_list_provider_service - init_host_provider_func() - - def _has_no_readers(self) -> bool: - if len(self._plugin_service.all_hosts) == 0: - return False - - for host in self._plugin_service.all_hosts: - if host.role == HostRole.READER: - return False - - return True - class AuroraInitialConnectionStrategyPluginFactory(PluginFactory): @staticmethod def get_instance(plugin_service: PluginService, props: Properties) -> Plugin: - return AuroraInitialConnectionStrategyPlugin(plugin_service) + return AuroraInitialConnectionStrategyPlugin(plugin_service, props) diff --git a/aws_advanced_python_wrapper/cluster_topology_monitor.py b/aws_advanced_python_wrapper/cluster_topology_monitor.py index 5321de5df..fabda15d9 100644 --- a/aws_advanced_python_wrapper/cluster_topology_monitor.py +++ b/aws_advanced_python_wrapper/cluster_topology_monitor.py @@ -52,11 +52,11 @@ class ClusterTopologyMonitor(ABC): @abstractmethod - def force_refresh(self, should_verify_writer: bool, timeout_sec: int) -> Topology: + def force_refresh(self, should_verify_writer: bool, timeout_sec: float) -> Topology: pass @abstractmethod - def force_refresh_with_connection(self, connection: Connection, timeout_sec: int) -> Topology: + def force_refresh_with_connection(self, connection: Connection, timeout_sec: float) -> Topology: pass @property @@ -133,7 +133,7 @@ def __init__(self, plugin_service: PluginService, topology_utils: TopologyUtils, self._start_monitoring() - def force_refresh(self, should_verify_writer: bool, timeout_sec: int) -> Topology: + def force_refresh(self, should_verify_writer: bool, timeout_sec: float) -> Topology: current_time_nano = time.time_ns() if (self._ignore_new_topology_requests_end_time_nano > 0 and current_time_nano < self._ignore_new_topology_requests_end_time_nano): @@ -149,12 +149,12 @@ def force_refresh(self, should_verify_writer: bool, timeout_sec: int) -> Topolog result = self._wait_till_topology_gets_updated(timeout_sec) return result - def force_refresh_with_connection(self, connection: Connection, timeout_sec: int) -> Topology: + def force_refresh_with_connection(self, connection: Connection, timeout_sec: float) -> Topology: if self._is_verified_writer_connection: return self._wait_till_topology_gets_updated(timeout_sec) return self._fetch_topology_and_update_cache(connection) - def _wait_till_topology_gets_updated(self, timeout_sec: int) -> Topology: + def _wait_till_topology_gets_updated(self, timeout_sec: float) -> Topology: current_hosts = self._get_stored_hosts() self._request_to_update_topology.set() diff --git a/aws_advanced_python_wrapper/host_list_provider.py b/aws_advanced_python_wrapper/host_list_provider.py index 06cf26c21..8d6164ca5 100644 --- a/aws_advanced_python_wrapper/host_list_provider.py +++ b/aws_advanced_python_wrapper/host_list_provider.py @@ -73,7 +73,7 @@ def get_current_topology(self, connection: Connection, initial_host_info: HostIn """ ... - def force_monitoring_refresh(self, should_verify_writer: bool, timeout_sec: int) -> Topology: + def force_monitoring_refresh(self, should_verify_writer: bool, timeout_sec: float) -> Topology: ... def get_cluster_id(self) -> str: @@ -248,7 +248,7 @@ def _get_or_create_monitor(self) -> Optional[ClusterTopologyMonitor]: ) ) - def _force_refresh_monitor(self, should_verify_writer: bool, timeout_sec: int) -> Optional[Topology]: + def _force_refresh_monitor(self, should_verify_writer: bool, timeout_sec: float) -> Optional[Topology]: """Force refresh using monitor - matches Java's forceRefreshMonitor""" monitor = self._get_or_create_monitor() if monitor is None: @@ -275,7 +275,7 @@ def get_current_topology(self, connection: Connection, initial_host_info: HostIn return hosts return () - def force_monitoring_refresh(self, should_verify_writer: bool, timeout_sec: int) -> Topology: + def force_monitoring_refresh(self, should_verify_writer: bool, timeout_sec: float) -> Topology: """Public API for forcing monitor refresh""" self._initialize() hosts = self._force_refresh_monitor(should_verify_writer, timeout_sec) @@ -363,7 +363,7 @@ def get_current_topology(self, connection: Connection, initial_host_info: HostIn self._initialize() return tuple(self._hosts) - def force_monitoring_refresh(self, should_verify_writer: bool, timeout_sec: int) -> Topology: + def force_monitoring_refresh(self, should_verify_writer: bool, timeout_sec: float) -> Topology: raise AwsWrapperError( Messages.get_formatted("HostListProvider.ForceMonitoringRefreshUnsupported", "ConnectionStringHostListProvider")) diff --git a/aws_advanced_python_wrapper/plugin_service.py b/aws_advanced_python_wrapper/plugin_service.py index 76eae5d28..5a2c03b50 100644 --- a/aws_advanced_python_wrapper/plugin_service.py +++ b/aws_advanced_python_wrapper/plugin_service.py @@ -263,7 +263,7 @@ def refresh_host_list(self, connection: Optional[Connection] = None): def force_refresh_host_list(self, connection: Optional[Connection] = None): ... - def force_monitoring_refresh_host_list(self, should_verify_writer: bool, timeout_ms: int) -> bool: + def force_monitoring_refresh_host_list(self, should_verify_writer: bool, timeout_sec: float) -> bool: ... def connect(self, host_info: HostInfo, props: Properties, plugin_to_skip: Optional[Plugin] = None) -> Connection: @@ -594,7 +594,7 @@ def force_refresh_host_list(self, connection: Optional[Connection] = None): self._update_host_availability(updated_host_list) self._update_hosts(updated_host_list) - def force_monitoring_refresh_host_list(self, should_verify_writer: bool, timeout_sec: int) -> bool: + def force_monitoring_refresh_host_list(self, should_verify_writer: bool, timeout_sec: float) -> bool: try: updated_host_list = self.host_list_provider.force_monitoring_refresh(should_verify_writer, timeout_sec) if updated_host_list is not None: 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 index ee8d6e3de..b01d488fb 100644 --- a/aws_advanced_python_wrapper/resources/aws_advanced_python_wrapper_messages.properties +++ b/aws_advanced_python_wrapper/resources/aws_advanced_python_wrapper_messages.properties @@ -17,8 +17,16 @@ AuroraPgDialect.AuroraUtils=[AuroraPgDialect] aurora_utils: {} AuroraPgDialect.HasTopologyTrue=[AuroraPgDialect] has_topology: True -AuroraInitialConnectionStrategyPlugin.RequireDynamicProvider=[AuroraInitialConnectionStrategyPlugin] Dynamic host list provider is required. AuroraInitialConnectionStrategyPlugin.UnsupportedStrategy=[AuroraInitialConnectionStrategyPlugin] Unsupported host selection strategy '{}'. +AuroraInitialConnectionStrategyPlugin.InvalidSettingForEndpoint=[AuroraInitialConnectionStrategyPlugin] Parameter '{}' cannot be set to '{}' when using a {} endpoint. Please see the Aurora Initial Connection Strategy Plugin documentation for valid setting/endpoint combinations. +AuroraInitialConnectionStrategyPlugin.InvalidSettingForInstanceEndpoint=[AuroraInitialConnectionStrategyPlugin] Parameter '{}' cannot be set when using an instance endpoint. Please see the Aurora Initial Connection Strategy Plugin documentation for valid setting/endpoint combinations. +AuroraInitialConnectionStrategyPlugin.InvalidPropertyValue=[AuroraInitialConnectionStrategyPlugin] Received an invalid value for parameter '{}'. Received '{}', valid values are {}. +AuroraInitialConnectionStrategyPlugin.Timeout=[AuroraInitialConnectionStrategyPlugin] The Aurora Initial Connection Strategy Plugin attempted to connect but timed out after {}ms. Please ensure that your URL is correct, there are no network issues, and you are connecting to the correct role if '{}' was set. +AuroraInitialConnectionStrategyPlugin.IncorrectRole=[AuroraInitialConnectionStrategyPlugin] The connection opened to '{}' did not have the expected role '{}'. Retrying. +AuroraInitialConnectionStrategyPlugin.VerifyReaderConfiguredButNoReadersExist=[AuroraInitialConnectionStrategyPlugin] Parameter '{}' was set to 'reader' but no readers were detected in the topology. The writer will be used as a fallback. +AuroraInitialConnectionStrategyPlugin.WaitingForTopology=[AuroraInitialConnectionStrategyPlugin] Waiting up to {}ms for the cluster topology of '{}' to be fetched before opening a new connection. +AuroraInitialConnectionStrategyPlugin.WaitForTopologyTimeout=[AuroraInitialConnectionStrategyPlugin] Timed out after {}ms while waiting for the cluster topology of '{}' to be fetched. Falling back to connecting via the provided endpoint. +AuroraInitialConnectionStrategyPlugin.FailedToConnectToSelectedInstance=[AuroraInitialConnectionStrategyPlugin] Failed to connect to the selected instance '{}'. Falling back to the connection opened via the provided endpoint. AdfsCredentialsProviderFactory.FailedLogin=[AdfsCredentialsProviderFactory] Failed login. Could not obtain SAML Assertion from ADFS SignOn Page POST response: '{}' AdfsCredentialsProviderFactory.GetSamlAssertionFailed=[AdfsCredentialsProviderFactory] Failed to get SAML Assertion due to exception: '{}' diff --git a/aws_advanced_python_wrapper/utils/properties.py b/aws_advanced_python_wrapper/utils/properties.py index 4401f7ca0..d21e45b83 100644 --- a/aws_advanced_python_wrapper/utils/properties.py +++ b/aws_advanced_python_wrapper/utils/properties.py @@ -585,12 +585,49 @@ class WrapperProperties: False, ) + # Deprecated. Use INITIAL_CONNECTION_HOST_SELECTOR_STRATEGY instead. READER_INITIAL_HOST_SELECTOR_STRATEGY = WrapperProperty( "reader_initial_connection_host_selector_strategy", "The strategy that should be used to select a new reader host while opening a new connection.", "random", ) + INITIAL_CONNECTION_HOST_SELECTOR_STRATEGY = WrapperProperty( + "initial_connection_host_selector_strategy", + "The strategy that should be used to select a host while opening a new connection.", + "random", + ) + + ENDPOINT_SUBSTITUTION_ROLE = WrapperProperty( + "endpoint_substitution_role", + "Defines whether or not the initial connection URL should be replaced with an instance URL from the " + "topology info when available, and if so, the role of the instance URL that should be selected. " + "Valid values are 'writer', 'reader', 'any', or 'none'.", + None, + ) + + INACTIVE_CLUSTER_WRITER_SUBSTITUTION_ROLE = WrapperProperty( + "inactive_cluster_writer_endpoint_substitution_role", + "Defines whether or not the inactive cluster writer endpoint in the initial connection URL should " + "be replaced with a writer instance URL from the topology info when available. " + "Valid values are 'writer' or 'none'.", + "writer", + ) + + VERIFY_OPENED_CONNECTION_ROLE = WrapperProperty( + "verify_opened_connection_type", + "Defines whether an opened connection should be verified to be a writer or reader, " + "or if no role verification should be performed. Valid values are 'writer', 'reader', or 'none'.", + None, + ) + + VERIFY_INACTIVE_CLUSTER_WRITER_CONNECTION_ROLE = WrapperProperty( + "verify_inactive_cluster_writer_endpoint_connection_type", + "Defines whether inactive cluster writer connection should be verified to be a writer, " + "or if no role verification should be performed. Valid values are 'writer' or 'none'.", + "writer", + ) + OPEN_CONNECTION_RETRY_TIMEOUT_MS = WrapperProperty( "open_connection_retry_timeout_ms", "Maximum allowed time for the retries opening a connection.", @@ -603,6 +640,19 @@ class WrapperProperties: 1000, ) + WAIT_FOR_INITIAL_TOPOLOGY_MS = WrapperProperty( + "wait_for_initial_topology_ms", + "Maximum allowed time, in milliseconds, to wait for the cluster topology to be fetched before opening a new " + "connection. When set to a value greater than 0 and the topology is not yet available, the plugin will block " + "until the topology has been discovered (or this timeout is reached) instead of falling back to connecting via " + "the initial endpoint in the connection string. This ensures host selection strategies such as 'round_robin' " + "distribute concurrent and connection-pool prefill connections across instances rather than routing them all " + "to a single instance resolved through DNS. The wait is scoped to the cluster the connection belongs to; " + "connections to other clusters are not affected. When set to 0 (the default) the previous behavior is " + "preserved.", + 0, + ) + # Simple Read/Write Splitting SRW_READ_ENDPOINT = WrapperProperty( "srw_read_endpoint", diff --git a/aws_advanced_python_wrapper/utils/rds_utils.py b/aws_advanced_python_wrapper/utils/rds_utils.py index ae914ad5e..d67727c91 100644 --- a/aws_advanced_python_wrapper/utils/rds_utils.py +++ b/aws_advanced_python_wrapper/utils/rds_utils.py @@ -202,6 +202,14 @@ def get_rds_region(self, host: Optional[str]): return elb_matcher.group(RdsUtils.REGION_GROUP) return None + def is_same_region(self, host1: Optional[str], host2: Optional[str]) -> bool: + if not host1 or not host1.strip() or not host2 or not host2.strip(): + return False + + host1_region = self.get_rds_region(host1) + host2_region = self.get_rds_region(host2) + return host1_region is not None and host2_region is not None and host1_region.casefold() == host2_region.casefold() + def is_writer_cluster_dns(self, host: str) -> bool: dns_group = self._get_dns_group(self._get_prepared_host(host)) return dns_group is not None and dns_group.casefold() == "cluster-" diff --git a/tests/unit/test_aurora_initial_connection_strategy_plugin.py b/tests/unit/test_aurora_initial_connection_strategy_plugin.py index 6e361e76c..08bf92e3a 100644 --- a/tests/unit/test_aurora_initial_connection_strategy_plugin.py +++ b/tests/unit/test_aurora_initial_connection_strategy_plugin.py @@ -14,29 +14,316 @@ from unittest.mock import MagicMock -from aws_advanced_python_wrapper.aurora_initial_connection_strategy_plugin import \ - AuroraInitialConnectionStrategyPlugin +import pytest + +from aws_advanced_python_wrapper.aurora_initial_connection_strategy_plugin import ( + AuroraInitialConnectionStrategyPlugin, InstanceSubstitutionStrategy) +from aws_advanced_python_wrapper.errors import AwsWrapperError +from aws_advanced_python_wrapper.hostinfo import HostInfo, HostRole from aws_advanced_python_wrapper.utils.properties import (Properties, WrapperProperties) +from aws_advanced_python_wrapper.utils.rds_url_type import RdsUrlType +WRITER_INSTANCE = "instance-1.xyz.us-east-1.rds.amazonaws.com" +READER_INSTANCE = "instance-2.xyz.us-east-1.rds.amazonaws.com" +WRITER_CLUSTER = "mycluster.cluster-xyz.us-east-1.rds.amazonaws.com" -def test_retry_deadline_uses_timeout_property(): - """Regression (parity review): the retry deadline previously reused - OPEN_CONNECTION_RETRY_INTERVAL_MS; it must come from - OPEN_CONNECTION_RETRY_TIMEOUT_MS.""" - plugin = AuroraInitialConnectionStrategyPlugin.__new__( - AuroraInitialConnectionStrategyPlugin) - plugin._plugin_service = MagicMock() - plugin._plugin_service.get_host_info_by_strategy = MagicMock(return_value=None) + +def _plugin(props, all_hosts=()): + plugin_service = MagicMock() + plugin_service.all_hosts = all_hosts + plugin_service.hosts = all_hosts + plugin_service.accepts_strategy.return_value = True + plugin_service.is_login_exception.return_value = False + plugin_service.is_network_exception.return_value = False + plugin_service.is_read_only_connection_exception.return_value = False + + plugin = AuroraInitialConnectionStrategyPlugin(plugin_service, props) plugin._host_list_provider_service = MagicMock() + return plugin, plugin_service + +def test_retry_deadline_uses_timeout_property(): props = Properties({}) - # Zero total budget: the retry loop must not run even once despite the - # 10-minute interval (with the old bug the deadline WAS the interval). + # Zero total budget: the retry loop must not run even once despite the long interval. WrapperProperties.OPEN_CONNECTION_RETRY_TIMEOUT_MS.set(props, "0") WrapperProperties.OPEN_CONNECTION_RETRY_INTERVAL_MS.set(props, "600000") + plugin, _ = _plugin(props) + connect_func = MagicMock() + + with pytest.raises(AwsWrapperError): + plugin.connect( + MagicMock(), MagicMock(), HostInfo(WRITER_CLUSTER), props, True, connect_func) + + connect_func.assert_not_called() + + +def test_wait_for_initial_topology_disabled_by_default(): + props = Properties({}) + plugin, plugin_service = _plugin(props) + # get_int returns -1 for an absent property; the plugin normalizes that to 0. + assert plugin._wait_for_initial_topology_ms == 0 + + fallback_conn = MagicMock() + connect_func = MagicMock(return_value=fallback_conn) + + host, conn = plugin._open_candidate_connection( + HostInfo(WRITER_CLUSTER), + RdsUrlType.RDS_WRITER_CLUSTER, + InstanceSubstitutionStrategy.SUBSTITUTE_WITH_WRITER, + props, + connect_func) + + assert conn is fallback_conn + assert host.host == WRITER_CLUSTER + plugin_service.force_refresh_host_list.assert_called_once_with(fallback_conn) + plugin_service.force_monitoring_refresh_host_list.assert_not_called() + + +def test_wait_for_initial_topology_connects_to_instance_after_wait(): + props = Properties({}) + WrapperProperties.WAIT_FOR_INITIAL_TOPOLOGY_MS.set(props, "5000") + + writer = HostInfo(WRITER_INSTANCE, role=HostRole.WRITER) + plugin, plugin_service = _plugin(props, all_hosts=()) + plugin_service.force_monitoring_refresh_host_list.return_value = True + + fallback_conn = MagicMock() + instance_conn = MagicMock() + connect_func = MagicMock(return_value=fallback_conn) + + # Topology is empty on the first selection attempt and populated after the wait. + def populate_topology(*_args, **_kwargs): + plugin_service.all_hosts = (writer,) + plugin_service.hosts = (writer,) + return True + + plugin_service.force_monitoring_refresh_host_list.side_effect = populate_topology + plugin_service.connect.return_value = instance_conn + + host, conn = plugin._open_candidate_connection( + HostInfo(WRITER_CLUSTER), + RdsUrlType.RDS_WRITER_CLUSTER, + InstanceSubstitutionStrategy.SUBSTITUTE_WITH_WRITER, + props, + connect_func) + + assert conn is instance_conn + assert host.host == WRITER_INSTANCE + # The timeout must reach the monitor in seconds. + plugin_service.force_monitoring_refresh_host_list.assert_called_once_with(True, 5.0) + fallback_conn.close.assert_called_once() + + +def test_wait_for_initial_topology_timeout_keeps_fallback(): + props = Properties({}) + WrapperProperties.WAIT_FOR_INITIAL_TOPOLOGY_MS.set(props, "5000") + + plugin, plugin_service = _plugin(props) + plugin_service.force_monitoring_refresh_host_list.return_value = False + + fallback_conn = MagicMock() + connect_func = MagicMock(return_value=fallback_conn) + + host, conn = plugin._open_candidate_connection( + HostInfo(WRITER_CLUSTER), + RdsUrlType.RDS_WRITER_CLUSTER, + InstanceSubstitutionStrategy.SUBSTITUTE_WITH_WRITER, + props, + connect_func) + + assert conn is fallback_conn + assert host.host == WRITER_CLUSTER + fallback_conn.close.assert_not_called() + + +def test_wait_for_initial_topology_survives_unsupported_provider(): + props = Properties({}) + WrapperProperties.WAIT_FOR_INITIAL_TOPOLOGY_MS.set(props, "5000") + + plugin, plugin_service = _plugin(props) + plugin_service.force_monitoring_refresh_host_list.side_effect = AwsWrapperError( + "Force monitoring refresh is not supported.") + + fallback_conn = MagicMock() + connect_func = MagicMock(return_value=fallback_conn) + + host, conn = plugin._open_candidate_connection( + HostInfo(WRITER_CLUSTER), + RdsUrlType.RDS_WRITER_CLUSTER, + InstanceSubstitutionStrategy.SUBSTITUTE_WITH_WRITER, + props, + connect_func) + + assert conn is fallback_conn + assert host.host == WRITER_CLUSTER + + +def test_wait_for_initial_topology_instance_connect_failure_keeps_fallback(): + props = Properties({}) + WrapperProperties.WAIT_FOR_INITIAL_TOPOLOGY_MS.set(props, "5000") + + writer = HostInfo(WRITER_INSTANCE, role=HostRole.WRITER) + # Topology must be empty initially so the instance is only selected after the wait. + plugin, plugin_service = _plugin(props, all_hosts=()) + plugin_service.connect.side_effect = AwsWrapperError("instance unreachable") + + def populate_topology(*_args, **_kwargs): + plugin_service.all_hosts = (writer,) + plugin_service.hosts = (writer,) + return True + + plugin_service.force_monitoring_refresh_host_list.side_effect = populate_topology + + fallback_conn = MagicMock() + connect_func = MagicMock(return_value=fallback_conn) + + host, conn = plugin._open_candidate_connection( + HostInfo(WRITER_CLUSTER), + RdsUrlType.RDS_WRITER_CLUSTER, + InstanceSubstitutionStrategy.SUBSTITUTE_WITH_WRITER, + props, + connect_func) + + assert conn is fallback_conn + assert host.host == WRITER_CLUSTER + fallback_conn.close.assert_not_called() + + +def test_do_not_substitute_skips_topology_wait_entirely(): + props = Properties({}) + WrapperProperties.WAIT_FOR_INITIAL_TOPOLOGY_MS.set(props, "5000") + + plugin, plugin_service = _plugin(props) + original_conn = MagicMock() + connect_func = MagicMock(return_value=original_conn) + + host, conn = plugin._open_candidate_connection( + HostInfo(WRITER_CLUSTER), + RdsUrlType.RDS_WRITER_CLUSTER, + InstanceSubstitutionStrategy.DO_NOT_SUBSTITUTE, + props, + connect_func) + + assert conn is original_conn + assert host.host == WRITER_CLUSTER + plugin_service.force_monitoring_refresh_host_list.assert_not_called() + plugin_service.force_refresh_host_list.assert_not_called() + + +def test_available_topology_connects_directly_without_wait(): + props = Properties({}) + WrapperProperties.WAIT_FOR_INITIAL_TOPOLOGY_MS.set(props, "5000") + + writer = HostInfo(WRITER_INSTANCE, role=HostRole.WRITER) + plugin, plugin_service = _plugin(props, all_hosts=(writer,)) + instance_conn = MagicMock() + plugin_service.connect.return_value = instance_conn connect_func = MagicMock() - result = plugin._get_verified_writer_connection(props, True, connect_func) - assert result is None + + host, conn = plugin._open_candidate_connection( + HostInfo(WRITER_CLUSTER), + RdsUrlType.RDS_WRITER_CLUSTER, + InstanceSubstitutionStrategy.SUBSTITUTE_WITH_WRITER, + props, + connect_func) + + assert conn is instance_conn + assert host.host == WRITER_INSTANCE connect_func.assert_not_called() + plugin_service.force_monitoring_refresh_host_list.assert_not_called() + + +def test_instance_connect_skips_this_plugin(): + props = Properties({}) + writer = HostInfo(WRITER_INSTANCE, role=HostRole.WRITER) + plugin, plugin_service = _plugin(props, all_hosts=(writer,)) + plugin_service.connect.return_value = MagicMock() + + plugin._open_candidate_connection( + HostInfo(WRITER_CLUSTER), + RdsUrlType.RDS_WRITER_CLUSTER, + InstanceSubstitutionStrategy.SUBSTITUTE_WITH_WRITER, + props, + MagicMock()) + + plugin_service.connect.assert_called_once() + assert plugin_service.connect.call_args.args[2] is plugin + + +def test_post_wait_instance_connect_skips_this_plugin(): + props = Properties({}) + WrapperProperties.WAIT_FOR_INITIAL_TOPOLOGY_MS.set(props, "5000") + + writer = HostInfo(WRITER_INSTANCE, role=HostRole.WRITER) + plugin, plugin_service = _plugin(props, all_hosts=()) + plugin_service.connect.return_value = MagicMock() + + def populate_topology(*_args, **_kwargs): + plugin_service.all_hosts = (writer,) + plugin_service.hosts = (writer,) + return True + + plugin_service.force_monitoring_refresh_host_list.side_effect = populate_topology + + plugin._open_candidate_connection( + HostInfo(WRITER_CLUSTER), + RdsUrlType.RDS_WRITER_CLUSTER, + InstanceSubstitutionStrategy.SUBSTITUTE_WITH_WRITER, + props, + MagicMock(return_value=MagicMock())) + + plugin_service.connect.assert_called_once() + assert plugin_service.connect.call_args.args[2] is plugin + + +def test_substitute_with_any_raises_unsupported_strategy(): + props = Properties({}) + reader = HostInfo(READER_INSTANCE, role=HostRole.READER) + plugin, plugin_service = _plugin(props, all_hosts=(reader,)) + + with pytest.raises(AwsWrapperError): + plugin._get_candidate_host( + HostInfo("mycluster.cluster-custom-xyz.us-east-1.rds.amazonaws.com"), + RdsUrlType.RDS_CUSTOM_CLUSTER, + InstanceSubstitutionStrategy.SUBSTITUTE_WITH_ANY) + + # The selector must never be consulted for an unsupported role. + plugin_service.get_host_info_by_strategy.assert_not_called() + + +def test_candidate_host_region_filter_uses_allow_block_filtered_hosts(): + props = Properties({}) + in_region = HostInfo(READER_INSTANCE, role=HostRole.READER) # us-east-1 + blocked = HostInfo("blocked-inst.xyz.us-east-1.rds.amazonaws.com", role=HostRole.READER) + plugin, plugin_service = _plugin(props, all_hosts=(in_region, blocked)) + # A custom endpoint is active: hosts excludes the blocked instance. + plugin_service.hosts = (in_region,) + plugin_service.get_host_info_by_strategy.return_value = in_region + + reader_cluster = HostInfo("mycluster.cluster-ro-xyz.us-east-1.rds.amazonaws.com") + plugin._get_candidate_host( + reader_cluster, + RdsUrlType.RDS_READER_CLUSTER, + InstanceSubstitutionStrategy.SUBSTITUTE_WITH_READER) + + passed_list = plugin_service.get_host_info_by_strategy.call_args.args[2] + assert in_region in passed_list + assert blocked not in passed_list + + +def test_candidate_host_substitute_with_reader_uses_reader_role(): + """The reader substitution passes HostRole.READER to the selector.""" + props = Properties({}) + reader = HostInfo(READER_INSTANCE, role=HostRole.READER) + plugin, plugin_service = _plugin(props, all_hosts=(reader,)) + plugin_service.get_host_info_by_strategy.return_value = reader + + result = plugin._get_candidate_host( + HostInfo("mycluster.cluster-ro-xyz.us-east-1.rds.amazonaws.com"), + RdsUrlType.RDS_READER_CLUSTER, + InstanceSubstitutionStrategy.SUBSTITUTE_WITH_READER) + + assert result is reader + assert plugin_service.get_host_info_by_strategy.call_args.args[0] == HostRole.READER From 03452a446f3aecfa4b76feb1c3f177283efd1077 Mon Sep 17 00:00:00 2001 From: karezche <64801825+karenc-bq@users.noreply.github.com> Date: Wed, 5 Aug 2026 16:33:29 -0700 Subject: [PATCH 2/2] chore: address comment --- .../aurora_initial_connection_strategy_plugin.py | 12 ++++++------ aws_advanced_python_wrapper/utils/properties.py | 2 +- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/aws_advanced_python_wrapper/aurora_initial_connection_strategy_plugin.py b/aws_advanced_python_wrapper/aurora_initial_connection_strategy_plugin.py index 59ba8a9aa..70b7f510e 100644 --- a/aws_advanced_python_wrapper/aurora_initial_connection_strategy_plugin.py +++ b/aws_advanced_python_wrapper/aurora_initial_connection_strategy_plugin.py @@ -88,7 +88,7 @@ def from_property_value(cls, value: Optional[str]) -> Optional[RoleVerificationS if setting is None: raise AwsWrapperError(Messages.get_formatted( "AuroraInitialConnectionStrategyPlugin.InvalidPropertyValue", - WrapperProperties.VERIFY_OPENED_CONNECTION_ROLE.name, + WrapperProperties.VERIFY_OPENED_CONNECTION_TYPE.name, value, ", ".join(item.value for item in cls))) return setting @@ -113,7 +113,7 @@ def __init__(self, plugin_service: PluginService, props: Properties): self._wait_for_initial_topology_ms: int = max( 0, WrapperProperties.WAIT_FOR_INITIAL_TOPOLOGY_MS.get_int(props)) - verify_role_value = WrapperProperties.VERIFY_OPENED_CONNECTION_ROLE.get(props) + verify_role_value = WrapperProperties.VERIFY_OPENED_CONNECTION_TYPE.get(props) self._verify_role_prop_value: Optional[str] = \ verify_role_value.lower() if verify_role_value is not None else None @@ -185,7 +185,7 @@ def connect( if self._verify_role_prop_value == RoleVerificationSetting.READER.value: logger.debug( "AuroraInitialConnectionStrategyPlugin.VerifyReaderConfiguredButNoReadersExist", - WrapperProperties.VERIFY_OPENED_CONNECTION_ROLE.name) + WrapperProperties.VERIFY_OPENED_CONNECTION_TYPE.name) self._set_initial_connection_host_info(is_initial_connection, candidate_host) return candidate_conn @@ -217,7 +217,7 @@ def connect( raise AwsWrapperError(Messages.get_formatted( "AuroraInitialConnectionStrategyPlugin.Timeout", self._open_connection_retry_timeout_ns // 1_000_000, - WrapperProperties.VERIFY_OPENED_CONNECTION_ROLE.name)) + WrapperProperties.VERIFY_OPENED_CONNECTION_TYPE.name)) def _open_candidate_connection( self, @@ -428,14 +428,14 @@ def _validate_verification_setting(self, setting: RoleVerificationSetting, url_t and url_type in (RdsUrlType.RDS_WRITER_CLUSTER, RdsUrlType.RDS_GLOBAL_WRITER_CLUSTER)): raise AwsWrapperError(Messages.get_formatted( "AuroraInitialConnectionStrategyPlugin.InvalidSettingForEndpoint", - WrapperProperties.VERIFY_OPENED_CONNECTION_ROLE.name, "reader", "writer cluster or global cluster")) + WrapperProperties.VERIFY_OPENED_CONNECTION_TYPE.name, "reader", "writer cluster or global cluster")) # A custom cluster can only be of type "reader" or "any". if (setting is RoleVerificationSetting.WRITER and url_type in (RdsUrlType.RDS_READER_CLUSTER, RdsUrlType.RDS_CUSTOM_CLUSTER)): raise AwsWrapperError(Messages.get_formatted( "AuroraInitialConnectionStrategyPlugin.InvalidSettingForEndpoint", - WrapperProperties.VERIFY_OPENED_CONNECTION_ROLE.name, "writer", "reader cluster or custom cluster")) + WrapperProperties.VERIFY_OPENED_CONNECTION_TYPE.name, "writer", "reader cluster or custom cluster")) def _get_candidate_host( self, diff --git a/aws_advanced_python_wrapper/utils/properties.py b/aws_advanced_python_wrapper/utils/properties.py index d21e45b83..25be0ecd5 100644 --- a/aws_advanced_python_wrapper/utils/properties.py +++ b/aws_advanced_python_wrapper/utils/properties.py @@ -614,7 +614,7 @@ class WrapperProperties: "writer", ) - VERIFY_OPENED_CONNECTION_ROLE = WrapperProperty( + VERIFY_OPENED_CONNECTION_TYPE = WrapperProperty( "verify_opened_connection_type", "Defines whether an opened connection should be verified to be a writer or reader, " "or if no role verification should be performed. Valid values are 'writer', 'reader', or 'none'.",