From f493bba5a7c6c6a5581bd8aa13b15af778f0d9c3 Mon Sep 17 00:00:00 2001 From: karezche <64801825+karenc-bq@users.noreply.github.com> Date: Mon, 5 Oct 2026 14:57:47 -0700 Subject: [PATCH] chore: update initial connection strategy plugin --- ...rora_initial_connection_strategy_plugin.py | 190 +++++++----- .../aio/plugin_service.py | 26 +- aws_advanced_python_wrapper/aio/wrapper.py | 12 + ...rora_initial_connection_strategy_plugin.py | 270 +++++++++++++----- aws_advanced_python_wrapper/plugin_service.py | 27 +- ...dvanced_python_wrapper_messages.properties | 5 +- ...heAuroraInitialConnectionStrategyPlugin.md | 3 +- .../test_aio_aurora_initial_connection.py | 125 +++++++- tests/unit/test_aio_plugin_service_extras.py | 43 ++- ...rora_initial_connection_strategy_plugin.py | 11 +- 10 files changed, 554 insertions(+), 158 deletions(-) diff --git a/aws_advanced_python_wrapper/aio/aurora_initial_connection_strategy_plugin.py b/aws_advanced_python_wrapper/aio/aurora_initial_connection_strategy_plugin.py index 20235f68b..08fba084b 100644 --- a/aws_advanced_python_wrapper/aio/aurora_initial_connection_strategy_plugin.py +++ b/aws_advanced_python_wrapper/aio/aurora_initial_connection_strategy_plugin.py @@ -41,7 +41,8 @@ from aws_advanced_python_wrapper.aio.plugin import AsyncPlugin from aws_advanced_python_wrapper.aurora_initial_connection_strategy_plugin import ( - InstanceSubstitutionStrategy, RoleVerificationSetting) + InstanceSubstitutionStrategy, RoleVerificationSetting, + _inactive_cluster_writer_role, _inactive_cluster_writer_strategy) from aws_advanced_python_wrapper.errors import AwsWrapperError from aws_advanced_python_wrapper.host_availability import HostAvailability from aws_advanced_python_wrapper.hostinfo import HostInfo, HostRole @@ -116,79 +117,111 @@ async def connect( connect_func: Callable[..., Awaitable[Any]]) -> Any: 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) loop = asyncio.get_running_loop() deadline = loop.time() + (self._open_connection_retry_timeout_ns / 1_000_000_000) + last_error: Optional[Exception] = None + + while True: + 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) + + if (substitution_strategy is InstanceSubstitutionStrategy.DO_NOT_SUBSTITUTE + and role_to_verify is None): + # Nothing to substitute and no role to verify, so retrying cannot change the + # outcome. + conn = await connect_func() + self._set_initial_connection_host_info(is_initial_connection, host_info) + return conn - while loop.time() < deadline: candidate_conn: Optional[Any] = None candidate_host: Optional[HostInfo] = None + attempted: List[HostInfo] = [] try: candidate_host, candidate_conn = await self._open_candidate_connection( - host_info, url_type, substitution_strategy, props, connect_func, driver_dialect) - - 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 role_to_verify is None: - # No verification required. - self._set_initial_connection_host_info(is_initial_connection, candidate_host) - return candidate_conn - - conn_role = await 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. - await 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_TYPE.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) - await self._close_connection(candidate_conn, driver_dialect) - await self._delay(self._retry_delay_ms) + host_info, url_type, substitution_strategy, props, connect_func, driver_dialect, + attempted) + + if candidate_conn is not None: + 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 = await 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. + await 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_TYPE.name) + self._set_initial_connection_host_info(is_initial_connection, candidate_host) + return candidate_conn + + last_error = AwsWrapperError(Messages.get_formatted( + "AuroraInitialConnectionStrategyPlugin.IncorrectRole", + candidate_host.host, role_to_verify.name.lower(), conn_role.name.lower())) + logger.debug( + "AuroraInitialConnectionStrategyPlugin.IncorrectRole", + candidate_host.host, role_to_verify.name.lower(), conn_role.name.lower()) + await self._close_connection(candidate_conn, driver_dialect) except Exception as e: await self._close_connection(candidate_conn, driver_dialect) if self._plugin_service.is_login_exception(error=e): raise - if candidate_host is not None: + last_error = e + # Marked only for a substituted candidate, and only once the candidate is known. + failed_host = attempted[-1] if attempted else None + if failed_host is not None: self._plugin_service.set_availability( - candidate_host.as_aliases(), HostAvailability.UNAVAILABLE) - - if self._plugin_service.is_network_exception(error=e): - # Retry connection. - continue + failed_host.as_aliases(), HostAvailability.UNAVAILABLE) - if (self._plugin_service.is_read_only_connection_exception(error=e) + retryable = ( + self._plugin_service.is_network_exception(error=e) + or (self._plugin_service.is_read_only_connection_exception(error=e) and (role_to_verify == HostRole.WRITER - or substitution_strategy is InstanceSubstitutionStrategy.SUBSTITUTE_WITH_WRITER)): - # Retry connection. - continue + or substitution_strategy is InstanceSubstitutionStrategy.SUBSTITUTE_WITH_WRITER))) + if not retryable: + raise - raise + # Checked after the attempt rather than before it, so a zero retry budget means + # "try once" rather than "never connect at all". + if not await self._delay_unless_expired(loop, deadline): + break + timeout_ms = self._open_connection_retry_timeout_ns // 1_000_000 + if last_error is None: + raise AwsWrapperError(Messages.get_formatted( + "AuroraInitialConnectionStrategyPlugin.Timeout", + timeout_ms, WrapperProperties.VERIFY_OPENED_CONNECTION_TYPE.name)) raise AwsWrapperError(Messages.get_formatted( - "AuroraInitialConnectionStrategyPlugin.Timeout", - self._open_connection_retry_timeout_ns // 1_000_000, - WrapperProperties.VERIFY_OPENED_CONNECTION_TYPE.name)) + "AuroraInitialConnectionStrategyPlugin.TimeoutWithCause", + timeout_ms, WrapperProperties.VERIFY_OPENED_CONNECTION_TYPE.name, str(last_error))) + + async def _delay_unless_expired(self, loop: asyncio.AbstractEventLoop, deadline: float) -> bool: + """Sleep the retry interval, reporting whether the retry budget allows another attempt. + + Returns ``False`` once the budget is exhausted. The sleep never overshoots the + deadline, so a budget shorter than the retry interval does not stretch the call. + """ + remaining_sec = deadline - loop.time() + if remaining_sec <= 0: + return False + + await self._delay(min(self._retry_delay_ms, int(remaining_sec * 1000))) + return loop.time() < deadline async def _open_candidate_connection( self, @@ -197,7 +230,8 @@ async def _open_candidate_connection( substitution_strategy: InstanceSubstitutionStrategy, props: Properties, connect_func: Callable[..., Awaitable[Any]], - driver_dialect: AsyncDriverDialect) -> Tuple[HostInfo, Optional[Any]]: + driver_dialect: AsyncDriverDialect, + attempted: Optional[List[HostInfo]] = None) -> Tuple[HostInfo, Optional[Any]]: """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 @@ -205,6 +239,10 @@ async def _open_candidate_connection( 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. + + Each substituted host is appended to ``attempted`` before it is connected to. The caller + cannot otherwise learn which host failed, because an exception here means its return value + was never unpacked -- which left the caller unable to mark anything unavailable. """ if substitution_strategy is InstanceSubstitutionStrategy.DO_NOT_SUBSTITUTE: return original_connect_host, await connect_func() @@ -214,6 +252,8 @@ async def _open_candidate_connection( # Topology is already available; connect to the selected instance. Skipping this plugin # avoids re-entering the pipeline here, while still letting the auth plugins (IAM, # Secrets, Federated, Okta) re-apply on the new connection. + if attempted is not None: + attempted.append(candidate_host) return candidate_host, await self._plugin_service.connect( candidate_host, props, plugin_to_skip=self) @@ -223,13 +263,17 @@ async def _open_candidate_connection( # topology, and it serves as a fallback connection if instance selection or connection fails. candidate_conn = await connect_func() - if self._wait_for_initial_topology_ms <= 0: - # Feature disabled. Preserve the previous behavior. - await self._plugin_service.force_refresh_host_list(candidate_conn) - return original_connect_host, candidate_conn + try: + if self._wait_for_initial_topology_ms <= 0: + # Feature disabled. Preserve the previous behavior. + await self._plugin_service.force_refresh_host_list(candidate_conn) + return original_connect_host, candidate_conn - return await self._wait_for_topology_and_connect_to_instance( - original_connect_host, url_type, substitution_strategy, props, candidate_conn, driver_dialect) + return await self._wait_for_topology_and_connect_to_instance( + original_connect_host, url_type, substitution_strategy, props, candidate_conn, driver_dialect) + except BaseException: + await self._close_connection(candidate_conn, driver_dialect) + raise async def _wait_for_topology_and_connect_to_instance( self, @@ -255,8 +299,11 @@ async def _wait_for_topology_and_connect_to_instance( # Unlike sync, this needs no guard: async providers without monitor support fall back to a # plain refresh and report whether any topology is held, rather than raising. + # + # Deliberately does not ask for the writer to be verified. Doing so closes the shared + # per-cluster monitoring connection and clears its verified-writer state. timeout_sec = self._wait_for_initial_topology_ms / 1000 - topology_fetched = await self._plugin_service.force_monitoring_refresh_host_list(True, timeout_sec) + topology_fetched = await self._plugin_service.force_monitoring_refresh_host_list(False, timeout_sec) if not topology_fetched: logger.debug( @@ -304,7 +351,7 @@ def _get_instance_substitution_strategy( 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 + return _inactive_cluster_writer_strategy(props) if self._rds_utils.is_same_region(writer.host, original_host): return InstanceSubstitutionStrategy.SUBSTITUTE_WITH_WRITER @@ -313,10 +360,7 @@ def _get_instance_substitution_strategy( # 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 _inactive_cluster_writer_strategy(props) if url_type == RdsUrlType.RDS_READER_CLUSTER: return InstanceSubstitutionStrategy.SUBSTITUTE_WITH_READER @@ -333,6 +377,14 @@ def _validate_substitution_strategy( "AuroraInitialConnectionStrategyPlugin.InvalidSettingForInstanceEndpoint", WrapperProperties.ENDPOINT_SUBSTITUTION_ROLE.name)) + if url_type in (RdsUrlType.RDS_PROXY, RdsUrlType.RDS_AURORA_LIMITLESS_DB_SHARD_GROUP): + # Substituting an instance host would route connections around the endpoint the + # user asked to connect through, losing its pooling, authentication handling and + # failover behaviour. + raise AwsWrapperError(Messages.get_formatted( + "AuroraInitialConnectionStrategyPlugin.InvalidSettingForManagedEndpoint", + WrapperProperties.ENDPOINT_SUBSTITUTION_ROLE.name, setting.value)) + if not url_type.is_rds_cluster: return @@ -390,9 +442,7 @@ def _get_role_to_verify( # 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 + return _inactive_cluster_writer_role(props) if url_type == RdsUrlType.RDS_READER_CLUSTER: return HostRole.READER diff --git a/aws_advanced_python_wrapper/aio/plugin_service.py b/aws_advanced_python_wrapper/aio/plugin_service.py index ba4628192..45a2d30a6 100644 --- a/aws_advanced_python_wrapper/aio/plugin_service.py +++ b/aws_advanced_python_wrapper/aio/plugin_service.py @@ -502,13 +502,35 @@ def set_availability( self, host_aliases: FrozenSet[str], availability: HostAvailability) -> None: - for alias in host_aliases: + if not host_aliases: + return + + hosts_to_change = { + host.url: host for host in self.all_hosts + if not host_aliases.isdisjoint(host.all_aliases) + } + if not hosts_to_change: + return + + changes: Dict[str, Set[HostEvent]] = {} + for url, host in hosts_to_change.items(): + previous_availability = host.get_raw_availability() + host.set_availability(availability) AsyncPluginServiceImpl._host_availability_expiring_cache.put( - alias, + url, availability, AsyncPluginServiceImpl._HOST_AVAILABILITY_EXPIRATION_NANO, ) + if previous_availability != availability: + changes[url] = { + HostEvent.WENT_UP if availability == HostAvailability.AVAILABLE else HostEvent.WENT_DOWN, + HostEvent.HOST_CHANGED, + } + + if changes and self._plugin_manager is not None: + self._plugin_manager.notify_host_list_changed(changes) + def get_availability(self, host_url: str) -> Optional[HostAvailability]: return AsyncPluginServiceImpl._host_availability_expiring_cache.get(host_url) diff --git a/aws_advanced_python_wrapper/aio/wrapper.py b/aws_advanced_python_wrapper/aio/wrapper.py index 4ec42b833..9a77e5884 100644 --- a/aws_advanced_python_wrapper/aio/wrapper.py +++ b/aws_advanced_python_wrapper/aio/wrapper.py @@ -57,6 +57,12 @@ TelemetryFactory +# A plugin belongs here when it reads the host list -- all_hosts, hosts, +# identify_connection, or a forced refresh. Omitting one does not disable the plugin, it +# hands it a static provider holding only the connection-string host, so the plugin either +# raises (host monitoring cannot identify the instance behind a cluster endpoint) or +# silently does nothing (initial_connection finds no instance host to substitute). Add new +# topology-reading plugins here; "srw" and "limitless" are deliberately absent because they don't require the topology. _TOPOLOGY_REQUIRING_PLUGINS = frozenset({ "failover", "failover_v2", @@ -65,6 +71,12 @@ "gdb_rw", "custom_endpoint", "aurora_connection_tracker", + "initial_connection", + "host_monitoring", + "host_monitoring_v2", + "stale_dns", + "fastest_response_strategy", + "bg", }) 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 8eae5cf7b..003cb3fdd 100644 --- a/aws_advanced_python_wrapper/aurora_initial_connection_strategy_plugin.py +++ b/aws_advanced_python_wrapper/aurora_initial_connection_strategy_plugin.py @@ -101,6 +101,57 @@ def from_property_value(cls, value: Optional[str]) -> Optional[RoleVerificationS item.value: item for item in RoleVerificationSetting } +_INACTIVE_SUBSTITUTION_STRATEGY_BY_KEY: Dict[str, InstanceSubstitutionStrategy] = { + InstanceSubstitutionStrategy.SUBSTITUTE_WITH_WRITER.value: InstanceSubstitutionStrategy.SUBSTITUTE_WITH_WRITER, + InstanceSubstitutionStrategy.DO_NOT_SUBSTITUTE.value: InstanceSubstitutionStrategy.DO_NOT_SUBSTITUTE, +} + +_INACTIVE_VERIFICATION_ROLE_BY_KEY: Dict[str, Optional[HostRole]] = { + "writer": HostRole.WRITER, + "none": None, +} + + +def _parse_restricted( + value: str, allowed: Dict, prop) -> object: + """Look ``value`` up in ``allowed``, raising with the allowed values when absent.""" + key = value.lower() + if key not in allowed: + raise AwsWrapperError(Messages.get_formatted( + "AuroraInitialConnectionStrategyPlugin.InvalidPropertyValue", + prop.name, value, ", ".join(sorted(allowed)))) + return allowed[key] + + +def _inactive_cluster_writer_strategy(props: Properties) -> InstanceSubstitutionStrategy: + """Return whether an inactive cluster writer endpoint should be substituted.""" + # return_default=False because the declared default is 'writer' + value = WrapperProperties.INACTIVE_CLUSTER_WRITER_SUBSTITUTION_ROLE.get(props, return_default=False) + if value is None: + return InstanceSubstitutionStrategy.SUBSTITUTE_WITH_WRITER + + strategy = _parse_restricted( + value, _INACTIVE_SUBSTITUTION_STRATEGY_BY_KEY, + WrapperProperties.INACTIVE_CLUSTER_WRITER_SUBSTITUTION_ROLE) + assert isinstance(strategy, InstanceSubstitutionStrategy) + return strategy + + +def _inactive_cluster_writer_role(props: Properties) -> Optional[HostRole]: + """Return the role a connection through an inactive cluster writer endpoint must have.""" + # return_default=False because the declared default is 'writer' + value = WrapperProperties.VERIFY_INACTIVE_CLUSTER_WRITER_CONNECTION_ROLE.get(props, return_default=False) + if value is not None: + role = _parse_restricted( + value, _INACTIVE_VERIFICATION_ROLE_BY_KEY, + WrapperProperties.VERIFY_INACTIVE_CLUSTER_WRITER_CONNECTION_ROLE) + assert role is None or isinstance(role, HostRole) + return role + + if _inactive_cluster_writer_strategy(props) is InstanceSubstitutionStrategy.DO_NOT_SUBSTITUTE: + return None + return HostRole.WRITER + class AuroraInitialConnectionStrategyPlugin(Plugin): _SUBSCRIBED_METHODS: Set[str] = {"init_host_provider", "connect"} @@ -151,77 +202,127 @@ def connect( 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 + # Carries the most recent attempt failure into the timeout error, so the caller is + # told what kept failing rather than only that time ran out. + last_error: Optional[Exception] = None + + while True: + # Both are re-derived every attempt because both read the topology, which the + # previous attempt may have refreshed. Deriving the strategy once pinned a + # connection whose topology was empty for the whole retry window, so it could + # never substitute the writer it had just discovered. The role can change across + # attempts too, and that is deliberate: a cold topology makes a writer cluster + # endpoint look inactive, and once a refresh names a same-region writer the + # endpoint is an active cluster writer endpoint whose role must be verified. + 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) + + if (substitution_strategy is InstanceSubstitutionStrategy.DO_NOT_SUBSTITUTE + and role_to_verify is None): + # Nothing to substitute and no role to verify, so retrying cannot change the + # outcome. Connect once and let the driver's own error surface instead of + # replacing it with a timeout that names a parameter the caller never set. + # The host is still recorded: the host list provider derives the initial + # connection from it, and skipping that moved topology resolution off the + # endpoint the caller named. + conn = connect_func() + self._set_initial_connection_host_info(is_initial_connection, host_info) + return conn - while perf_counter_ns() < end_time_ns: candidate_conn: Optional[Connection] = None candidate_host: Optional[HostInfo] = None + attempted: List[HostInfo] = [] try: candidate_host, candidate_conn = self._open_candidate_connection( - host_info, url_type, substitution_strategy, props, connect_func) - - 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 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_TYPE.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) + host_info, url_type, substitution_strategy, props, connect_func, attempted) + + if candidate_conn is not None: + 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_TYPE.name) + self._set_initial_connection_host_info(is_initial_connection, candidate_host) + return candidate_conn + + last_error = AwsWrapperError(Messages.get_formatted( + "AuroraInitialConnectionStrategyPlugin.IncorrectRole", + candidate_host.host, role_to_verify.name.lower(), conn_role.name.lower())) + logger.debug( + "AuroraInitialConnectionStrategyPlugin.IncorrectRole", + candidate_host.host, role_to_verify.name.lower(), conn_role.name.lower()) + self._close_connection(candidate_conn) except Exception as e: self._close_connection(candidate_conn) if self._plugin_service.is_login_exception(e): raise - if candidate_host is not None: + last_error = e + # Marked only for a substituted candidate, and only once the candidate is known. + # Reading it back from the unpacked return value never worked, because an + # exception means there was nothing to unpack -- so a reader that refused a + # connection was never marked, and the next attempt's selector could keep + # choosing it for the whole retry window. Marking the host the user named in + # the connection string is deliberately skipped: that would hand the selectors + # a verdict about a host they were never asked to choose. + failed_host = attempted[-1] if attempted else None + if failed_host is not None: self._plugin_service.set_availability( - candidate_host.as_aliases(), HostAvailability.UNAVAILABLE) - - if self._plugin_service.is_network_exception(e): - # Retry connection. - continue + failed_host.as_aliases(), HostAvailability.UNAVAILABLE) - if (self._plugin_service.is_read_only_connection_exception(e) + retryable = ( + self._plugin_service.is_network_exception(e) + or (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 + or substitution_strategy is InstanceSubstitutionStrategy.SUBSTITUTE_WITH_WRITER))) + if not retryable: + raise - raise + # Checked after the attempt rather than before it, so a zero retry budget means + # "try once" rather than "never connect at all". + if not self._delay_unless_expired(end_time_ns): + break + timeout_ms = self._open_connection_retry_timeout_ns // 1_000_000 + if last_error is None: + raise AwsWrapperError(Messages.get_formatted( + "AuroraInitialConnectionStrategyPlugin.Timeout", + timeout_ms, WrapperProperties.VERIFY_OPENED_CONNECTION_TYPE.name)) raise AwsWrapperError(Messages.get_formatted( - "AuroraInitialConnectionStrategyPlugin.Timeout", - self._open_connection_retry_timeout_ns // 1_000_000, - WrapperProperties.VERIFY_OPENED_CONNECTION_TYPE.name)) + "AuroraInitialConnectionStrategyPlugin.TimeoutWithCause", + timeout_ms, WrapperProperties.VERIFY_OPENED_CONNECTION_TYPE.name, str(last_error))) + + def _delay_unless_expired(self, end_time_ns: int) -> bool: + """Sleep the retry interval, reporting whether the retry budget allows another attempt. + + Returns ``False`` once the budget is exhausted. The sleep never overshoots the + deadline, so a budget shorter than the retry interval does not stretch the call. + """ + remaining_ns = end_time_ns - perf_counter_ns() + if remaining_ns <= 0: + return False + + self._delay(min(self._retry_delay_ms, remaining_ns // 1_000_000)) + return perf_counter_ns() < end_time_ns def _open_candidate_connection( self, @@ -229,7 +330,8 @@ def _open_candidate_connection( url_type: RdsUrlType, substitution_strategy: InstanceSubstitutionStrategy, props: Properties, - connect_func: Callable) -> Tuple[HostInfo, Optional[Connection]]: + connect_func: Callable, + attempted: Optional[List[HostInfo]] = None) -> 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 @@ -237,6 +339,10 @@ def _open_candidate_connection( 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. + + Each substituted host is appended to ``attempted`` before it is connected to. The caller + cannot otherwise learn which host failed, because an exception here means its return value + was never unpacked -- which left the caller unable to mark anything unavailable. """ if substitution_strategy is InstanceSubstitutionStrategy.DO_NOT_SUBSTITUTE: return original_connect_host, connect_func() @@ -244,6 +350,8 @@ def _open_candidate_connection( 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. + if attempted is not None: + attempted.append(candidate_host) 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. @@ -252,13 +360,17 @@ def _open_candidate_connection( # 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 + try: + 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) + return self._wait_for_topology_and_connect_to_instance( + original_connect_host, url_type, substitution_strategy, props, candidate_conn) + except Exception: + self._close_connection(candidate_conn) + raise def _wait_for_topology_and_connect_to_instance( self, @@ -285,8 +397,18 @@ def _wait_for_topology_and_connect_to_instance( # 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: + # Deliberately does not ask for the writer to be verified. Doing so closes the + # shared per-cluster monitoring connection and clears its verified-writer state, + # which is right after a failover but not at connect time, when nothing has failed + # over: under a pool prefill every arriving connection would tear down what the + # previous one had just established. The role of the connection returned here is + # verified directly, on the real connection, by the caller. + topology_fetched = self._plugin_service.force_monitoring_refresh_host_list(False, timeout_sec) + except Exception as e: + # A provider without monitor support raises rather than returning a host list, and + # so does a genuine refresh failure. Both fall back to the endpoint connection, but + # a failure is not a timeout and must not be reported as one. + logger.debug("AuroraInitialConnectionStrategyPlugin.WaitForTopologyError", str(e)) topology_fetched = False if not topology_fetched: @@ -332,7 +454,12 @@ def _get_instance_substitution_strategy( if url_type == RdsUrlType.RDS_WRITER_CLUSTER: writer = self._find_writer(self._plugin_service.all_hosts) if writer is None or not self._rds_utils.is_rds_instance(writer.host): - return InstanceSubstitutionStrategy.DO_NOT_SUBSTITUTE + # With no topology there is no way to tell an active cluster writer endpoint + # from an inactive one, so this defers to the same setting the cross-region + # case below uses. Without that, a request not to substitute an inactive + # endpoint would be overridden on the first connection of every process, and + # then overridden again with a writer in another region once topology arrived. + return _inactive_cluster_writer_strategy(props) if self._rds_utils.is_same_region(writer.host, original_host): return InstanceSubstitutionStrategy.SUBSTITUTE_WITH_WRITER @@ -341,10 +468,7 @@ def _get_instance_substitution_strategy( # 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 _inactive_cluster_writer_strategy(props) if url_type == RdsUrlType.RDS_READER_CLUSTER: return InstanceSubstitutionStrategy.SUBSTITUTE_WITH_READER @@ -361,6 +485,14 @@ def _validate_substitution_strategy( "AuroraInitialConnectionStrategyPlugin.InvalidSettingForInstanceEndpoint", WrapperProperties.ENDPOINT_SUBSTITUTION_ROLE.name)) + if url_type in (RdsUrlType.RDS_PROXY, RdsUrlType.RDS_AURORA_LIMITLESS_DB_SHARD_GROUP): + # Substituting an instance host would route connections around the endpoint the + # user asked to connect through, losing its pooling, authentication handling and + # failover behaviour. + raise AwsWrapperError(Messages.get_formatted( + "AuroraInitialConnectionStrategyPlugin.InvalidSettingForManagedEndpoint", + WrapperProperties.ENDPOINT_SUBSTITUTION_ROLE.name, setting.value)) + if not url_type.is_rds_cluster: return @@ -418,9 +550,7 @@ def _get_role_to_verify( # 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 + return _inactive_cluster_writer_role(props) if url_type == RdsUrlType.RDS_READER_CLUSTER: return HostRole.READER diff --git a/aws_advanced_python_wrapper/plugin_service.py b/aws_advanced_python_wrapper/plugin_service.py index f6296d876..a78da6f9e 100644 --- a/aws_advanced_python_wrapper/plugin_service.py +++ b/aws_advanced_python_wrapper/plugin_service.py @@ -323,6 +323,7 @@ def get_status(self, clazz: Type[StatusType], key: str) -> Optional[StatusType]: class PluginServiceImpl(PluginService, HostListProviderService, CanReleaseResources): _STATUS_CACHE_EXPIRATION_NANO = 60 * 60 * 1_000_000_000 # one hour + _HOST_AVAILABILITY_EXPIRATION_NANO: ClassVar[int] = 5 * 60 * 1_000_000_000 # five minutes _host_availability_expiring_cache: CacheMap[str, HostAvailability] = CacheMap() _executor_name: ClassVar[str] = "PluginServiceImplExecutor" @@ -618,7 +619,31 @@ def force_connect( self._target_func, self._driver_dialect, host_info, props, self.current_connection is None, plugin_to_skip) def set_availability(self, host_aliases: FrozenSet[str], availability: HostAvailability): - ... + if not host_aliases: + return + + hosts_to_change = { + host.url: host for host in self.all_hosts + if not host_aliases.isdisjoint(host.all_aliases) + } + if not hosts_to_change: + return + + changes: Dict[str, Set[HostEvent]] = {} + for url, host in hosts_to_change.items(): + previous_availability = host.get_raw_availability() + host.set_availability(availability) + PluginServiceImpl._host_availability_expiring_cache.put( + url, availability, PluginServiceImpl._HOST_AVAILABILITY_EXPIRATION_NANO) + + if previous_availability != availability: + changes[url] = { + HostEvent.WENT_UP if availability == HostAvailability.AVAILABLE else HostEvent.WENT_DOWN, + HostEvent.HOST_CHANGED, + } + + if changes: + self._container.plugin_manager.notify_host_list_changed(changes) def identify_connection(self, connection: Optional[Connection] = None) -> Optional[HostInfo]: connection = self.current_connection if connection is None else connection 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 911056ccf..efdeea189 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 @@ -22,7 +22,10 @@ AuroraInitialConnectionStrategyPlugin.InvalidSettingForEndpoint=[AuroraInitialCo 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.TimeoutWithCause=[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. The last failure was: {} +AuroraInitialConnectionStrategyPlugin.InvalidSettingForManagedEndpoint=[AuroraInitialConnectionStrategyPlugin] Parameter '{}' cannot be set to '{}' when using a proxy or Limitless shard group endpoint, because substituting an instance host would bypass that endpoint. +AuroraInitialConnectionStrategyPlugin.WaitForTopologyError=[AuroraInitialConnectionStrategyPlugin] An error occurred while waiting for the cluster topology to be fetched: {}. Falling back to connecting via the provided endpoint. +AuroraInitialConnectionStrategyPlugin.IncorrectRole=[AuroraInitialConnectionStrategyPlugin] The connection opened to '{}' should have the role '{}' but has the role '{}'. 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. diff --git a/docs/using-the-python-wrapper/using-plugins/UsingTheAuroraInitialConnectionStrategyPlugin.md b/docs/using-the-python-wrapper/using-plugins/UsingTheAuroraInitialConnectionStrategyPlugin.md index ad94447ca..d63e261bd 100644 --- a/docs/using-the-python-wrapper/using-plugins/UsingTheAuroraInitialConnectionStrategyPlugin.md +++ b/docs/using-the-python-wrapper/using-plugins/UsingTheAuroraInitialConnectionStrategyPlugin.md @@ -31,7 +31,7 @@ The following properties can be used to configure the Aurora Initial Connection | `endpoint_substitution_role` | String | No | Defines whether the initial connection URL should be replaced with an instance URL from the topology when available, and if so, the role of the instance URL to select. Set this only when using a URL that resolves to a cluster endpoint (global writer, writer, reader, or custom).

