Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view

Large diffs are not rendered by default.

26 changes: 24 additions & 2 deletions aws_advanced_python_wrapper/aio/plugin_service.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down
12 changes: 12 additions & 0 deletions aws_advanced_python_wrapper/aio/wrapper.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand All @@ -65,6 +71,12 @@
"gdb_rw",
"custom_endpoint",
"aurora_connection_tracker",
"initial_connection",
"host_monitoring",
"host_monitoring_v2",
"stale_dns",
"fastest_response_strategy",
"bg",
})


Expand Down
270 changes: 200 additions & 70 deletions aws_advanced_python_wrapper/aurora_initial_connection_strategy_plugin.py

Large diffs are not rendered by default.

27 changes: 26 additions & 1 deletion aws_advanced_python_wrapper/plugin_service.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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). <br><br> 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.<br><br>`reader` for reader cluster endpoints.<br><br>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. <br><br> 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.<br><br>`reader` for reader cluster endpoints.<br><br>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`. <br><br> 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` |
Expand All @@ -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.

Expand Down
125 changes: 122 additions & 3 deletions tests/unit/test_aio_aurora_initial_connection.py
Original file line number Diff line number Diff line change
Expand Up @@ -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),
Expand Down Expand Up @@ -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()
Loading
Loading