apply review comments

sophia-bq · sophia-bq · commit e863f98d3a05 · 2025-12-15T10:18:49.000-08:00
diff --git a/aws_advanced_python_wrapper/read_write_splitting_plugin.py b/aws_advanced_python_wrapper/read_write_splitting_plugin.py
@@ -56,11 +56,11 @@ def __init__(
         self,
         plugin_service: PluginService,
         props: Properties,
-        connection_handler: ConnectionHandler,
+        connection_handler: ReadWriteConnectionHandler,
     ):
         self._plugin_service: PluginService = plugin_service
         self._properties: Properties = props
-        self._connection_handler: ConnectionHandler = connection_handler
+        self._connection_handler: ReadWriteConnectionHandler = connection_handler
         self._writer_connection: Optional[Connection] = None
         self._reader_connection: Optional[Connection] = None
         self._writer_host_info: Optional[HostInfo] = None
@@ -310,7 +310,7 @@ def _switch_to_reader_connection(self):
 
         if (
             self._reader_connection is not None
-            and not self._connection_handler.old_reader_can_be_used(
+            and not self._connection_handler.can_host_be_used(
                 self._reader_host_info
             )
         ):
@@ -344,7 +344,7 @@ def _switch_to_reader_connection(self):
             self._close_connection_if_idle(self._writer_connection)
 
     def _initialize_reader_connection(self):
-        if self._connection_handler.need_connect_to_writer():
+        if self._connection_handler.has_no_readers():
             if not self._is_connection_usable(
                 self._writer_connection, self._plugin_service.driver_dialect
             ):
@@ -437,7 +437,7 @@ def close_connection(conn: Optional[Connection], driver_dialect: DriverDialect):
                 return
 
 
-class ConnectionHandler(Protocol):
+class ReadWriteConnectionHandler(Protocol):
     """Protocol for handling writer/reader connection logic."""
 
     @property
@@ -491,15 +491,15 @@ def is_writer_host(self, current_host: HostInfo) -> bool:
         ...
 
     def is_reader_host(self, current_host: HostInfo) -> bool:
-        """Return true if the current host fits the criteria of a writer host."""
+        """Return true if the current host fits the criteria of a reader host."""
         ...
 
-    def old_reader_can_be_used(self, reader_host_info: HostInfo) -> bool:
-        """Return true if the current host can be used to switch connection to."""
+    def can_host_be_used(self, host_info: HostInfo) -> bool:
+        """Returns true if connections can be switched to the given host"""
         ...
 
-    def need_connect_to_writer(self) -> bool:
-        """Return true if switching to reader should instead connect to writer."""
+    def has_no_readers(self) -> bool:
+        """Return true if there are no readers in the host list"""
         ...
 
     def refresh_and_store_host_list(
@@ -509,7 +509,7 @@ def refresh_and_store_host_list(
         ...
 
 
-class TopologyBasedConnectionHandler(ConnectionHandler):
+class TopologyBasedConnectionHandler(ReadWriteConnectionHandler):
     """Topology based implementation of connection handling logic."""
 
     def __init__(self, plugin_service: PluginService, props: Properties):
@@ -615,11 +615,11 @@ def get_verified_initial_connection(
 
         return current_conn
 
-    def old_reader_can_be_used(self, reader_host_info: HostInfo) -> bool:
+    def can_host_be_used(self, host_info: HostInfo) -> bool:
         hostnames = [host_info.host for host_info in self._hosts]
-        return reader_host_info is not None and reader_host_info.host in hostnames
+        return host_info.host in hostnames
 
-    def need_connect_to_writer(self) -> bool:
+    def has_no_readers(self) -> bool:
         if len(self._hosts) == 1:
             return self._get_writer() is not None
         return False
diff --git a/aws_advanced_python_wrapper/simple_read_write_splitting_plugin.py b/aws_advanced_python_wrapper/simple_read_write_splitting_plugin.py
@@ -19,7 +19,7 @@
 
 from aws_advanced_python_wrapper.host_availability import HostAvailability
 from aws_advanced_python_wrapper.read_write_splitting_plugin import (
-    ConnectionHandler, ReadWriteSplittingConnectionManager)
+    ReadWriteConnectionHandler, ReadWriteSplittingConnectionManager)
 from aws_advanced_python_wrapper.utils.rds_url_type import RdsUrlType
 from aws_advanced_python_wrapper.utils.rdsutils import RdsUtils
 
@@ -37,14 +37,14 @@
 from aws_advanced_python_wrapper.utils.properties import WrapperProperties
 
 
-class EndpointBasedConnectionHandler(ConnectionHandler):
+class EndpointBasedConnectionHandler(ReadWriteConnectionHandler):
     """Endpoint based implementation of connection handling logic."""
 
     def __init__(self, plugin_service: PluginService, props: Properties):
-        self._read_endpoint: str = EndpointBasedConnectionHandler._verify_parameter(
+        read_endpoint: str = EndpointBasedConnectionHandler._verify_parameter(
             WrapperProperties.SRW_READ_ENDPOINT, props, str, required=True
         )
-        self._write_endpoint: str = EndpointBasedConnectionHandler._verify_parameter(
+        write_endpoint: str = EndpointBasedConnectionHandler._verify_parameter(
             WrapperProperties.SRW_WRITE_ENDPOINT, props, str, required=True
         )
 
@@ -60,17 +60,19 @@ def __init__(self, plugin_service: PluginService, props: Properties):
                 WrapperProperties.SRW_CONNECT_RETRY_INTERVAL_MS, props, int, lambda x: x > 0
             )
 
-            self._verify_opened_connection_type: Optional[HostRole] = (
-                EndpointBasedConnectionHandler._parse_connection_type(
+            self._verify_initial_connection_type: Optional[HostRole] = (
+                EndpointBasedConnectionHandler._parse_role(
                     WrapperProperties.SRW_VERIFY_INITIAL_CONNECTION_TYPE.get(props)
                 )
             )
 
         self._plugin_service: PluginService = plugin_service
         self._rds_utils: RdsUtils = RdsUtils()
         self._host_list_provider_service: Optional[HostListProviderService] = None
-        self._write_endpoint_host_info: HostInfo = self._create_host_info(self._write_endpoint, HostRole.WRITER)
-        self._read_endpoint_host_info: HostInfo = self._create_host_info(self._read_endpoint, HostRole.READER)
+        self._write_endpoint_host_info: HostInfo = self._create_host_info(write_endpoint, HostRole.WRITER)
+        self._read_endpoint_host_info: HostInfo = self._create_host_info(read_endpoint, HostRole.READER)
+        self._write_endpoint = write_endpoint.casefold()
+        self._read_endpoint = read_endpoint.casefold()
 
     @property
     def host_list_provider_service(self) -> Optional[HostListProviderService]:
@@ -116,12 +118,12 @@ def get_verified_initial_connection(
 
         if (
             url_type == RdsUrlType.RDS_WRITER_CLUSTER
-            or self._verify_opened_connection_type == HostRole.WRITER
+            or self._verify_initial_connection_type == HostRole.WRITER
         ):
             conn = self._get_verified_connection(host_info, HostRole.WRITER, plugin_service_connect_func, connect_func)
         elif (
             url_type == RdsUrlType.RDS_READER_CLUSTER
-            or self._verify_opened_connection_type == HostRole.READER
+            or self._verify_initial_connection_type == HostRole.READER
         ):
             conn = self._get_verified_connection(host_info, HostRole.READER, plugin_service_connect_func, connect_func)
 
@@ -154,10 +156,8 @@ def _get_verified_connection(
             try:
                 if connect_func is not None:
                     candidate_conn = connect_func()
-                elif host_info is not None:
-                    candidate_conn = plugin_service_connect_func(host_info)
                 else:
-                    return None
+                    candidate_conn = plugin_service_connect_func(host_info)
 
                 if candidate_conn is None:
                     self._delay()
@@ -178,12 +178,12 @@ def _get_verified_connection(
 
         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.
+    def can_host_be_used(self, host_info: HostInfo) -> bool:
+        # Assume that the host 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.
+    def has_no_readers(self) -> bool:
+        # SetReadOnly(true) will always connect to the read_endpoint, regardless of number of readers.
         return False
 
     def refresh_and_store_host_list(
@@ -218,14 +218,14 @@ def should_update_reader_with_current_conn(
 
     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()
+            current_host.host.casefold() == self._write_endpoint
+            or current_host.url.casefold() == self._write_endpoint
         )
 
     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()
+            current_host.host.casefold() == self._read_endpoint
+            or current_host.url.casefold() == self._read_endpoint
         )
 
     def _create_host_info(self, endpoint: str, role: HostRole) -> HostInfo:
@@ -272,14 +272,14 @@ def _delay(self):
         sleep(self._connect_retry_interval_ms / 1000)
 
     @staticmethod
-    def _parse_connection_type(phase_str: Optional[str]) -> HostRole:
-        if not phase_str:
+    def _parse_role(role_str: Optional[str]) -> HostRole:
+        if not role_str:
             return HostRole.UNKNOWN
 
-        phase_upper = phase_str.lower()
-        if phase_upper == "reader":
+        phase_lower = role_str.lower()
+        if phase_lower == "reader":
             return HostRole.READER
-        elif phase_upper == "writer":
+        elif phase_lower == "writer":
             return HostRole.WRITER
         else:
             raise ValueError(
diff --git a/tests/integration/container/conftest.py b/tests/integration/container/conftest.py
@@ -140,6 +140,7 @@ def pytest_runtest_setup(item):
         CustomEndpointPlugin._monitors.clear()
         CustomEndpointMonitor._custom_endpoint_info_cache.clear()
         MonitoringThreadContainer.clean_up()
+        ConnectionProviderManager.release_resources()
 
         ConnectionProviderManager.reset_provider()
         DatabaseDialectManager.reset_custom_dialect()