For writer cluster or global writer endpoints, valid values are `writer` and `none`. For reader endpoints, valid values are `reader` and `none`. For custom cluster endpoints, valid values are `reader` and `none`. If set to `none`, the initial URL is not replaced. | `reader` | `writer` for writer/global writer cluster endpoints.

`reader` for reader cluster endpoints.

Otherwise: `none` (no substitution). | | `verify_opened_connection_type` | String | No | Defines whether an opened connection should be verified to be a writer or reader after connecting, or if no role verification should be performed.

For writer or global writer endpoints, valid values are `writer` and `none`. For reader and custom endpoints, valid values are `reader` and `none`. The value `none` performs no role verification. | `reader` | `writer` for writer/global writer cluster endpoints.

`reader` for reader cluster endpoints.

Otherwise: `none`. | | `inactive_cluster_writer_endpoint_substitution_role` | String | No | Applicable to Aurora Global Databases. Defines whether the inactive cluster writer endpoint in the initial connection URL should be replaced with a writer instance URL from the topology when available. Region-bound cluster writer endpoints may be inactive depending on the Global Database primary region; this parameter configures the desired behavior for them. Valid values are `writer` and `none`. | `none` | `writer` | -| `verify_inactive_cluster_writer_endpoint_connection_type` | String | No | Applicable to Aurora Global Databases. Defines whether a connection opened via an inactive cluster writer endpoint should be verified to be a writer, or if no role verification should be performed. Valid values are `writer` and `none`. | `none` | `writer` | +| `verify_inactive_cluster_writer_endpoint_connection_type` | String | No | Applicable to Aurora Global Databases. Defines whether a connection opened via an inactive cluster writer endpoint should be verified to be a writer, or if no role verification should be performed. Valid values are `writer` and `none`.

