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
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,8 @@

from enum import Enum
from time import perf_counter_ns, sleep
from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Set, Tuple
from typing import (TYPE_CHECKING, Callable, Dict, FrozenSet, List, Optional,
Sequence, Set, Tuple)

if TYPE_CHECKING:
from aws_advanced_python_wrapper.driver_dialect import DriverDialect
Expand All @@ -29,6 +30,8 @@
from aws_advanced_python_wrapper.host_availability import HostAvailability
from aws_advanced_python_wrapper.hostinfo import HostInfo, HostRole
from aws_advanced_python_wrapper.plugin import Plugin, PluginFactory
from aws_advanced_python_wrapper.utils.accessible_regions import \
parse as parse_accessible_regions
from aws_advanced_python_wrapper.utils.log import Logger
from aws_advanced_python_wrapper.utils.messages import Messages
from aws_advanced_python_wrapper.utils.properties import (Properties,
Expand Down Expand Up @@ -106,6 +109,7 @@ def __init__(self, plugin_service: PluginService, props: Properties):
self._plugin_service: PluginService = plugin_service
self._rds_utils = RdsUtils()
self._host_list_provider_service: Optional[HostListProviderService] = None
self._accessible_regions: Optional[FrozenSet[str]] = parse_accessible_regions(props)

self._retry_delay_ms: int = WrapperProperties.OPEN_CONNECTION_RETRY_INTERVAL_MS.get_int(props)
self._open_connection_retry_timeout_ns: int = \
Expand Down Expand Up @@ -277,7 +281,7 @@ def _wait_for_topology_and_connect_to_instance(
"AuroraInitialConnectionStrategyPlugin.WaitingForTopology",
self._wait_for_initial_topology_ms, original_connect_host.host)

# Deviation from JDBC: force_monitoring_refresh_host_list takes seconds, and host list
# force_monitoring_refresh_host_list takes seconds, and host list
# providers without monitor support raise instead of returning their host list.
timeout_sec = self._wait_for_initial_topology_ms / 1000
try:
Expand Down Expand Up @@ -326,7 +330,7 @@ def _get_instance_substitution_strategy(
return InstanceSubstitutionStrategy.SUBSTITUTE_WITH_WRITER

if url_type == RdsUrlType.RDS_WRITER_CLUSTER:
writer = self._get_writer()
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

Expand Down Expand Up @@ -405,7 +409,7 @@ def _get_role_to_verify(
return HostRole.WRITER

if url_type == RdsUrlType.RDS_WRITER_CLUSTER:
writer = self._get_writer()
writer = self._find_writer(self._plugin_service.all_hosts)
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.
Expand Down Expand Up @@ -446,7 +450,12 @@ def _get_candidate_host(
return original_connect_host

if substitution_strategy is InstanceSubstitutionStrategy.SUBSTITUTE_WITH_WRITER:
return self._get_writer()
# Filter by accessible regions BEFORE picking the writer so a writer in
# an unreachable region is never selected (no-op unless
# gdb_accessible_regions is set on a Global Aurora dialect); the
# candidate host is chosen from the filtered host list here.
available_hosts = self._filter_by_accessible_regions(self._plugin_service.all_hosts)
return self._find_writer(available_hosts)

# SUBSTITUTE_WITH_ANY has no specific target role, so to_target_role() returns None
target_role = substitution_strategy.to_target_role()
Expand All @@ -457,17 +466,23 @@ def _get_candidate_host(
"AuroraInitialConnectionStrategyPlugin.UnsupportedStrategy", self._selection_strategy))

try:
# Filter to accessible regions BEFORE any strategy/region selection so
# a candidate is never chosen from an unreachable region (no-op unless
# gdb_accessible_regions is set on a Global Aurora dialect).
available_hosts = self._filter_by_accessible_regions(self._plugin_service.hosts)

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
host for host in available_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)
return self._plugin_service.get_host_info_by_strategy(
target_role, self._selection_strategy, available_hosts)
except Exception:
# Unable to find a candidate host.
return None
Expand All @@ -479,8 +494,14 @@ def _set_initial_connection_host_info(
and host_info is not None):
self._host_list_provider_service.initial_connection_host_info = host_info

def _get_writer(self) -> Optional[HostInfo]:
for host in self._plugin_service.all_hosts:
@staticmethod
def _find_writer(hosts: Sequence[HostInfo]) -> Optional[HostInfo]:
"""Return the first WRITER in ``hosts``, or ``None``.

Does NOT filter by accessible regions — the caller decides whether to
pass an already-filtered list.
"""
for host in hosts:
if host.role == HostRole.WRITER:
return host
return None
Expand All @@ -502,6 +523,21 @@ def _close_connection(self, connection: Optional[Connection]):
def _delay(self, delay_ms: int):
sleep(delay_ms / 1000)

def _filter_by_accessible_regions(self, hosts: Sequence[HostInfo]) -> List[HostInfo]:
"""Filter hosts down to the configured ``gdb_accessible_regions``.

Returns the list unchanged when no accessible-regions restriction is
set. Filtering is delegated to the dialect's ``filter_available_hosts``
(a no-op default; Global Aurora dialects filter by region), so this is a
pass-through for non-Global clusters.
"""
if self._accessible_regions is None:
return list(hosts)
dialect = self._plugin_service.database_dialect
if dialect is None:
return list(hosts)
return dialect.filter_available_hosts(hosts, self._accessible_regions)


class AuroraInitialConnectionStrategyPluginFactory(PluginFactory):
@staticmethod
Expand Down
Loading
Loading