From 16070b52014c8230aedb34ac5611f2b8f9db67e6 Mon Sep 17 00:00:00 2001
From: karezche <64801825+karenc-bq@users.noreply.github.com>
Date: Tue, 4 Aug 2026 17:27:51 -0700
Subject: [PATCH] chore: update the async initial strategy plugin
---
...rora_initial_connection_strategy_plugin.py | 609 ++++++++++--------
.../aio/plugin_factory.py | 2 +-
...heAuroraInitialConnectionStrategyPlugin.md | 72 ++-
.../test_aio_aurora_initial_connection.py | 388 +++++++++--
4 files changed, 767 insertions(+), 304 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 cdfe7951f..6f42acd76 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
@@ -26,18 +26,22 @@
so auth plugins (IAM, Secrets, Federated, Okta) and connection
tracker re-apply on the new connection just like a user-driven
connect.
-* ``identify_connection`` approximated via :meth:`get_host_role` plus a
- topology scan (async ``PluginService`` doesn't yet expose
- ``identify_connection``).
+* ``initial_connection_host_info`` is set on the plugin service directly:
+ ``AsyncPluginServiceImpl`` implements ``AsyncHostListProviderService``
+ and ``aio/wrapper.py`` passes it as the host list provider service, so
+ the two are the same object.
"""
from __future__ import annotations
import asyncio
-from typing import (TYPE_CHECKING, Any, Awaitable, Callable, Optional, Set,
- Tuple)
+from contextlib import suppress
+from typing import (TYPE_CHECKING, Any, Awaitable, Callable, List, Optional,
+ Set, Tuple)
from aws_advanced_python_wrapper.aio.plugin import AsyncPlugin
+from aws_advanced_python_wrapper.aurora_initial_connection_strategy_plugin import (
+ InstanceSubstitutionStrategy, RoleVerificationSetting)
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
@@ -67,13 +71,28 @@ class AsyncAuroraInitialConnectionStrategyPlugin(AsyncPlugin):
DbApiMethod.CONNECT.method_name,
}
- _DEFAULT_RETRY_TIMEOUT_MS = 30000
- _DEFAULT_RETRY_INTERVAL_MS = 1000
-
- def __init__(self, plugin_service: AsyncPluginService) -> None:
+ def __init__(self, plugin_service: AsyncPluginService, props: Properties) -> None:
self._plugin_service = plugin_service
self._rds_utils = RdsUtils()
+ 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)
+
@property
def subscribed_methods(self) -> Set[str]:
return set(self._SUBSCRIBED)
@@ -83,9 +102,8 @@ def init_host_provider(
props: Properties,
host_list_provider_service: Any,
init_host_provider_func: Callable) -> None:
- # Sync parity (aurora_initial_connection_strategy_plugin.py:231-233):
- # capture the service and continue the chain.
- self._host_list_provider_service = host_list_provider_service
+ # Sync captures the service here; async sets initial_connection_host_info on the plugin
+ # service directly so there is nothing to hold.
init_host_provider_func()
async def connect(
@@ -96,279 +114,366 @@ async def connect(
props: Properties,
is_initial_connection: bool,
connect_func: Callable[..., Awaitable[Any]]) -> Any:
- url_type: RdsUrlType = self._rds_utils.identify_rds_type(host_info.host)
- if not url_type.is_rds_cluster:
- return await connect_func()
-
- if url_type in (RdsUrlType.RDS_WRITER_CLUSTER,
- RdsUrlType.RDS_GLOBAL_WRITER_CLUSTER):
- verified = await self._get_verified_writer(
- driver_dialect, host_info, props, connect_func,
- is_initial_connection, target_driver_func)
- if verified is not None:
- return verified
- return await connect_func()
+ 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)
- if url_type == RdsUrlType.RDS_READER_CLUSTER:
- verified = await self._get_verified_reader(
- driver_dialect, host_info, props, connect_func,
- is_initial_connection, target_driver_func)
- if verified is not None:
- return verified
- return await connect_func()
-
- return await connect_func()
-
- # ---- writer verification ----
-
- async def _get_verified_writer(
- self,
- driver_dialect: AsyncDriverDialect,
- host_info: HostInfo,
- props: Properties,
- connect_func: Callable[..., Awaitable[Any]],
- is_initial_connection: bool,
- target_driver_func: Callable) -> Optional[Any]:
- timeout_ms, interval_ms = self._retry_bounds(props)
loop = asyncio.get_running_loop()
- deadline = loop.time() + (timeout_ms / 1000.0)
+ deadline = loop.time() + (self._open_connection_retry_timeout_ns / 1_000_000_000)
while loop.time() < deadline:
candidate_conn: Optional[Any] = None
+ candidate_host: Optional[HostInfo] = None
+
try:
- writer_candidate = self._pick_writer()
- if (writer_candidate is None
- or self._rds_utils.is_rds_cluster_dns(writer_candidate.host)):
- # Topology is stale -- open via cluster endpoint, refresh, identify.
- candidate_conn = await connect_func()
- try:
- await self._plugin_service.force_refresh_host_list(candidate_conn)
- except Exception: # noqa: BLE001
- pass
- writer_candidate = await self._identify_host_role(
- candidate_conn, HostRole.WRITER)
- if writer_candidate is None:
- await self._close_best_effort(candidate_conn, driver_dialect)
- await asyncio.sleep(interval_ms / 1000.0)
- continue
- if is_initial_connection:
- self._plugin_service.initial_connection_host_info = writer_candidate
+ 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
- # Open directly to the writer instance.
- candidate_conn = await self._open_direct(
- driver_dialect, writer_candidate, props, target_driver_func)
- actual = await self._plugin_service.get_host_role(candidate_conn)
- if actual != HostRole.WRITER:
- try:
- await self._plugin_service.force_refresh_host_list(candidate_conn)
- except Exception: # noqa: BLE001
- pass
- await self._close_best_effort(candidate_conn, driver_dialect)
- await asyncio.sleep(interval_ms / 1000.0)
+ # 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_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)
+ await self._close_connection(candidate_conn, driver_dialect)
+ await self._delay(self._retry_delay_ms)
+ 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:
+ self._plugin_service.set_availability(
+ candidate_host.as_aliases(), HostAvailability.UNAVAILABLE)
+
+ if self._plugin_service.is_network_exception(error=e):
+ # Retry connection.
continue
- if is_initial_connection:
- self._plugin_service.initial_connection_host_info = writer_candidate
- return candidate_conn
- except Exception: # noqa: BLE001 - retry on any error
- await self._close_best_effort(candidate_conn, driver_dialect)
- await asyncio.sleep(interval_ms / 1000.0)
- return None
- # ---- reader verification ----
+ if (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
- async def _get_verified_reader(
+ raise
+
+ return await connect_func()
+
+ async def _open_candidate_connection(
self,
- driver_dialect: AsyncDriverDialect,
- host_info: HostInfo,
+ original_connect_host: HostInfo,
+ url_type: RdsUrlType,
+ substitution_strategy: InstanceSubstitutionStrategy,
props: Properties,
connect_func: Callable[..., Awaitable[Any]],
+ driver_dialect: AsyncDriverDialect) -> 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
+ 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, await 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. 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.
+ return candidate_host, await self._plugin_service.connect(
+ candidate_host, props, plugin_to_skip=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 = 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
+
+ return await self._wait_for_topology_and_connect_to_instance(
+ original_connect_host, url_type, substitution_strategy, props, candidate_conn, driver_dialect)
+
+ async def _wait_for_topology_and_connect_to_instance(
+ self,
+ original_connect_host: HostInfo,
+ url_type: RdsUrlType,
+ substitution_strategy: InstanceSubstitutionStrategy,
+ props: Properties,
+ fallback_conn: Any,
+ driver_dialect: AsyncDriverDialect) -> Tuple[HostInfo, Optional[Any]]:
+ """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)
+
+ # 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.
+ timeout_sec = self._wait_for_initial_topology_ms / 1000
+ topology_fetched = await self._plugin_service.force_monitoring_refresh_host_list(True, timeout_sec)
+
+ 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 = await self._plugin_service.connect(
+ instance_host, props, plugin_to_skip=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.
+ await self._close_connection(fallback_conn, driver_dialect)
+ return instance_host, instance_conn
+
+ # ---- strategy and role resolution (ported from sync) ----
+
+ def _get_instance_substitution_strategy(
+ self,
+ props: Properties,
+ url_type: RdsUrlType,
is_initial_connection: bool,
- target_driver_func: Callable) -> Optional[Any]:
- timeout_ms, interval_ms = self._retry_bounds(props)
- loop = asyncio.get_running_loop()
- deadline = loop.time() + (timeout_ms / 1000.0)
+ 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
- while loop.time() < deadline:
- candidate_conn: Optional[Any] = None
- reader_candidate: Optional[HostInfo] = None
- try:
- reader_candidate = self._pick_reader(props, host_info)
- if (reader_candidate is None
- or self._rds_utils.is_rds_cluster_dns(reader_candidate.host)):
- candidate_conn = await connect_func()
- try:
- await self._plugin_service.force_refresh_host_list(candidate_conn)
- except Exception: # noqa: BLE001
- pass
- actual_host = await self._identify_host_role(
- candidate_conn, HostRole.READER)
- if actual_host is None:
- if self._has_no_readers():
- # Cluster has no readers -- simulate Aurora reader cluster
- # endpoint and return the current writer connection.
- if is_initial_connection:
- self._plugin_service.initial_connection_host_info = \
- self._pick_writer() or host_info
- return candidate_conn
- await self._close_best_effort(candidate_conn, driver_dialect)
- await asyncio.sleep(interval_ms / 1000.0)
- continue
- if is_initial_connection:
- self._plugin_service.initial_connection_host_info = actual_host
- return candidate_conn
+ if url_type == RdsUrlType.RDS_READER_CLUSTER:
+ return InstanceSubstitutionStrategy.SUBSTITUTE_WITH_READER
- # Connect directly to the picked reader instance.
- candidate_conn = await self._open_direct(
- driver_dialect, reader_candidate, props, target_driver_func)
- actual = await self._plugin_service.get_host_role(candidate_conn)
- if actual != HostRole.READER:
- try:
- await self._plugin_service.force_refresh_host_list(candidate_conn)
- except Exception: # noqa: BLE001
- pass
- if self._has_no_readers():
- if is_initial_connection:
- self._plugin_service.initial_connection_host_info = reader_candidate
- return candidate_conn
- await self._close_best_effort(candidate_conn, driver_dialect)
- await asyncio.sleep(interval_ms / 1000.0)
- continue
- if is_initial_connection:
- self._plugin_service.initial_connection_host_info = reader_candidate
- return candidate_conn
- except AwsWrapperError:
- # Configuration errors (e.g., unsupported strategy) should
- # surface immediately rather than loop to exhaustion.
- await self._close_best_effort(candidate_conn, driver_dialect)
- raise
- except Exception as e: # noqa: BLE001
- await self._close_best_effort(candidate_conn, driver_dialect)
- # On non-login failure, mark reader UNAVAILABLE so the next
- # iteration picks a different one.
- if (reader_candidate is not None
- and not self._plugin_service.is_login_exception(error=e)):
- self._plugin_service.set_availability(
- reader_candidate.as_aliases(),
- HostAvailability.UNAVAILABLE)
- await asyncio.sleep(interval_ms / 1000.0)
- return None
+ return InstanceSubstitutionStrategy.DO_NOT_SUBSTITUTE
- # ---- helpers ----
+ def _validate_substitution_strategy(
+ self, setting: InstanceSubstitutionStrategy, url_type: RdsUrlType):
+ if setting is InstanceSubstitutionStrategy.DO_NOT_SUBSTITUTE:
+ return
- @staticmethod
- def _retry_bounds(props: Properties) -> Tuple[int, int]:
- timeout_ms = (
- WrapperProperties.OPEN_CONNECTION_RETRY_TIMEOUT_MS.get_int(props)
- or AsyncAuroraInitialConnectionStrategyPlugin._DEFAULT_RETRY_TIMEOUT_MS)
- interval_ms = (
- WrapperProperties.OPEN_CONNECTION_RETRY_INTERVAL_MS.get_int(props)
- or AsyncAuroraInitialConnectionStrategyPlugin._DEFAULT_RETRY_INTERVAL_MS)
- return timeout_ms, interval_ms
-
- def _pick_writer(self) -> Optional[HostInfo]:
- for h in self._plugin_service.all_hosts:
- if h.role == HostRole.WRITER:
- return h
- return None
+ if url_type == RdsUrlType.RDS_INSTANCE:
+ raise AwsWrapperError(Messages.get_formatted(
+ "AuroraInitialConnectionStrategyPlugin.InvalidSettingForInstanceEndpoint",
+ WrapperProperties.ENDPOINT_SUBSTITUTION_ROLE.name))
- def _pick_reader(
+ 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,
- connect_host: Optional[HostInfo] = None) -> Optional[HostInfo]:
- strategy = WrapperProperties.READER_INITIAL_HOST_SELECTOR_STRATEGY.get(props)
- if not strategy:
+ original_host: str) -> Optional[HostRole]:
+ if not is_initial_connection:
return None
- if not self._plugin_service.accepts_strategy(HostRole.READER, strategy):
- raise AwsWrapperError(Messages.get_formatted(
- "AuroraInitialConnectionStrategyPlugin.UnsupportedStrategy",
- strategy))
- try:
- readers = [h for h in self._plugin_service.all_hosts
- if h.role == HostRole.READER]
- readers = self._filter_readers_by_region(readers, connect_host)
- return self._plugin_service.get_host_info_by_strategy(
- HostRole.READER, strategy, readers)
- except Exception: # noqa: BLE001
+
+ 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
- def _filter_readers_by_region(
- self,
- readers: list,
- connect_host: Optional[HostInfo]) -> list:
- """Restrict reader candidates to the connect URL's region.
-
- Sync parity: aurora_initial_connection_strategy_plugin.py:210-224 --
- when the connect URL encodes an AWS region (e.g. a Global Database
- cluster endpoint), only topology hosts in that same region are
- eligible; without a region the full reader list is used. Sync reads
- the URL from ``plugin_service.current_host_info``, which at initial
- connect is the connect-URL host -- here the connect pipeline passes
- that host in directly.
- """
- if connect_host is None:
- return readers
- url_type = self._rds_utils.identify_rds_type(connect_host.host)
- if not url_type.has_region:
- return readers
- aws_region = self._rds_utils.get_rds_region(connect_host.host)
- if not aws_region:
- return readers
- return [h for h in readers
- if (self._rds_utils.get_rds_region(h.host) or "").lower()
- == aws_region.lower()]
-
- async def _identify_host_role(
+ if url_type == RdsUrlType.RDS_READER_CLUSTER:
+ return HostRole.READER
+
+ return None
+
+ 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,
- conn: Any,
- expected_role: HostRole) -> Optional[HostInfo]:
- """Probe ``conn``'s role; return the matching HostInfo from topology
- if role matches, else None."""
+ 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:
- actual = await self._plugin_service.get_host_role(conn)
- except Exception: # noqa: BLE001
- return None
- if actual != expected_role:
+ 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
- # Find a matching topology entry. Best-effort -- topology may not
- # contain the exact host we resolved to.
- for h in self._plugin_service.all_hosts:
- if h.role == expected_role:
- return h
+
+ 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_no_readers(self) -> bool:
- if not self._plugin_service.all_hosts:
- return False
- return not any(h.role == HostRole.READER
- for h in self._plugin_service.all_hosts)
+ def _has_hosts(self) -> bool:
+ return len(self._plugin_service.all_hosts) > 0
- async def _open_direct(
- self,
- driver_dialect: AsyncDriverDialect,
- host_info: HostInfo,
- props: Properties,
- target_driver_func: Callable) -> Any:
- new_props = Properties(dict(props))
- new_props["host"] = host_info.host
- if host_info.is_port_specified():
- new_props["port"] = str(host_info.port)
- # Routes through the plugin pipeline (skipping this plugin to
- # avoid recursion) so auth plugins (IAM, Secrets, Federated,
- # Okta) re-apply on the new instance connection.
- return await self._plugin_service.connect(
- host_info, new_props, plugin_to_skip=self)
+ def _has_readers(self) -> bool:
+ return any(host.role == HostRole.READER for host in self._plugin_service.all_hosts)
- @staticmethod
- async def _close_best_effort(
- conn: Optional[Any],
- driver_dialect: AsyncDriverDialect) -> None:
- if conn is None:
+ def _set_initial_connection_host_info(
+ self, is_initial_connection: bool, host_info: Optional[HostInfo]):
+ # AsyncPluginServiceImpl implements AsyncHostListProviderService and is passed as the host
+ # list provider service, so this is the same attribute sync writes through that service.
+ if is_initial_connection and host_info is not None:
+ self._plugin_service.initial_connection_host_info = host_info
+
+ async def _close_connection(
+ self, connection: Optional[Any], driver_dialect: AsyncDriverDialect) -> None:
+ if connection is None:
return
- try:
- await driver_dialect.abort_connection(conn)
- except Exception: # noqa: BLE001
- pass
+ with suppress(Exception):
+ await driver_dialect.abort_connection(connection)
+
+ @staticmethod
+ async def _delay(delay_ms: int) -> None:
+ await asyncio.sleep(delay_ms / 1000)
__all__ = ["AsyncAuroraInitialConnectionStrategyPlugin"]
diff --git a/aws_advanced_python_wrapper/aio/plugin_factory.py b/aws_advanced_python_wrapper/aio/plugin_factory.py
index bb4a695d9..41c16df97 100644
--- a/aws_advanced_python_wrapper/aio/plugin_factory.py
+++ b/aws_advanced_python_wrapper/aio/plugin_factory.py
@@ -219,7 +219,7 @@ def get_instance(
# Local import keeps module load-order cheap.
from aws_advanced_python_wrapper.aio.aurora_initial_connection_strategy_plugin import \
AsyncAuroraInitialConnectionStrategyPlugin
- return AsyncAuroraInitialConnectionStrategyPlugin(plugin_service)
+ return AsyncAuroraInitialConnectionStrategyPlugin(plugin_service, props)
class _FastestResponseFactory:
diff --git a/docs/using-the-python-wrapper/using-plugins/UsingTheAuroraInitialConnectionStrategyPlugin.md b/docs/using-the-python-wrapper/using-plugins/UsingTheAuroraInitialConnectionStrategyPlugin.md
index 81194b07e..92d1776e6 100644
--- a/docs/using-the-python-wrapper/using-plugins/UsingTheAuroraInitialConnectionStrategyPlugin.md
+++ b/docs/using-the-python-wrapper/using-plugins/UsingTheAuroraInitialConnectionStrategyPlugin.md
@@ -6,13 +6,16 @@ The following sequence diagram describes the default plugin behaviour if no cust
The AWS Advanced Python Wrapper may retry the connection attempts multiple times until it is able to connect to a valid reader instance or a valid writer instance.
You can configure how often to retry a connection and the maximum allowed time to obtain a connection using the `open_connection_retry_interval_ms` and the `open_connection_retry_timeout_ms` parameters respectively.
-When this plugin is enabled, if the initial connection is to a reader cluster endpoint, the connected reader host will be chosen based on selection strategy specified using the `reader_initial_connection_host_selector_strategy` parameter.
+When this plugin is enabled, if the initial connection is to a reader cluster or custom cluster endpoint, the connected host will be chosen based on the configured selection strategy specified using the `initial_connection_host_selector_strategy` parameter.
See [initial connection strategy](../ReaderSelectionStrategies.md) for all possible strategies.
This plugin also helps retrieve connections more reliably. When a user connects to a cluster endpoint, the actual instance for a new connection is resolved by DNS.
During failover, the cluster elects another instance to be the writer. While DNS is updating, which can take up to 40-60 seconds, if a user tries to connect to the cluster endpoint, they may be connecting to an old host.
This plugin helps by replacing the out of date endpoint if DNS is updating.
+When using Aurora Global Database, the user has an option to use an [Aurora Global Writer Endpoint](https://docs.aws.amazon.com/AmazonRDS/latest/AuroraUserGuide/aurora-global-database-connecting.html).
+The Global Writer Endpoint makes application configuration easier, but like the cluster writer endpoint it can be affected by DNS updates. The plugin recognizes an Aurora Global Writer Endpoint and substitutes it with the current writer endpoint.
+
## Enabling the Aurora Initial Connection Strategy Plugin
To enable the Aurora Initial Connection Strategy Plugin, add `initial_connection` to the [`plugins`](../UsingThePythonWrapper.md#connection-plugin-manager-parameters) value.
@@ -21,8 +24,65 @@ To enable the Aurora Initial Connection Strategy Plugin, add `initial_connection
The following properties can be used to configure the Aurora Initial Connection Strategy Plugin.
-| Parameter | Value | Required | Description | Example | Default Value |
-|----------------------------------------------------|:-------:|:--------:|--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|--------------------|---------------|
-| `reader_initial_connection_host_selector_strategy` | String | No | The strategy that will be used to select a new reader host when opening a new connection.
For more information on the available reader selection strategies, see this [table](../ReaderSelectionStrategies.md). | `leastConnections` | `random` |
-| `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` |
+| Parameter | Value | Required | Description | Example | Default Value |
+|-----------------------------------------------------------|:-------:|:--------:|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|--------------------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
+| `initial_connection_host_selector_strategy` | String | No | The strategy that will be used to select a host when opening a new connection. A host will be selected according to the host role implied by `endpoint_substitution_role`.
For more information on the available selection strategies, see this [table](../ReaderSelectionStrategies.md). | `least_connections`| `random` |
+| ~~`reader_initial_connection_host_selector_strategy`~~ | String | No | **Deprecated. Use `initial_connection_host_selector_strategy` instead.** During the migration period, the value of `reader_initial_connection_host_selector_strategy` is used only when `initial_connection_host_selector_strategy` is omitted. | `least_connections`| `random` |
+| `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` |
+| `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` |
+
+### Valid setting/endpoint combinations
+
+`endpoint_substitution_role` and `verify_opened_connection_type` accept different values depending on the endpoint the connection URL resolves to:
+
+| Endpoint type | `endpoint_substitution_role` | `verify_opened_connection_type` |
+|----------------------------------|------------------------------|---------------------------------|
+| Writer cluster / global writer | `writer`, `none` | `writer`, `none` |
+| Reader cluster | `reader`, `none` | `reader`, `none` |
+| Custom cluster | `reader`, `none` | `reader`, `none` |
+| Instance | `none` only | `none` only |
+
+Setting a value outside the valid set for the given endpoint raises an error.
+
+> **Note:** `endpoint_substitution_role=any` is accepted only for a custom cluster endpoint, but host selection for the `any` role is not currently supported — it raises an unsupported-strategy error at connect time. Use `reader` or `none` for custom cluster endpoints.
+
+## Examples
+
+Disable endpoint URL substitution. By default the plugin substitutes reader cluster URLs with a reader instance URL; setting the role to `none` removes this behavior:
+
+```python
+params = {
+ "plugins": "initial_connection",
+ "host": "mydb.cluster-ro-XYZ.us-east-1.rds.amazonaws.com",
+ "endpoint_substitution_role": "none",
+}
+conn = AwsWrapperConnection.connect(psycopg.Connection.connect, **params)
+```
+
+Disable reader host verification when using a reader cluster URL. By default the plugin verifies that the opened connection landed on a reader:
+
+```python
+params = {
+ "plugins": "initial_connection",
+ "host": "mydb.cluster-ro-XYZ.us-east-1.rds.amazonaws.com",
+ "verify_opened_connection_type": "none",
+}
+conn = AwsWrapperConnection.connect(psycopg.Connection.connect, **params)
+```
+
+Distribute concurrent and connection-pool prefill connections across readers. `wait_for_initial_topology_ms` blocks the first connections until the cluster topology is available, so `round_robin` can spread them across instances instead of all falling back to the single DNS-resolved reader cluster endpoint:
+
+```python
+params = {
+ "plugins": "initial_connection",
+ "host": "mydb.cluster-ro-XYZ.us-east-1.rds.amazonaws.com",
+ "initial_connection_host_selector_strategy": "round_robin",
+ "wait_for_initial_topology_ms": "30000",
+}
+conn = AwsWrapperConnection.connect(psycopg.Connection.connect, **params)
+```
diff --git a/tests/unit/test_aio_aurora_initial_connection.py b/tests/unit/test_aio_aurora_initial_connection.py
index 4a2d47eae..052f2eb01 100644
--- a/tests/unit/test_aio_aurora_initial_connection.py
+++ b/tests/unit/test_aio_aurora_initial_connection.py
@@ -12,20 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-"""Unit tests for :class:`AsyncAuroraInitialConnectionStrategyPlugin`.
-
-Covers the load-bearing branches ported from sync:
-
-1. Non-RDS-cluster URL passes through with no verification.
-2. Writer cluster URL + already connected to writer -> returns original conn.
-3. Writer cluster URL + connected to reader -> retries via writer instance.
-4. Reader cluster URL + connected to reader -> returns conn.
-5. Reader cluster URL + no readers in topology -> returns writer conn
- (simulated Aurora reader fallback).
-6. Timeout exhausted -> returns None -> falls back to plain connect_func.
-7. READER_INITIAL_HOST_SELECTOR_STRATEGY not supported -> raises AwsWrapperError.
-8. Reader probe failure on non-login exception -> marks host UNAVAILABLE.
-"""
+"""Unit tests for :class:`AsyncAuroraInitialConnectionStrategyPlugin`."""
from __future__ import annotations
@@ -39,10 +26,14 @@
AsyncAuroraInitialConnectionStrategyPlugin
from aws_advanced_python_wrapper.aio.plugin_service import \
AsyncPluginServiceImpl
+from aws_advanced_python_wrapper.aurora_initial_connection_strategy_plugin import \
+ InstanceSubstitutionStrategy
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
from aws_advanced_python_wrapper.utils.properties import Properties
+from aws_advanced_python_wrapper.utils.rds_url_type import RdsUrlType
+from aws_advanced_python_wrapper.utils.rds_utils import RdsUtils
# ---- Helpers -----------------------------------------------------------
@@ -106,7 +97,7 @@ def _build(
# over unchanged.
svc.connect = driver_dialect.connect # type: ignore[method-assign]
- plugin = AsyncAuroraInitialConnectionStrategyPlugin(svc)
+ plugin = AsyncAuroraInitialConnectionStrategyPlugin(svc, props)
return plugin, svc, driver_dialect
@@ -140,8 +131,9 @@ async def _run():
# No verification path taken.
driver_dialect.connect.assert_not_awaited()
svc.get_host_role.assert_not_awaited()
- # initial_connection_host_info untouched (plugin didn't set it).
- assert svc.initial_connection_host_info is None
+ # A non-cluster URL needs no role verification, so the loop records the
+ # original host and returns on the first pass.
+ assert svc.initial_connection_host_info is host
# ---- 2. Writer cluster URL + already writer -> direct writer conn ------
@@ -258,7 +250,7 @@ async def _run():
def test_reader_cluster_no_readers_returns_writer_fallback():
writer = _writer_host()
- # Topology has only a writer. _pick_reader returns None
+ # Topology has only a writer. _get_candidate_host returns None
# (strategy_pick=None), so the plugin falls through the "topology
# stale" branch, opens via connect_func, probes, and since no
# readers exist it returns that connection unmodified.
@@ -286,9 +278,9 @@ async def _run():
result = asyncio.run(_run())
assert result is cluster_conn
- # No-readers fallback sets initial_connection_host_info to the
- # writer (via _pick_writer).
- assert svc.initial_connection_host_info is writer
+ # The no-readers fallback records the candidate host, which on the stale
+ # path is the original connect host rather than the writer.
+ assert svc.initial_connection_host_info is host
# ---- 6. Timeout exhausted -> falls back to plain connect_func ----------
@@ -360,7 +352,9 @@ async def _run():
# ---- 8. Reader non-login exception -> host marked UNAVAILABLE ----------
-def test_reader_non_login_exception_marks_unavailable():
+def test_network_exception_retries_then_falls_back():
+ """A network exception is retried until the budget expires, then the plain
+ connect_func result is returned."""
reader = _reader_host()
writer = _writer_host()
plugin, svc, driver_dialect = _build(
@@ -368,10 +362,8 @@ def test_reader_non_login_exception_marks_unavailable():
role=HostRole.READER,
strategy_pick=reader,
)
- # driver_dialect.connect raises on every direct attempt -> reader
- # candidate gets marked UNAVAILABLE each iteration until the loop
- # exits. is_login_exception is already mocked to return False.
driver_dialect.connect.side_effect = RuntimeError("network-down")
+ svc.is_network_exception = MagicMock(return_value=True) # type: ignore[method-assign]
host = _cluster_host_info(_READER_CLUSTER)
fallback_conn = MagicMock(name="fallback_conn")
@@ -393,7 +385,43 @@ async def _run():
# Timeout expired -> plain connect_func fallback.
assert result is fallback_conn
- # At least one UNAVAILABLE mark was written for the picked reader.
+ # Retried rather than giving up on the first network failure.
+ assert driver_dialect.connect.await_count > 1
+
+
+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."""
+ reader = _reader_host()
+ plugin, svc, driver_dialect = _build(
+ all_hosts=(_writer_host(), reader),
+ role=HostRole.READER,
+ strategy_pick=reader,
+ )
+ # Candidate connect succeeds; the role probe then fails with a network error.
+ svc.connect = AsyncMock(return_value=MagicMock(name="instance_conn")) # type: ignore[method-assign]
+ svc.get_host_role = AsyncMock( # type: ignore[method-assign]
+ side_effect=RuntimeError("probe failed"))
+ svc.is_network_exception = MagicMock(return_value=True) # type: ignore[method-assign]
+
+ host = _cluster_host_info(_READER_CLUSTER)
+
+ async def _connect_func():
+ return MagicMock(name="fallback_conn")
+
+ async def _run():
+ return await plugin.connect(
+ target_driver_func=MagicMock(),
+ driver_dialect=driver_dialect,
+ host_info=host,
+ props=svc.props,
+ is_initial_connection=True,
+ connect_func=_connect_func,
+ )
+
+ asyncio.run(_run())
+
assert svc.set_availability.call_count >= 1
called_aliases = svc.set_availability.call_args_list[0][0][0]
assert reader.host in "".join(called_aliases)
@@ -401,6 +429,135 @@ async def _run():
HostAvailability.UNAVAILABLE
+def test_login_exception_raises_without_retrying():
+ """A login failure raises on the first attempt rather than being retried
+ to exhaustion."""
+ reader = _reader_host()
+ plugin, svc, driver_dialect = _build(
+ all_hosts=(_writer_host(), reader),
+ role=HostRole.READER,
+ strategy_pick=reader,
+ )
+ driver_dialect.connect.side_effect = RuntimeError("bad password")
+ svc.is_login_exception = MagicMock(return_value=True) # type: ignore[method-assign]
+
+ host = _cluster_host_info(_READER_CLUSTER)
+
+ async def _connect_func():
+ return MagicMock(name="fallback_conn")
+
+ async def _run():
+ return await plugin.connect(
+ target_driver_func=MagicMock(),
+ driver_dialect=driver_dialect,
+ host_info=host,
+ props=svc.props,
+ is_initial_connection=True,
+ connect_func=_connect_func,
+ )
+
+ with pytest.raises(RuntimeError, match="bad password"):
+ asyncio.run(_run())
+
+ # Exactly one attempt -- no retry loop on a credentials failure.
+ assert driver_dialect.connect.await_count == 1
+
+
+def test_unclassified_exception_propagates():
+ """An error that is neither login, network, nor read-only surfaces rather
+ than being swallowed."""
+ reader = _reader_host()
+ plugin, svc, driver_dialect = _build(
+ all_hosts=(_writer_host(), reader),
+ role=HostRole.READER,
+ strategy_pick=reader,
+ )
+ driver_dialect.connect.side_effect = RuntimeError("something unexpected")
+
+ host = _cluster_host_info(_READER_CLUSTER)
+
+ async def _connect_func():
+ return MagicMock(name="fallback_conn")
+
+ async def _run():
+ return await plugin.connect(
+ target_driver_func=MagicMock(),
+ driver_dialect=driver_dialect,
+ host_info=host,
+ props=svc.props,
+ is_initial_connection=True,
+ connect_func=_connect_func,
+ )
+
+ with pytest.raises(RuntimeError, match="something unexpected"):
+ asyncio.run(_run())
+
+
+def test_instance_connect_skips_this_plugin():
+ """Regression: the instance connect must skip this plugin so the pipeline
+ does not re-enter it. Async has always done this; the rewrite must keep it
+ (sync still omits it at both call sites)."""
+ writer = _writer_host()
+ plugin, svc, driver_dialect = _build(
+ all_hosts=(writer,), role=HostRole.WRITER, strategy_pick=writer)
+ svc.connect = AsyncMock(return_value=MagicMock(name="instance_conn")) # type: ignore[method-assign]
+
+ host = _cluster_host_info(_WRITER_CLUSTER)
+
+ async def _connect_func():
+ return MagicMock(name="cluster_conn")
+
+ async def _run():
+ return await plugin.connect(
+ target_driver_func=MagicMock(),
+ driver_dialect=driver_dialect,
+ host_info=host,
+ props=svc.props,
+ is_initial_connection=True,
+ connect_func=_connect_func,
+ )
+
+ asyncio.run(_run())
+
+ svc.connect.assert_awaited()
+ assert svc.connect.await_args.kwargs["plugin_to_skip"] is plugin
+
+
+def test_stale_topology_uses_is_rds_instance_not_cluster_dns():
+ """A custom-cluster host in the topology is not an instance endpoint, so
+ the plugin falls back to the initial endpoint rather than connecting
+ directly to it."""
+ custom_cluster_host = HostInfo(
+ host="my-cluster.cluster-custom-XYZ.us-east-1.rds.amazonaws.com",
+ port=5432, role=HostRole.WRITER)
+ plugin, svc, driver_dialect = _build(
+ all_hosts=(custom_cluster_host,), role=HostRole.WRITER)
+ svc.connect = AsyncMock(name="should_not_be_used") # type: ignore[method-assign]
+ cluster_conn = MagicMock(name="cluster_conn")
+
+ host = _cluster_host_info(_WRITER_CLUSTER)
+
+ async def _connect_func():
+ return cluster_conn
+
+ async def _run():
+ return await plugin.connect(
+ target_driver_func=MagicMock(),
+ driver_dialect=driver_dialect,
+ host_info=host,
+ props=svc.props,
+ is_initial_connection=True,
+ connect_func=_connect_func,
+ )
+
+ result = asyncio.run(_run())
+
+ # The cluster endpoint connection is used; the custom-cluster topology host
+ # is never treated as an instance to connect directly to.
+ assert result is cluster_conn
+ svc.connect.assert_not_awaited()
+
+
# ---- 9. E3: region-aware reader filtering -------------------------------
@@ -414,9 +571,8 @@ def _other_region_reader_host() -> HostInfo:
def test_reader_candidates_restricted_to_connect_url_region():
- """E3: sync parity (aurora_initial_connection_strategy_plugin.py:210-224)
- -- when the connect URL encodes a region, only readers in that region
- are offered to the selection strategy."""
+ """When the connect URL encodes a region, only readers in that region are
+ offered to the selection strategy."""
writer = _writer_host()
in_region_reader = _reader_host() # us-east-1
out_of_region_reader = _other_region_reader_host() # eu-west-1
@@ -455,26 +611,168 @@ async def _run():
assert out_of_region_reader not in candidate_list
-def test_filter_readers_by_region_no_connect_host_keeps_all():
- plugin, svc, _ = _build()
- readers = [_reader_host(), _other_region_reader_host()]
- assert plugin._filter_readers_by_region(readers, None) == readers
-
-
-def test_filter_readers_by_region_keeps_all_when_no_region_in_url():
+def test_candidate_host_keeps_all_when_no_region_in_url():
"""A connect URL without a region (e.g. a bare hostname) must not
- restrict the reader candidates."""
- plugin, svc, _ = _build()
- readers = [_reader_host(), _other_region_reader_host()]
+ restrict the candidates -- the unfiltered selector call is used."""
+ in_region = _reader_host()
+ out_of_region = _other_region_reader_host()
+ plugin, svc, _ = _build(all_hosts=(in_region, out_of_region),
+ strategy_pick=in_region)
no_region_host = HostInfo(host="some-random.example.com", port=5432)
- assert plugin._filter_readers_by_region(readers, no_region_host) == readers
+
+ plugin._get_candidate_host(
+ no_region_host,
+ RdsUtils().identify_rds_type(no_region_host.host),
+ InstanceSubstitutionStrategy.SUBSTITUTE_WITH_READER)
+
+ # No host_list argument => selector sees the full topology.
+ assert svc.get_host_info_by_strategy.call_count == 1
+ assert len(svc.get_host_info_by_strategy.call_args_list[0].args) == 2
-def test_filter_readers_by_region_filters_cross_region_readers():
- plugin, svc, _ = _build()
+def test_candidate_host_filters_cross_region_readers():
in_region = _reader_host()
out_of_region = _other_region_reader_host()
+ plugin, svc, _ = _build(all_hosts=(in_region, out_of_region),
+ strategy_pick=in_region)
connect_host = _cluster_host_info(_READER_CLUSTER) # us-east-1
- filtered = plugin._filter_readers_by_region(
- [in_region, out_of_region], connect_host)
- assert filtered == [in_region]
+
+ plugin._get_candidate_host(
+ connect_host,
+ RdsUrlType.RDS_READER_CLUSTER,
+ InstanceSubstitutionStrategy.SUBSTITUTE_WITH_READER)
+
+ assert svc.get_host_info_by_strategy.call_count == 1
+ candidate_list = svc.get_host_info_by_strategy.call_args_list[0].args[2]
+ assert in_region in candidate_list
+ assert out_of_region not in candidate_list
+
+
+def test_candidate_host_substitute_with_writer_ignores_strategy():
+ """SUBSTITUTE_WITH_WRITER short-circuits to the topology writer without
+ consulting the host selection strategy at all."""
+ writer = _writer_host()
+ plugin, svc, _ = _build(all_hosts=(writer, _reader_host()))
+
+ result = plugin._get_candidate_host(
+ _cluster_host_info(_WRITER_CLUSTER),
+ RdsUrlType.RDS_WRITER_CLUSTER,
+ InstanceSubstitutionStrategy.SUBSTITUTE_WITH_WRITER)
+
+ assert result == writer
+ svc.get_host_info_by_strategy.assert_not_called()
+
+
+def test_candidate_host_substitute_with_any_raises_unsupported_strategy():
+ """SUBSTITUTE_WITH_ANY has no target role, so it raises unsupportedStrategy
+ rather than being coerced to a reader-only selection."""
+ reader = _reader_host()
+ plugin, svc, _ = _build(all_hosts=(reader,))
+
+ with pytest.raises(AwsWrapperError):
+ plugin._get_candidate_host(
+ _cluster_host_info(
+ "my-cluster.cluster-custom-XYZ.us-east-1.rds.amazonaws.com"),
+ RdsUrlType.RDS_CUSTOM_CLUSTER,
+ InstanceSubstitutionStrategy.SUBSTITUTE_WITH_ANY)
+
+ svc.get_host_info_by_strategy.assert_not_called()
+
+def test_endpoint_substitution_role_on_instance_endpoint_raises():
+ """Substitution cannot be requested for an instance endpoint."""
+ props_overrides = {"endpoint_substitution_role": "writer"}
+ plugin, _, _ = _build(all_hosts=(_writer_host(),),
+ props_overrides=props_overrides)
+ props = Properties(props_overrides)
+
+ with pytest.raises(AwsWrapperError):
+ plugin._get_instance_substitution_strategy(
+ props, RdsUrlType.RDS_INSTANCE, True, _WRITER_INSTANCE)
+
+
+def test_invalid_verify_opened_connection_type_raises():
+ """A typo in verify_opened_connection_type must not be swallowed."""
+ plugin, _, _ = _build(
+ all_hosts=(_writer_host(),),
+ props_overrides={"verify_opened_connection_type": "bogus_value"})
+
+ with pytest.raises(AwsWrapperError):
+ plugin._get_role_to_verify(
+ RdsUrlType.RDS_WRITER_CLUSTER, True, Properties({}), _WRITER_CLUSTER)
+
+
+def test_verify_reader_on_writer_cluster_raises():
+ """Verifying 'reader' against a writer cluster endpoint is invalid."""
+ plugin, _, _ = _build(
+ all_hosts=(_writer_host(),),
+ props_overrides={"verify_opened_connection_type": "reader"})
+
+ with pytest.raises(AwsWrapperError):
+ plugin._get_role_to_verify(
+ RdsUrlType.RDS_WRITER_CLUSTER, True, Properties({}), _WRITER_CLUSTER)
+
+
+def test_global_writer_cluster_substitutes_and_verifies_writer():
+ """A global writer cluster endpoint resolves to writer substitution and
+ writer verification, rather than falling through unhandled."""
+ plugin, _, _ = _build(all_hosts=(_writer_host(),))
+ props = Properties({})
+ global_host = "my-global.global-XYZ.global.rds.amazonaws.com"
+
+ strategy = plugin._get_instance_substitution_strategy(
+ props, RdsUrlType.RDS_GLOBAL_WRITER_CLUSTER, True, global_host)
+ role = plugin._get_role_to_verify(
+ RdsUrlType.RDS_GLOBAL_WRITER_CLUSTER, True, props, global_host)
+
+ assert strategy is InstanceSubstitutionStrategy.SUBSTITUTE_WITH_WRITER
+ assert role == HostRole.WRITER
+
+
+def test_inactive_cluster_writer_substitution_role_honored():
+ """When the cluster writer endpoint is in a different region than the
+ current writer (Aurora Global secondary), the inactive-writer setting
+ decides whether to substitute."""
+ # Writer lives in us-west-2; the connect URL is us-east-1.
+ out_of_region_writer = HostInfo(
+ host="my-cluster-inst-1.XYZ.us-west-2.rds.amazonaws.com",
+ port=5432, role=HostRole.WRITER)
+ plugin, _, _ = _build(all_hosts=(out_of_region_writer,))
+ props = Properties({"inactive_cluster_writer_endpoint_substitution_role": "none"})
+
+ strategy = plugin._get_instance_substitution_strategy(
+ props, RdsUrlType.RDS_WRITER_CLUSTER, True, _WRITER_CLUSTER)
+
+ assert strategy is InstanceSubstitutionStrategy.DO_NOT_SUBSTITUTE
+
+
+def test_initial_connection_host_selector_strategy_overrides_reader_variant():
+ """The non-deprecated property wins when explicitly set."""
+ plugin, _, _ = _build(props_overrides={
+ "reader_initial_connection_host_selector_strategy": "random",
+ "initial_connection_host_selector_strategy": "round_robin",
+ })
+ assert plugin._selection_strategy == "round_robin"
+
+
+def test_reader_strategy_falls_back_to_deprecated_property():
+ plugin, _, _ = _build(props_overrides={
+ "reader_initial_connection_host_selector_strategy": "least_connections",
+ })
+ assert plugin._selection_strategy == "least_connections"
+
+
+def test_retry_bounds_honor_explicit_zero():
+ """An explicit 0 for the retry bounds must be taken literally rather than
+ falling back to the default."""
+ plugin, _, _ = _build(props_overrides={
+ "open_connection_retry_timeout_ms": "0",
+ "open_connection_retry_interval_ms": "0",
+ })
+ assert plugin._open_connection_retry_timeout_ns == 0
+ assert plugin._retry_delay_ms == 0
+
+
+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