When this parameter is not set and `inactive_cluster_writer_endpoint_substitution_role` is `none`, no role is verified: requiring a writer while forbidding the substitution that would produce one cannot be satisfied, so the two are coupled. Set this parameter explicitly to verify a role regardless. | `none` | `writer` | | `wait_for_initial_topology_ms` | Integer | No | Maximum time, in milliseconds, to wait for the cluster topology to be fetched before opening a new connection. When set greater than `0` and the topology is not yet available, the plugin blocks until the topology is discovered (or this timeout is reached) instead of falling back to connecting via the initial endpoint. This lets host selection strategies such as `round_robin` distribute concurrent and connection-pool prefill connections across instances rather than routing them all to a single DNS-resolved instance. The wait is scoped to the cluster. When set to `0` (the default) the previous behavior is preserved. | `30000` | `0` | | `open_connection_retry_timeout_ms` | Integer | No | The maximum allowed time for retries when opening a connection in milliseconds. | `40000` | `30000` | | `open_connection_retry_interval_ms` | Integer | No | The time between retries when opening a connection in milliseconds. | `2000` | `1000` | @@ -46,6 +46,7 @@ The following properties can be used to configure the Aurora Initial Connection | Reader cluster | `reader`, `none` | `reader`, `none` | | Custom cluster | `reader`, `none` | `reader`, `none` | | Instance | `none` only | `none` only | +| Proxy / Limitless shard group | `none` only | `none` only | Setting a value outside the valid set for the given endpoint raises an error. diff --git a/tests/unit/test_aio_aurora_initial_connection.py b/tests/unit/test_aio_aurora_initial_connection.py index afd2fa944..67ce52658 100644 --- a/tests/unit/test_aio_aurora_initial_connection.py +++ b/tests/unit/test_aio_aurora_initial_connection.py @@ -386,9 +386,11 @@ async def _run(): def test_availability_marked_when_verification_fails_after_connecting(): - """set_availability is reached only when the failure happens after - _open_candidate_connection returned a host. When the candidate connect - itself throws, candidate_host is still None and the mark is skipped.""" + """set_availability is reached for any failure against a substituted candidate. + + The candidate is recorded before it is connected to, so a failure in the connect + itself is marked too. Relying on the unpacked return value meant an exception from + the connect left the host unknown and the mark was silently skipped.""" reader = _reader_host() plugin, svc, driver_dialect = _build( all_hosts=(_writer_host(), reader), @@ -774,3 +776,120 @@ def test_wait_for_initial_topology_defaults_to_zero(): """get_int returns -1 for an absent property; the plugin normalizes it to 0.""" plugin, _, _ = _build() assert plugin._wait_for_initial_topology_ms == 0 + + +# ---- Regression tests for the ported defect fixes ----------------------- + + +def test_topology_wait_does_not_ask_the_monitor_to_verify_the_writer(): + """Verifying the writer closes the shared per-cluster monitoring connection. + + That is right after a failover but not at connect time, when nothing has failed over: + under a pool prefill every arriving connection would tear down what the previous one + had just established, which is the opposite of what the wait exists to achieve. The + role of the returned connection is verified directly on the connection instead. + """ + # The topology is deliberately empty: with an instance already available the plugin + # connects straight to it and the wait is never reached. + plugin, svc, driver_dialect = _build( + all_hosts=(), + role=HostRole.WRITER, + strategy_pick=None, + props_overrides={"wait_for_initial_topology_ms": "5000"}, + ) + svc.force_monitoring_refresh_host_list = AsyncMock( # type: ignore[method-assign] + return_value=False) + + async def _connect_func(): + return MagicMock(name="endpoint_conn") + + async def _run(): + return await plugin.connect( + target_driver_func=MagicMock(), + driver_dialect=driver_dialect, + host_info=_cluster_host_info(_WRITER_CLUSTER), + props=svc.props, + is_initial_connection=True, + connect_func=_connect_func, + ) + + asyncio.run(_run()) + + svc.force_monitoring_refresh_host_list.assert_awaited_once_with(False, 5.0) + + +def test_zero_retry_budget_still_makes_one_attempt(): + """A zero budget means "do not retry", not "do not connect". + + The budget is checked after an attempt rather than before one. Checking it first meant + any deployment carrying 0 failed every connection without ever calling the driver. + """ + plugin, svc, driver_dialect = _build( + all_hosts=(), + role=HostRole.READER, + props_overrides={ + "open_connection_retry_timeout_ms": "0", + "open_connection_retry_interval_ms": "600000", + }, + ) + + calls = 0 + + async def _connect_func(): + nonlocal calls + calls += 1 + return MagicMock(name="endpoint_conn") + + async def _run(): + return await plugin.connect( + target_driver_func=MagicMock(), + driver_dialect=driver_dialect, + host_info=_cluster_host_info(_WRITER_CLUSTER), + props=svc.props, + is_initial_connection=True, + connect_func=_connect_func, + ) + + with pytest.raises(AwsWrapperError): + asyncio.run(_run()) + + assert calls == 1 + + +def test_inactive_cluster_writer_none_alone_does_not_deadlock_async(): + """Declining to substitute must also stop the role being verified. + + Both parameters default to ``writer``, so leaving them independent made the request + unsatisfiable and burned the whole retry window on every connection. + """ + remote_writer = HostInfo( + host="my-cluster-inst-9.XYZ.eu-west-1.rds.amazonaws.com", port=5432, role=HostRole.WRITER) + plugin, svc, driver_dialect = _build( + all_hosts=(remote_writer,), + role=HostRole.READER, + props_overrides={"inactive_cluster_writer_endpoint_substitution_role": "none"}, + ) + + expected = MagicMock(name="endpoint_conn") + calls = 0 + + async def _connect_func(): + nonlocal calls + calls += 1 + return expected + + async def _run(): + return await plugin.connect( + target_driver_func=MagicMock(), + driver_dialect=driver_dialect, + host_info=_cluster_host_info(_WRITER_CLUSTER), + props=svc.props, + is_initial_connection=True, + connect_func=_connect_func, + ) + + result = asyncio.run(_run()) + + assert result is expected + assert calls == 1, "the connection must not be retried to the timeout" + svc.get_host_role.assert_not_awaited() diff --git a/tests/unit/test_aio_plugin_service_extras.py b/tests/unit/test_aio_plugin_service_extras.py index 3e013a65e..956549c96 100644 --- a/tests/unit/test_aio_plugin_service_extras.py +++ b/tests/unit/test_aio_plugin_service_extras.py @@ -72,17 +72,46 @@ def test_get_availability_returns_none_when_unset(): def test_set_then_get_availability(): + """The mark lands on the live host and is cached under the host's url. + + The url is the key because that is what ``_update_host_availability`` looks up when a + topology query replaces the host objects; an alias-keyed entry was never read back. + """ svc = _make_service() - svc.set_availability(frozenset({"host-1.cluster.example"}), HostAvailability.UNAVAILABLE) - assert svc.get_availability("host-1.cluster.example") == HostAvailability.UNAVAILABLE + host = HostInfo("host-1.cluster.example", 5432) + svc._all_hosts = (host,) + + # as_aliases() is what every caller passes, and it spells the host as "host:port". + svc.set_availability(host.as_aliases(), HostAvailability.UNAVAILABLE) + + assert host.get_raw_availability() == HostAvailability.UNAVAILABLE + assert svc.get_availability(host.url) == HostAvailability.UNAVAILABLE -def test_set_availability_covers_all_aliases(): +def test_set_availability_matches_a_host_through_any_of_its_aliases(): svc = _make_service() - aliases = frozenset({"h1.example", "h1-alias.example"}) - svc.set_availability(aliases, HostAvailability.UNAVAILABLE) - for alias in aliases: - assert svc.get_availability(alias) == HostAvailability.UNAVAILABLE + host = HostInfo("h1.example", 5432) + host.add_alias("h1-alias.example") + svc._all_hosts = (host,) + + svc.set_availability(frozenset({"h1-alias.example"}), HostAvailability.UNAVAILABLE) + + assert host.get_raw_availability() == HostAvailability.UNAVAILABLE + assert svc.get_availability(host.url) == HostAvailability.UNAVAILABLE + + +def test_set_availability_ignores_a_host_outside_the_topology(): + """A mark for a host the service does not hold is dropped rather than cached. + + Otherwise a stale entry could be re-hydrated onto an unrelated host that later joins the + topology under the same address. + """ + svc = _make_service() + svc._all_hosts = () + + svc.set_availability(frozenset({"not-in-topology.example"}), HostAvailability.UNAVAILABLE) + + assert svc.get_availability("not-in-topology.example/") is None def test_default_plugin_accepts_random_strategy(): diff --git a/tests/unit/test_aurora_initial_connection_strategy_plugin.py b/tests/unit/test_aurora_initial_connection_strategy_plugin.py index 08bf92e3a..019541344 100644 --- a/tests/unit/test_aurora_initial_connection_strategy_plugin.py +++ b/tests/unit/test_aurora_initial_connection_strategy_plugin.py @@ -45,7 +45,9 @@ def _plugin(props, all_hosts=()): def test_retry_deadline_uses_timeout_property(): props = Properties({}) - # Zero total budget: the retry loop must not run even once despite the long interval. + # A zero total budget means "do not retry", not "do not connect". The budget is checked + # after an attempt rather than before one, so exactly one attempt is made however long + # the retry interval is -- and the interval is never slept, because there is no retry. WrapperProperties.OPEN_CONNECTION_RETRY_TIMEOUT_MS.set(props, "0") WrapperProperties.OPEN_CONNECTION_RETRY_INTERVAL_MS.set(props, "600000") @@ -56,7 +58,7 @@ def test_retry_deadline_uses_timeout_property(): plugin.connect( MagicMock(), MagicMock(), HostInfo(WRITER_CLUSTER), props, True, connect_func) - connect_func.assert_not_called() + assert connect_func.call_count == 1 def test_wait_for_initial_topology_disabled_by_default(): @@ -112,7 +114,10 @@ def populate_topology(*_args, **_kwargs): 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) + # Deliberately False: asking the monitor to verify the writer closes the shared + # per-cluster monitoring connection, which under a pool prefill means every arriving + # connection tears down what the previous one established. + plugin_service.force_monitoring_refresh_host_list.assert_called_once_with(False, 5.0) fallback_conn.close.assert_called_once()