fix: incorrect writer host details in ClusterTopologyMonitorImpl (#1095)

Author: aaron-congo

aws_advanced_python_wrapper/cluster_topology_monitor.py

@@ -20,10 +20,12 @@
 from concurrent.futures import ThreadPoolExecutor
 from typing import TYPE_CHECKING, Dict, Optional, Tuple
 
+from aws_advanced_python_wrapper.host_availability import HostAvailability
 from aws_advanced_python_wrapper.hostinfo import HostInfo
 from aws_advanced_python_wrapper.utils.atomic import AtomicReference
 from aws_advanced_python_wrapper.utils.cache_map import CacheMap
 from aws_advanced_python_wrapper.utils.messages import Messages
+from aws_advanced_python_wrapper.utils.rdsutils import RdsUtils
 from aws_advanced_python_wrapper.utils.thread_safe_connection_holder import \
     ThreadSafeConnectionHolder
 from aws_advanced_python_wrapper.utils.utils import LogUtils
@@ -87,6 +89,7 @@ def __init__(self, plugin_service: PluginService, topology_utils: TopologyUtils,
         self._refresh_rate_nano = refresh_rate_nano
         self._high_refresh_rate_nano = high_refresh_rate_nano
 
+        self._rds_utils = RdsUtils()
         self._writer_host_info: AtomicReference[Optional[HostInfo]] = AtomicReference(None)
         self._monitoring_connection: ThreadSafeConnectionHolder = ThreadSafeConnectionHolder(None)
 
@@ -300,16 +303,34 @@ def _open_any_connection_and_update_topology(self) -> Tuple[HostInfo, ...]:
             try:
                 conn = self._plugin_service.force_connect(self._initial_host_info, self._monitoring_properties)
                 self._monitoring_connection.set(conn, close_previous=False)
-                logger.debug("ClusterTopologyMonitorImpl.OpenedMonitoringConnection", self._cluster_id, self._initial_host_info.host)
+                logger.debug("ClusterTopologyMonitorImpl.OpenedMonitoringConnection",
+                             self._cluster_id, self._initial_host_info.host)
 
                 try:
-                    writer_host = self._topology_utils.get_writer_host_if_connected(
-                            conn, self._plugin_service.driver_dialect)
-                    if writer_host:
+                    writer_id = self._topology_utils.get_writer_id_if_connected(
+                        conn, self._plugin_service.driver_dialect)
+                    if writer_id:
                         self._is_verified_writer_connection = True
                         writer_verified_by_this_thread = True
-                        self._writer_host_info.set(HostInfo(writer_host, self._initial_host_info.port))
-                        logger.debug("ClusterTopologyMonitorImpl.WriterMonitoringConnection", self._cluster_id, writer_host)
+
+                        if self._rds_utils.is_rds_instance(self._initial_host_info.host):
+                            writer_host_info = self._initial_host_info
+                            self._writer_host_info.set(writer_host_info)
+                        else:
+                            writer_host = self._instance_template.host.replace("?", writer_id)
+                            port = self._instance_template.port \
+                                if self._instance_template.is_port_specified() \
+                                else self._initial_host_info.port
+                            writer_host_info = HostInfo(
+                                writer_host,
+                                port,
+                                HostRole.WRITER,
+                                HostAvailability.AVAILABLE,
+                                host_id=writer_id)
+                            self._writer_host_info.set(writer_host_info)
+
+                        logger.debug("ClusterTopologyMonitorImpl.WriterMonitoringConnection",
+                                     self._cluster_id, writer_host_info.host)
                 except Exception:
                     pass
             except Exception:
@@ -321,7 +342,7 @@ def _open_any_connection_and_update_topology(self) -> Tuple[HostInfo, ...]:
                 self._ignore_new_topology_requests_end_time_nano = 0
             else:
                 self._ignore_new_topology_requests_end_time_nano = (
-                    time.time_ns() + self.IGNORE_TOPOLOGY_REQUEST_NANO)
+                        time.time_ns() + self.IGNORE_TOPOLOGY_REQUEST_NANO)
 
         if len(hosts) == 0:
             self._monitoring_connection.clear()
@@ -472,7 +493,7 @@ def __call__(self) -> None:
                 if connection is not None:
                     is_writer = False
                     try:
-                        is_writer = self._monitor._topology_utils.get_writer_host_if_connected(
+                        is_writer = self._monitor._topology_utils.get_writer_id_if_connected(
                             connection, self._monitor._plugin_service.driver_dialect) is not None
                     except Exception:
                         self._monitor._close_connection(connection)

aws_advanced_python_wrapper/host_list_provider.py

@@ -606,7 +606,7 @@ def _get_host_id(self, conn: Connection):
             cursor.execute(self._dialect.host_id_query)
             return cursor.fetchone()
 
-    def get_writer_host_if_connected(self, connection: Connection, driver_dialect: DriverDialect) -> Optional[str]:
+    def get_writer_id_if_connected(self, connection: Connection, driver_dialect: DriverDialect) -> Optional[str]:
         try:
             cursor_execute_func_with_timeout = preserve_transaction_status_with_timeout(
                 self._thread_pool, self._max_timeout_sec, driver_dialect, connection)(self._get_writer_id)

tests/unit/test_cluster_topology_monitor.py

@@ -44,7 +44,7 @@ def topology_utils_mock():
         HostInfo("reader1.com", 5432, HostRole.READER),
         HostInfo("reader2.com", 5432, HostRole.READER)
     )
-    mock.get_writer_host_if_connected.return_value = "writer.com"
+    mock.get_writer_id_if_connected.return_value = "writer.com"
     mock.get_host_role.return_value = HostRole.WRITER
     return mock
 
@@ -204,7 +204,7 @@ def test_open_any_connection_and_update_topology_verifies_writer(self, cluster_m
         connection_mock = MagicMock()
         cluster_monitor._writer_host_info.set(None)
         cluster_monitor._plugin_service.force_connect.return_value = connection_mock
-        topology_utils_mock.get_writer_host_if_connected.return_value = "writer"
+        topology_utils_mock.get_writer_id_if_connected.return_value = "writer"
         cluster_monitor._fetch_topology_and_update_cache = MagicMock(return_value=expected_hosts)
 
         result = cluster_monitor._open_any_connection_and_update_topology()
@@ -248,7 +248,7 @@ def test_call_connection_success_writer_detected(self, monitor_impl_mock, topolo
         monitor = HostMonitor(monitor_impl_mock, host_info, None)
         connection_mock = MagicMock()
         monitor_impl_mock._plugin_service.force_connect.return_value = connection_mock
-        topology_utils_mock.get_writer_host_if_connected.return_value = "writer.com"
+        topology_utils_mock.get_writer_id_if_connected.return_value = "writer.com"
         topology_utils_mock.get_host_role.return_value = HostRole.WRITER
 
         call_count = [0]
@@ -272,7 +272,7 @@ def test_call_connection_success_reader_detected(self, monitor_impl_mock, topolo
         monitor = HostMonitor(monitor_impl_mock, host_info, None)
         connection_mock = MagicMock()
         monitor_impl_mock._plugin_service.force_connect.return_value = connection_mock
-        topology_utils_mock.get_writer_host_if_connected.return_value = None
+        topology_utils_mock.get_writer_id_if_connected.return_value = None
 
         call_count = [0]