chore: update aurora initial connection strategy plugin implementation

aws_advanced_python_wrapper/cluster_topology_monitor.py

@@ -52,11 +52,11 @@
 
 class ClusterTopologyMonitor(ABC):
     @abstractmethod
-    def force_refresh(self, should_verify_writer: bool, timeout_sec: int) -> Topology:
+    def force_refresh(self, should_verify_writer: bool, timeout_sec: float) -> Topology:
         pass
 
     @abstractmethod
-    def force_refresh_with_connection(self, connection: Connection, timeout_sec: int) -> Topology:
+    def force_refresh_with_connection(self, connection: Connection, timeout_sec: float) -> Topology:
         pass
 
     @property
@@ -133,7 +133,7 @@ def __init__(self, plugin_service: PluginService, topology_utils: TopologyUtils,
 
         self._start_monitoring()
 
-    def force_refresh(self, should_verify_writer: bool, timeout_sec: int) -> Topology:
+    def force_refresh(self, should_verify_writer: bool, timeout_sec: float) -> Topology:
         current_time_nano = time.time_ns()
         if (self._ignore_new_topology_requests_end_time_nano > 0 and
                 current_time_nano < self._ignore_new_topology_requests_end_time_nano):
@@ -149,12 +149,12 @@ def force_refresh(self, should_verify_writer: bool, timeout_sec: int) -> Topolog
         result = self._wait_till_topology_gets_updated(timeout_sec)
         return result
 
-    def force_refresh_with_connection(self, connection: Connection, timeout_sec: int) -> Topology:
+    def force_refresh_with_connection(self, connection: Connection, timeout_sec: float) -> Topology:
         if self._is_verified_writer_connection:
             return self._wait_till_topology_gets_updated(timeout_sec)
         return self._fetch_topology_and_update_cache(connection)
 
-    def _wait_till_topology_gets_updated(self, timeout_sec: int) -> Topology:
+    def _wait_till_topology_gets_updated(self, timeout_sec: float) -> Topology:
         current_hosts = self._get_stored_hosts()
 
         self._request_to_update_topology.set()

aws_advanced_python_wrapper/host_list_provider.py

@@ -73,7 +73,7 @@ def get_current_topology(self, connection: Connection, initial_host_info: HostIn
         """
         ...
 
-    def force_monitoring_refresh(self, should_verify_writer: bool, timeout_sec: int) -> Topology:
+    def force_monitoring_refresh(self, should_verify_writer: bool, timeout_sec: float) -> Topology:
         ...
 
     def get_cluster_id(self) -> str:
@@ -248,7 +248,7 @@ def _get_or_create_monitor(self) -> Optional[ClusterTopologyMonitor]:
             )
         )
 
-    def _force_refresh_monitor(self, should_verify_writer: bool, timeout_sec: int) -> Optional[Topology]:
+    def _force_refresh_monitor(self, should_verify_writer: bool, timeout_sec: float) -> Optional[Topology]:
         """Force refresh using monitor - matches Java's forceRefreshMonitor"""
         monitor = self._get_or_create_monitor()
         if monitor is None:
@@ -275,7 +275,7 @@ def get_current_topology(self, connection: Connection, initial_host_info: HostIn
             return hosts
         return ()
 
-    def force_monitoring_refresh(self, should_verify_writer: bool, timeout_sec: int) -> Topology:
+    def force_monitoring_refresh(self, should_verify_writer: bool, timeout_sec: float) -> Topology:
         """Public API for forcing monitor refresh"""
         self._initialize()
         hosts = self._force_refresh_monitor(should_verify_writer, timeout_sec)
@@ -363,7 +363,7 @@ def get_current_topology(self, connection: Connection, initial_host_info: HostIn
         self._initialize()
         return tuple(self._hosts)
 
-    def force_monitoring_refresh(self, should_verify_writer: bool, timeout_sec: int) -> Topology:
+    def force_monitoring_refresh(self, should_verify_writer: bool, timeout_sec: float) -> Topology:
         raise AwsWrapperError(
                 Messages.get_formatted("HostListProvider.ForceMonitoringRefreshUnsupported", "ConnectionStringHostListProvider"))
 

aws_advanced_python_wrapper/plugin_service.py

@@ -263,7 +263,7 @@ def refresh_host_list(self, connection: Optional[Connection] = None):
     def force_refresh_host_list(self, connection: Optional[Connection] = None):
         ...
 
-    def force_monitoring_refresh_host_list(self, should_verify_writer: bool, timeout_ms: int) -> bool:
+    def force_monitoring_refresh_host_list(self, should_verify_writer: bool, timeout_sec: float) -> bool:
         ...
 
     def connect(self, host_info: HostInfo, props: Properties, plugin_to_skip: Optional[Plugin] = None) -> Connection:
@@ -594,7 +594,7 @@ def force_refresh_host_list(self, connection: Optional[Connection] = None):
             self._update_host_availability(updated_host_list)
             self._update_hosts(updated_host_list)
 
-    def force_monitoring_refresh_host_list(self, should_verify_writer: bool, timeout_sec: int) -> bool:
+    def force_monitoring_refresh_host_list(self, should_verify_writer: bool, timeout_sec: float) -> bool:
         try:
             updated_host_list = self.host_list_provider.force_monitoring_refresh(should_verify_writer, timeout_sec)
             if updated_host_list is not None:

aws_advanced_python_wrapper/resources/aws_advanced_python_wrapper_messages.properties

@@ -17,8 +17,16 @@
 AuroraPgDialect.AuroraUtils=[AuroraPgDialect] aurora_utils: {}
 AuroraPgDialect.HasTopologyTrue=[AuroraPgDialect] has_topology: True
 
-AuroraInitialConnectionStrategyPlugin.RequireDynamicProvider=[AuroraInitialConnectionStrategyPlugin] Dynamic host list provider is required.
 AuroraInitialConnectionStrategyPlugin.UnsupportedStrategy=[AuroraInitialConnectionStrategyPlugin] Unsupported host selection strategy '{}'.
+AuroraInitialConnectionStrategyPlugin.InvalidSettingForEndpoint=[AuroraInitialConnectionStrategyPlugin] Parameter '{}' cannot be set to '{}' when using a {} endpoint. Please see the Aurora Initial Connection Strategy Plugin documentation for valid setting/endpoint combinations.
+AuroraInitialConnectionStrategyPlugin.InvalidSettingForInstanceEndpoint=[AuroraInitialConnectionStrategyPlugin] Parameter '{}' cannot be set when using an instance endpoint. Please see the Aurora Initial Connection Strategy Plugin documentation for valid setting/endpoint combinations.
+AuroraInitialConnectionStrategyPlugin.InvalidPropertyValue=[AuroraInitialConnectionStrategyPlugin] Received an invalid value for parameter '{}'. Received '{}', valid values are {}.
+AuroraInitialConnectionStrategyPlugin.Timeout=[AuroraInitialConnectionStrategyPlugin] The Aurora Initial Connection Strategy Plugin attempted to connect but timed out after {}ms. Please ensure that your URL is correct, there are no network issues, and you are connecting to the correct role if '{}' was set.
+AuroraInitialConnectionStrategyPlugin.IncorrectRole=[AuroraInitialConnectionStrategyPlugin] The connection opened to '{}' did not have the expected role '{}'. Retrying.
+AuroraInitialConnectionStrategyPlugin.VerifyReaderConfiguredButNoReadersExist=[AuroraInitialConnectionStrategyPlugin] Parameter '{}' was set to 'reader' but no readers were detected in the topology. The writer will be used as a fallback.
+AuroraInitialConnectionStrategyPlugin.WaitingForTopology=[AuroraInitialConnectionStrategyPlugin] Waiting up to {}ms for the cluster topology of '{}' to be fetched before opening a new connection.
+AuroraInitialConnectionStrategyPlugin.WaitForTopologyTimeout=[AuroraInitialConnectionStrategyPlugin] Timed out after {}ms while waiting for the cluster topology of '{}' to be fetched. Falling back to connecting via the provided endpoint.
+AuroraInitialConnectionStrategyPlugin.FailedToConnectToSelectedInstance=[AuroraInitialConnectionStrategyPlugin] Failed to connect to the selected instance '{}'. Falling back to the connection opened via the provided endpoint.
 
 AdfsCredentialsProviderFactory.FailedLogin=[AdfsCredentialsProviderFactory] Failed login. Could not obtain SAML Assertion from ADFS SignOn Page POST response: '{}'
 AdfsCredentialsProviderFactory.GetSamlAssertionFailed=[AdfsCredentialsProviderFactory] Failed to get SAML Assertion due to exception: '{}'

aws_advanced_python_wrapper/utils/properties.py

@@ -585,12 +585,49 @@ class WrapperProperties:
         False,
     )
 
+    # Deprecated. Use INITIAL_CONNECTION_HOST_SELECTOR_STRATEGY instead.
     READER_INITIAL_HOST_SELECTOR_STRATEGY = WrapperProperty(
         "reader_initial_connection_host_selector_strategy",
         "The strategy that should be used to select a new reader host while opening a new connection.",
         "random",
     )
 
+    INITIAL_CONNECTION_HOST_SELECTOR_STRATEGY = WrapperProperty(
+        "initial_connection_host_selector_strategy",
+        "The strategy that should be used to select a host while opening a new connection.",
+        "random",
+    )
+
+    ENDPOINT_SUBSTITUTION_ROLE = WrapperProperty(
+        "endpoint_substitution_role",
+        "Defines whether or not the initial connection URL should be replaced with an instance URL from the "
+        "topology info when available, and if so, the role of the instance URL that should be selected. "
+        "Valid values are 'writer', 'reader', 'any', or 'none'.",
+        None,
+    )
+
+    INACTIVE_CLUSTER_WRITER_SUBSTITUTION_ROLE = WrapperProperty(
+        "inactive_cluster_writer_endpoint_substitution_role",
+        "Defines whether or not the inactive cluster writer endpoint in the initial connection URL should "
+        "be replaced with a writer instance URL from the topology info when available. "
+        "Valid values are 'writer' or 'none'.",
+        "writer",
+    )
+
+    VERIFY_OPENED_CONNECTION_ROLE = WrapperProperty(
+        "verify_opened_connection_type",
+        "Defines whether an opened connection should be verified to be a writer or reader, "
+        "or if no role verification should be performed. Valid values are 'writer', 'reader', or 'none'.",
+        None,
+    )
+
+    VERIFY_INACTIVE_CLUSTER_WRITER_CONNECTION_ROLE = WrapperProperty(
+        "verify_inactive_cluster_writer_endpoint_connection_type",
+        "Defines whether inactive cluster writer connection should be verified to be a writer, "
+        "or if no role verification should be performed. Valid values are 'writer' or 'none'.",
+        "writer",
+    )
+
     OPEN_CONNECTION_RETRY_TIMEOUT_MS = WrapperProperty(
         "open_connection_retry_timeout_ms",
         "Maximum allowed time for the retries opening a connection.",
@@ -603,6 +640,19 @@ class WrapperProperties:
         1000,
     )
 
+    WAIT_FOR_INITIAL_TOPOLOGY_MS = WrapperProperty(
+        "wait_for_initial_topology_ms",
+        "Maximum allowed time, in milliseconds, to wait for the cluster topology to be fetched before opening a new "
+        "connection. When set to a value greater than 0 and the topology is not yet available, the plugin will block "
+        "until the topology has been discovered (or this timeout is reached) instead of falling back to connecting via "
+        "the initial endpoint in the connection string. This ensures host selection strategies such as 'round_robin' "
+        "distribute concurrent and connection-pool prefill connections across instances rather than routing them all "
+        "to a single instance resolved through DNS. The wait is scoped to the cluster the connection belongs to; "
+        "connections to other clusters are not affected. When set to 0 (the default) the previous behavior is "
+        "preserved.",
+        0,
+    )
+
     # Simple Read/Write Splitting
     SRW_READ_ENDPOINT = WrapperProperty(
         "srw_read_endpoint",

aws_advanced_python_wrapper/utils/rds_utils.py

@@ -202,6 +202,14 @@ def get_rds_region(self, host: Optional[str]):
             return elb_matcher.group(RdsUtils.REGION_GROUP)
         return None
 
+    def is_same_region(self, host1: Optional[str], host2: Optional[str]) -> bool:
+        if not host1 or not host1.strip() or not host2 or not host2.strip():
+            return False
+
+        host1_region = self.get_rds_region(host1)
+        host2_region = self.get_rds_region(host2)
+        return host1_region is not None and host2_region is not None and host1_region.casefold() == host2_region.casefold()
+
     def is_writer_cluster_dns(self, host: str) -> bool:
         dns_group = self._get_dns_group(self._get_prepared_host(host))
         return dns_group is not None and dns_group.casefold() == "cluster-"