fix(aio): pin the instance for cluster-endpoint connections in aurora_connection_tracker

aws_advanced_python_wrapper/aio/aurora_connection_tracker.py

@@ -270,6 +270,8 @@ def __init__(
         self._plugin_service = plugin_service
         self._tracker = tracker or AsyncOpenedConnectionTracker()
         self._current_writer: Optional[HostInfo] = None
+        # The instance a cluster-endpoint connection landed on (#1276).
+        self._connected_instance: Optional[HostInfo] = None
         self._pending_invalidations: Set[asyncio.Task] = set()
 
         # Telemetry counter -- the sync tracker plugin doesn't emit this,
@@ -342,6 +344,13 @@ async def _pin_current_writer(self, host_info: HostInfo, conn: Any) -> None:
             role = await self._plugin_service.get_host_role(conn)
             if role == HostRole.WRITER:
                 host = self._plugin_service.current_host_info
+                # Topology names instances, so a cluster or custom endpoint
+                # never matches the writer and the first execute would
+                # invalidate this connection (#1276). Pin the instance the
+                # connection landed on; if it can't be identified, keep the
+                # topology pin.
+                if host is not None and not rds.is_rds_instance(host.host):
+                    host = self._connected_instance
                 if host is not None and (
                         self._current_writer is None
                         or not self._same_host(host, self._current_writer)):
@@ -376,6 +385,7 @@ async def _fill_instance_alias(
             identified = await self._plugin_service.identify_connection(conn)
             if identified is None:
                 return
+            self._connected_instance = identified
             host_info.reset_aliases()
             host_info.add_alias(host_info.as_alias())
             host_info.add_alias(*identified.as_aliases())
@@ -505,9 +515,15 @@ async def _invalidate_departed_host(
         """
         if pre_failover_host is None:
             return
+        # A cluster-endpoint connection departed the instance it landed on,
+        # which is what the pin and the tracker's instance keys name (#1276).
+        if (self._connected_instance is not None
+                and not RdsUtils().is_rds_instance(pre_failover_host.host)):
+            pre_failover_host = self._connected_instance
         post = self._plugin_service.current_host_info
         if post is None or self._same_host(pre_failover_host, post):
             return
+        self._connected_instance = post
         if (self._current_writer is not None
                 and not self._same_host(pre_failover_host, self._current_writer)
                 and not self._same_host(post, self._current_writer)):

tests/unit/test_aio_aurora_connection_tracker.py

@@ -1005,3 +1005,128 @@ async def _raising():
         assert plugin._current_writer.host == new_writer.host
 
     asyncio.run(_run())
+
+
+def test_cluster_endpoint_connection_survives_its_first_execute():
+    """Issue #1276: the connect-time live-probe pin used
+    ``current_host_info``, which for a cluster-endpoint connection is the
+    cluster endpoint. The first execute compared that pin against the
+    topology writer (an instance endpoint), saw a writer change, and
+    ``invalidate_all`` closed the connection about to run its statement."""
+    plugin, svc, driver_dialect, tracker = _build()
+    cluster = "mydb.cluster-xyz.us-east-1.rds.amazonaws.com"
+    instance = "inst-1.xyz.us-east-1.rds.amazonaws.com"
+    writer = HostInfo(host=instance, port=5432, role=HostRole.WRITER)
+
+    async def _refresh(*a, **k):
+        svc._all_hosts = (writer,)
+    svc.refresh_host_list = AsyncMock(side_effect=_refresh)
+    svc.identify_connection = AsyncMock(
+        return_value=HostInfo(host=instance, port=5432))
+    svc.get_host_role = AsyncMock(return_value=HostRole.WRITER)
+
+    conn = _plain_conn("conn")
+
+    async def _connect_func():
+        return conn
+
+    async def _run():
+        svc._current_host_info = HostInfo(host=cluster, port=5432)
+        await plugin.connect(
+            target_driver_func=MagicMock(), driver_dialect=driver_dialect,
+            host_info=HostInfo(host=cluster, port=5432), props=svc.props,
+            is_initial_connection=True, connect_func=_connect_func)
+        assert plugin._current_writer is not None
+        assert plugin._current_writer.host == instance
+
+        result = await plugin.execute(
+            MagicMock(), "Cursor.execute", AsyncMock(return_value="ok"))
+        assert result == "ok"
+        conn.close.assert_not_called()
+
+    asyncio.run(_run())
+
+
+def _connect_via(plugin, svc, driver_dialect, endpoint, conn):
+    async def _connect_func():
+        return conn
+
+    host_info = HostInfo(host=endpoint, port=5432)
+    svc._current_host_info = host_info
+    return plugin.connect(
+        target_driver_func=MagicMock(), driver_dialect=driver_dialect,
+        host_info=host_info, props=svc.props,
+        is_initial_connection=True, connect_func=_connect_func)
+
+
+def test_failover_invalidates_departed_instance_of_cluster_endpoint_connection():
+    """A cluster-endpoint connection departs the INSTANCE it landed on, not
+    the cluster endpoint. With topology still lagged (old writer named), the
+    departed-host path must resolve the cluster endpoint to that instance so
+    idle connections to the demoted writer are closed."""
+    plugin, svc, driver_dialect, tracker = _build()
+    cluster = "mydb.cluster-xyz.us-east-1.rds.amazonaws.com"
+    old = "inst-1.xyz.us-east-1.rds.amazonaws.com"
+    new = "inst-2.xyz.us-east-1.rds.amazonaws.com"
+    lagged = (HostInfo(old, 5432, role=HostRole.WRITER),
+              HostInfo(new, 5432, role=HostRole.READER))
+
+    async def _refresh(*a, **k):
+        svc._all_hosts = lagged
+    svc.refresh_host_list = AsyncMock(side_effect=_refresh)
+    svc.force_refresh_host_list = AsyncMock(side_effect=_refresh)
+    svc.identify_connection = AsyncMock(return_value=HostInfo(old, 5432))
+    svc.get_host_role = AsyncMock(return_value=HostRole.WRITER)
+
+    idle_conn = _plain_conn("idle")
+    active_conn = _plain_conn("active")
+
+    async def _run():
+        await _connect_via(plugin, svc, driver_dialect, cluster, idle_conn)
+        await _connect_via(plugin, svc, driver_dialect, cluster, active_conn)
+
+        async def _raising():
+            svc._current_host_info = HostInfo(new, 5432, role=HostRole.WRITER)
+            raise FailoverSuccessError("failover")
+        with pytest.raises(FailoverSuccessError):
+            await plugin.execute(MagicMock(), "Cursor.execute", _raising)
+        idle_conn.close.assert_called()
+
+    asyncio.run(_run())
+
+
+def test_reader_endpoint_failover_does_not_invalidate_the_writer():
+    """A reader-endpoint connection that fails over between readers departs
+    a reader, so writer connections must survive."""
+    plugin, svc, driver_dialect, tracker = _build()
+    reader_ep = "mydb.cluster-ro-xyz.us-east-1.rds.amazonaws.com"
+    writer = "inst-1.xyz.us-east-1.rds.amazonaws.com"
+    reader_1 = "inst-2.xyz.us-east-1.rds.amazonaws.com"
+    reader_2 = "inst-3.xyz.us-east-1.rds.amazonaws.com"
+    topology = (HostInfo(writer, 5432, role=HostRole.WRITER),
+                HostInfo(reader_1, 5432, role=HostRole.READER),
+                HostInfo(reader_2, 5432, role=HostRole.READER))
+
+    async def _refresh(*a, **k):
+        svc._all_hosts = topology
+    svc.refresh_host_list = AsyncMock(side_effect=_refresh)
+    svc.force_refresh_host_list = AsyncMock(side_effect=_refresh)
+    svc.identify_connection = AsyncMock(return_value=HostInfo(reader_1, 5432))
+    svc.get_host_role = AsyncMock(return_value=HostRole.READER)
+
+    writer_conn = _plain_conn("writer")
+    reader_conn = _plain_conn("reader")
+
+    async def _run():
+        tracker.track(HostInfo(writer, 5432), writer_conn)
+        await _connect_via(plugin, svc, driver_dialect, reader_ep, reader_conn)
+        assert plugin._current_writer.host == writer
+
+        async def _raising():
+            svc._current_host_info = HostInfo(reader_2, 5432, role=HostRole.READER)
+            raise FailoverSuccessError("failover")
+        with pytest.raises(FailoverSuccessError):
+            await plugin.execute(MagicMock(), "Cursor.execute", _raising)
+        writer_conn.close.assert_not_called()
+
+    asyncio.run(_run())