feat: srw

Author: sophia-bq

aws_advanced_python_wrapper/plugin_service.py

@@ -80,8 +80,8 @@
 from aws_advanced_python_wrapper.hostinfo import HostInfo, HostRole
 from aws_advanced_python_wrapper.iam_plugin import IamAuthPluginFactory
 from aws_advanced_python_wrapper.plugin import CanReleaseResources
-from aws_advanced_python_wrapper.read_write_splitting_plugin import \
-    ReadWriteSplittingPluginFactory
+from aws_advanced_python_wrapper.read_write_splitting_plugin import ReadWriteSplittingPluginFactory
+from aws_advanced_python_wrapper.simple_read_write_splitting_plugin import SimpleReadWriteSplittingPluginFactory
 from aws_advanced_python_wrapper.stale_dns_plugin import StaleDnsPluginFactory
 from aws_advanced_python_wrapper.utils.cache_map import CacheMap
 from aws_advanced_python_wrapper.utils.decorators import \
@@ -760,6 +760,7 @@ class PluginManager(CanReleaseResources):
         "host_monitoring_v2": HostMonitoringV2PluginFactory,
         "failover": FailoverPluginFactory,
         "read_write_splitting": ReadWriteSplittingPluginFactory,
+        "srw": SimpleReadWriteSplittingPluginFactory,
         "fastest_response_strategy": FastestResponseStrategyPluginFactory,
         "stale_dns": StaleDnsPluginFactory,
         "custom_endpoint": CustomEndpointPluginFactory,
@@ -784,6 +785,7 @@ class PluginManager(CanReleaseResources):
         AuroraConnectionTrackerPluginFactory: 100,
         StaleDnsPluginFactory: 200,
         ReadWriteSplittingPluginFactory: 300,
+        SimpleReadWriteSplittingPluginFactory: 310,
         FailoverPluginFactory: 400,
         HostMonitoringPluginFactory: 500,
         HostMonitoringV2PluginFactory: 510,

aws_advanced_python_wrapper/read_write_splitting_plugin.py

@@ -286,7 +286,7 @@ OpenTelemetryFactory.WrongParameterType="[OpenTelemetryFactory] Wrong parameter
 
 Plugin.UnsupportedMethod=[Plugin] '{}' is not supported by this plugin.
 
-PluginManager.ConfigurationProfileNotFound=PluginManager] Configuration profile '{}' not found.
+PluginManager.ConfigurationProfileNotFound=[PluginManager] Configuration profile '{}' not found.
 PluginManager.InvalidPlugin=[PluginManager] Invalid plugin requested: '{}'.
 PluginManager.MethodInvokedAgainstOldConnection = [PluginManager] The internal connection has changed since '{}' was created. This is likely due to failover or read-write splitting functionality. To ensure you are using the updated connection, please re-create Cursor objects after failover and/or setting readonly.
 PluginManager.PipelineNone=[PluginManager] A pipeline was requested but the created pipeline evaluated to None.
