diff --git a/aws_advanced_python_wrapper/aio/aurora_connection_tracker.py b/aws_advanced_python_wrapper/aio/aurora_connection_tracker.py index 93e287ee..e926a42b 100644 --- a/aws_advanced_python_wrapper/aio/aurora_connection_tracker.py +++ b/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)): diff --git a/tests/unit/test_aio_aurora_connection_tracker.py b/tests/unit/test_aio_aurora_connection_tracker.py index 9056d959..9aecd49f 100644 --- a/tests/unit/test_aio_aurora_connection_tracker.py +++ b/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())