@@ -357,8 +357,9 @@ ReadWriteSplittingPlugin.ErrorVerifyingInitialHostSpecRole=[ReadWriteSplittingPl
 ReadWriteSplittingPlugin.ExceptionWhileExecutingCommand=[ReadWriteSplittingPlugin] Detected an exception while executing a command: '{}'
 ReadWriteSplittingPlugin.ExecutingAgainstOldConnection=[ReadWriteSplittingPlugin] Executing method against old connection: '{}'
 ReadWriteSplittingPlugin.FailedToConnectToReader=[ReadWriteSplittingPlugin] Failed to connect to reader host: '{}'
+ReadWriteSplittingPlugin.FailedToConnectToWriter=[ReadWriteSplittingPlugin] Failed to connect to writer host: '{}'
 ReadWriteSplittingPlugin.FailoverExceptionWhileExecutingCommand=[ReadWriteSplittingPlugin] Detected a failover exception while executing a command: '{}'
-ReadWriteSplittingPlugin.FallbackToWriter=[ReadWriteSplittingPlugin] Failed to switch to a reader; the current writer will be used as a fallback: '{}'
+ReadWriteSplittingPlugin.FallbackToCurrentConnection=[ReadWriteSplittingPlugin] Failed to switch to a reader; the current connection will be used as a fallback: '{}'
 ReadWriteSplittingPlugin.NoReadersAvailable=[ReadWriteSplittingPlugin] The plugin was unable to establish a reader connection to any reader instance.
 ReadWriteSplittingPlugin.NoReadersFound=[ReadWriteSplittingPlugin] A reader instance was requested via set_read_only, but there are no readers in the host list. The current writer will be used as a fallback: '{}'
 ReadWriteSplittingPlugin.NoWriterFound=[ReadWriteSplittingPlugin] No writer was found in the current host list. This may occur if the writer is not in the list of allowed hosts.
@@ -382,6 +383,9 @@ RoundRobinHostSelector.RoundRobinInvalidHostWeightPairs= [RoundRobinHostSelector
 WeightedRandomHostSelector.WeightedRandomInvalidHostWeightPairs= [WeightedRandomHostSelector] The provided host weight pairs have not been configured correctly. Please ensure the provided host weight pairs is a comma separated list of pairs, each pair in the format of <host>:<weight>. Weight values must be an integer greater than or equal to the default weight value of 1. Weight pair: '{}'
 WeightedRandomHostSelector.WeightedRandomInvalidDefaultWeight=[WeightedRandomHostSelector] The provided default weight value is not valid. Weight values must be an integer greater than or equal to 1.
 
+SimpleReadWriteSplittingPlugin.MissingRequiredConfigParameter=[SimpleReadWriteSplittingPlugin] Configuration parameter {} is required.
+SimpleReadWriteSplittingPlugin.IncorrectConfiguration=[SimpleReadWriteSplittingPlugin] Unable to verify connections with this current configuration. Ensure a correct value is provided to the configuration parameter {}.
+
 SqlAlchemyPooledConnectionProvider.PoolNone=[SqlAlchemyPooledConnectionProvider] Attempted to find or create a pool for '{}' but the result of the attempt evaluated to None.
 SqlAlchemyPooledConnectionProvider.UnableToCreateDefaultKey=[SqlAlchemyPooledConnectionProvider] Unable to create a default key for internal connection pools. By default, the user parameter is used, but the given user evaluated to None or the empty string (""). Please ensure you have passed a valid user in the connection properties.
 

aws_advanced_python_wrapper/resources/aws_advanced_python_wrapper_messages.properties

@@ -0,0 +1,243 @@
+#  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License").
+#  You may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#  http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+from __future__ import annotations
+
+from enum import Enum, auto
+from time import perf_counter_ns, sleep
+from typing import TYPE_CHECKING, Callable, Optional
+
+from aws_advanced_python_wrapper.host_availability import HostAvailability
+from aws_advanced_python_wrapper.read_write_splitting_plugin import ReadWriteSplittingConnectionManager, ConnectionHandler
+from aws_advanced_python_wrapper.utils.rds_url_type import RdsUrlType
+from aws_advanced_python_wrapper.utils.rdsutils import RdsUtils
+
+if TYPE_CHECKING:
+    from aws_advanced_python_wrapper.driver_dialect import DriverDialect
+    from aws_advanced_python_wrapper.host_list_provider import HostListProviderService
+    from aws_advanced_python_wrapper.pep249 import Connection
+    from aws_advanced_python_wrapper.plugin_service import PluginService
+    from aws_advanced_python_wrapper.utils.properties import Properties
+
+from aws_advanced_python_wrapper.errors import AwsWrapperError
+from aws_advanced_python_wrapper.hostinfo import HostInfo, HostRole
+from aws_advanced_python_wrapper.plugin import PluginFactory
+from aws_advanced_python_wrapper.utils.messages import Messages
+from aws_advanced_python_wrapper.utils.properties import WrapperProperties
+
+class EndpointBasedConnectionHandler(ConnectionHandler):
+    """Endpoint based implementation of connection handling logic."""
+    def __init__(self, plugin_service: PluginService, props: Properties):
+        srw_read_endpoint = WrapperProperties.SRW_READ_ENDPOINT.get(props)
+        if srw_read_endpoint is None:
+            raise AwsWrapperError(Messages.get_formatted("SimpleReadWriteSplittingPlugin.MissingRequiredConfigParameter", WrapperProperties.SRW_READ_ENDPOINT.name))
+        self.read_endpoint = srw_read_endpoint
+
+        srw_write_endpoint = WrapperProperties.SRW_WRITE_ENDPOINT.get(props)
+        if srw_write_endpoint is None:
+            raise AwsWrapperError(Messages.get_formatted("SimpleReadWriteSplittingPlugin.MissingRequiredConfigParameter", WrapperProperties.SRW_WRITE_ENDPOINT.name))
+        self.write_endpoint = srw_write_endpoint
+    
+        self.verify_new_connections = WrapperProperties.SRW_VERIFY_NEW_CONNECTIONS.get_bool(props)
+        if self.verify_new_connections is True:
+            srw_connect_retry_timeout_ms = WrapperProperties.SRW_CONNECT_RETRY_TIMEOUT_MS.get_int(props)
+            if srw_connect_retry_timeout_ms <= 0:
+                raise ValueError(Messages.get_formatted("SimpleReadWriteSplittingPlugin.IncorrectConfiguration", WrapperProperties.SRW_CONNECT_RETRY_TIMEOUT_MS.name))
+            self.connect_retry_timeout_ms = srw_connect_retry_timeout_ms
+
+            srw_connect_retry_interval_ms = WrapperProperties.SRW_CONNECT_RETRY_INTERVAL_MS.get_int(props)
+            if srw_connect_retry_interval_ms <= 0:
+                raise ValueError(Messages.get_formatted("SimpleReadWriteSplittingPlugin.IncorrectConfiguration", WrapperProperties.SRW_CONNECT_RETRY_INTERVAL_MS.name))
+            self.connect_retry_interval_ms = srw_connect_retry_interval_ms
+
+            self.verify_opened_connection_type = EndpointBasedConnectionHandler._parse_connection_type(WrapperProperties.SRW_VERIFY_INITIAL_CONNECTION_TYPE.get(props))
+    
+        self._plugin_service = plugin_service
+        self._properties = props
+        self._rds_utils = RdsUtils()
+        self._host_list_provider_service: Optional[HostListProviderService] = None
+        self._write_endpoint_host_info = None 
+        self._read_endpoint_host_info = None 
+
+    def open_new_writer_connection(self) -> Optional[tuple[Connection, HostInfo]]: 
+        if self._write_endpoint_host_info is None:
+            self._write_endpoint_host_info = self._create_host_info(self.write_endpoint, HostRole.WRITER)
+
+        conn: Optional[Connection] = None
+        if self.verify_new_connections:
+            conn = self._get_verified_connection(self._properties, self._write_endpoint_host_info, HostRole.WRITER)
+        else:
+            conn = self._plugin_service.connect(self._write_endpoint_host_info, self._properties, None)
+
+        return conn, self._write_endpoint_host_info
+    
+    def open_new_reader_connection(self) -> Optional[tuple[Connection, HostInfo]]:
+        if self._read_endpoint_host_info is None:
+            self._read_endpoint_host_info = self._create_host_info(self.read_endpoint, HostRole.READER)
+
+        conn: Optional[Connection] = None
+        if self.verify_new_connections:
+            conn = self._get_verified_connection(self._properties, self._read_endpoint_host_info, HostRole.READER)
+        else:
+            conn = self._plugin_service.connect(self._read_endpoint_host_info, self._properties, None)
+
+        return conn, self._read_endpoint_host_info
+
+    def get_verified_initial_connection(self, host_info: HostInfo, props: Properties, is_initial_connection: bool, connect_func: Callable) -> Optional[Connection]:
+        if not is_initial_connection or not self.verify_new_connections:
+            return connect_func()
+        
+        url_type: RdsUrlType = self._rds_utils.identify_rds_type(host_info.host)
+
+        conn: Optional[Connection] = None
+
+        if url_type == RdsUrlType.RDS_WRITER_CLUSTER or (self.verify_opened_connection_type is not None and self.verify_opened_connection_type == HostRole.WRITER):
+            conn = self._get_verified_connection(props, host_info, HostRole.WRITER, connect_func)
+        elif url_type == RdsUrlType.RDS_READER_CLUSTER or (self.verify_opened_connection_type is not None and self.verify_opened_connection_type == HostRole.READER):
+            conn = self._get_verified_connection(props, host_info, HostRole.READER, connect_func)
+
+        if conn is None:
+            conn = connect_func()
+            
+        self._set_initial_connection_host_info(conn, host_info)
+        return conn
+
+    def _set_initial_connection_host_info(self, conn: Connection, host_info: HostInfo):
+        if self.set_host_list_provider_service is None:
+            return
+        
+        if host_info is None:
+            try:
+                host_info = self._plugin_service.identify_connection(conn)
+            except Exception as e:
+                return
+            
+        if host_info is not None and self._host_list_provider_service is not None:
+            self._host_list_provider_service.initial_connection_host_info = host_info
+
+    def _get_verified_connection(self, props: Properties, host_info: HostInfo, role: HostRole, connect_func: Callable = None) -> Connection:
+        end_time_nano = perf_counter_ns() + (self.connect_retry_timeout_ms * 1000000)
+
+        candidate_conn: Optional[Connection]
+        attempt = 0
+
+        while perf_counter_ns() < end_time_nano:
+            attempt += 1
+            candidate_conn = None
+
+            try:
+                if connect_func is not None:
+                    candidate_conn = connect_func()
+                elif host_info is not None:
+                    candidate_conn = self._plugin_service.connect(host_info, props, None)
+                else:
+                    return None
+
+                if candidate_conn is None:
+                    self._delay()
+                    continue
+                
+                actual_role = self._plugin_service.get_host_role(candidate_conn)
+
+                instance_connected_to = self._plugin_service.identify_connection(candidate_conn)
+
+                if actual_role != role:
+                    ReadWriteSplittingConnectionManager.close_connection(candidate_conn)
+                    self._delay()
+                    continue
+
+                return candidate_conn
+                
+            except Exception as e:
+                ReadWriteSplittingConnectionManager.close_connection(candidate_conn)
+                self._delay()
+
+        return None
+
+    def old_reader_can_be_used(self, reader_host_info: HostInfo) -> bool:
+        # Assume that the old reader can always be used, no topology-based information to check.
+        return True
+
+    def need_connect_to_writer(self) -> bool:
+        # SetReadOnly(true) will always connect to the read_endpoint, and not the writer.
+        return False
+
+    def refresh_and_store_host_list(self, current_conn: Connection, driver_dialect: DriverDialect):
+        # Endpoint based connections do not require a host list.
+        return
+
+    def should_update_writer_with_current_conn(self, current_conn: Connection, current_host: HostInfo, writer_conn: Connection) -> bool:
+        return self.is_writer_host(current_host) and current_conn != writer_conn and (not self.verify_new_connections or self._plugin_service.get_host_role(current_conn) == HostRole.WRITER)
+
+    def should_update_reader_with_current_conn(self, current_conn: Connection, current_host: HostInfo, reader_conn: Connection) -> bool:
+        return self.is_reader_host(current_host) and current_conn != reader_conn and (not self.verify_new_connections or self._plugin_service.get_host_role(current_conn) == HostRole.READER)
+    
+    def is_writer_host(self, current_host: HostInfo) -> bool:
+        return current_host.host.casefold() == self.write_endpoint.casefold() or current_host.url.casefold() == self.write_endpoint.casefold()
+
+    def is_reader_host(self, current_host: HostInfo) -> bool:
+        return current_host.host.casefold() == self.read_endpoint.casefold() or current_host.url.casefold() == self.read_endpoint.casefold()
+    
+    def _create_host_info(self, endpoint, role: HostRole) -> HostInfo:
+        endpoint = endpoint.strip()
+        host = endpoint
+        port = self._plugin_service.database_dialect.default_port
+        colon_index = endpoint.rfind(":")
+
+        if colon_index != -1:
+            port_str = endpoint[colon_index + 1:]
+            if port_str.isdigit():
+                host = endpoint[:colon_index]
+                port = int(port_str)
+            else: 
+                if (self.set_host_list_provider_service is not None and self.set_host_list_provider_service.initial_connection_host_info is not None and 
+                    self.set_host_list_provider_service.initial_connection_host_info.port != HostInfo.NO_PORT):
+                    port = self.set_host_list_provider_service.initial_connection_host_info.port
+
+        return HostInfo(
+            host=host,
+            port=port,
+            role=role,
+            availability=HostAvailability.AVAILABLE)
+
+    def _delay(self):
+        sleep(self.connect_retry_interval_ms / 1000)
+
+    @staticmethod
+    def _parse_connection_type(phase_str: Optional[str]) -> HostRole:
+        if not phase_str:
+            return None 
+        
+        phase_upper = phase_str.lower()
+        if phase_upper == "reader":
+            return HostRole.READER
+        elif phase_upper == "writer":
+            return HostRole.WRITER
+        else:
+            raise ValueError(Messages.get_formatted("SimpleReadWriteSplittingPlugin.IncorrectConfiguration", WrapperProperties.SRW_VERIFY_INITIAL_CONNECTION_TYPE.name))
+    
+    
+class SimpleReadWriteSplittingPlugin(ReadWriteSplittingConnectionManager):    
+    def __init__(self, plugin_service, props: Properties):
+        # The simple read/write splitting plugin handles connections based on configuration parameter endpoints.
+        connection_handler = EndpointBasedConnectionHandler(
+            plugin_service, 
+            props,
+        )
+        
+        super().__init__(plugin_service, props, connection_handler)
+
+class SimpleReadWriteSplittingPluginFactory(PluginFactory):
+    def get_instance(self, plugin_service, props: Properties):
+        return SimpleReadWriteSplittingPlugin(plugin_service, props)

aws_advanced_python_wrapper/simple_read_write_splitting_plugin.py

@@ -462,6 +462,37 @@ class WrapperProperties:
         1000
     )
 
+    # Simple Read/Write Splitting
+    SRW_READ_ENDPOINT = WrapperProperty(
+        "srw_read_endpoint",
+        "The read-only endpoint that should be used to connect to a reader.",
+        None)
+    
+    SRW_WRITE_ENDPOINT = WrapperProperty(
+        "srw_write_endpoint",
+        "The read-write/cluster endpoint that should be used to connect to the writer.",
+        None)
+    
+    SRW_VERIFY_NEW_CONNECTIONS = WrapperProperty(
+        "srw_verify_new_connections",
+        "Enables role-verification for new connections made by the Simple Read/Write Splitting Plugin..",
+        True)
+    
+    SRW_VERIFY_INITIAL_CONNECTION_TYPE = WrapperProperty(
+        "srw_verify_initial_connection_type",
+        "Force to verify an initial connection to be either a writer or a reader.",
+        None)
+    
+    SRW_CONNECT_RETRY_TIMEOUT_MS = WrapperProperty(
+        "srw_connect_retry_timeout_ms",
+        "Maximum allowed time in milliseconds for the plugin to retry opening a connection.",
+        60000)
+    
+    SRW_CONNECT_RETRY_INTERVAL_MS = WrapperProperty(
+        "srw_connect_retry_interval_ms",
+        "Time in milliseconds between each retry of opening a connection.",
+        1000)
+
 
 class PropertiesUtils:
     _MONITORING_PROPERTY_PREFIX = "monitoring-"

aws_advanced_python_wrapper/utils/properties.py

@@ -75,6 +75,7 @@ build-backend = "poetry.core.masonry.api"
 [tool.pytest.ini_options]
 filterwarnings = [
     'ignore:cache could not write path',
-    'ignore:could not create cache path'
+    'ignore:could not create cache path',
+    'ignore:Exception during reset or similar:pytest.PytestUnhandledThreadExceptionWarning'
 ]