client-routes: preserve mixed event route state

Author: dkropachev

cassandra/client_routes.py

@@ -294,7 +294,7 @@ def handle_client_routes_change(self, connection: 'Connection', timeout: float,
             return
 
         routes = self._query_routes_for_change_event(connection, timeout, pairs)
-        self._routes.merge(routes, affected_host_ids=set(host_uuids))
+        self._routes.merge(routes, affected_host_ids={host_id for _, host_id in pairs})
 
     def _query_all_routes_for_connections(self, connection: 'Connection', timeout: float,
                                           connection_ids: Set[str]) -> List[_Route]:

cassandra/connection.py

@@ -468,21 +468,25 @@ def resolve(self) -> Tuple[str, int]:
     def __eq__(self, other):
         return (isinstance(other, ClientRoutesEndPoint) and
                 self._host_id == other._host_id and
-                self._original_address == other._original_address)
+                self._original_address == other._original_address and
+                self._original_port == other._original_port)
 
     def __hash__(self):
-        return hash((self._host_id, self._original_address))
+        return hash((self._host_id, self._original_address, self._original_port))
+
+    def _comparison_key(self):
+        return (self._host_id, self._original_address,
+                self._original_port is None, self._original_port)
 
     def __lt__(self, other):
-        return ((self._host_id, self._original_address) <
-                (other._host_id, other._original_address))
+        return self._comparison_key() < other._comparison_key()
 
     def __str__(self):
         return str("%s (host_id=%s)" % (self._original_address, self._host_id))
 
     def __repr__(self):
-        return "<%s: host_id=%s, original_addr=%s>" % (
-            self.__class__.__name__, self._host_id, self._original_address)
+        return "<%s: host_id=%s, original_addr=%s, original_port=%s>" % (
+            self.__class__.__name__, self._host_id, self._original_address, self._original_port)
 
 
 class _Frame(object):

tests/unit/test_client_routes.py

@@ -233,6 +233,35 @@ def test_handle_change_merges_when_host_ids_present(self, mock_query):
         self.assertIsNotNone(handler._routes.get_by_host_id(existing_host))
         self.assertIsNotNone(handler._routes.get_by_host_id(new_host))
 
+    @patch.object(_ClientRoutesHandler, '_query_routes_for_change_event')
+    def test_handle_change_preserves_routes_for_unrelated_connection_ids(self, mock_query):
+        """Routes for unrelated connection_ids in mixed events should not be removed."""
+        handler = _ClientRoutesHandler(self.config)
+        mock_conn = Mock()
+
+        conn_id = str(self.conn_id)
+        changed_host = uuid.uuid4()
+        unrelated_host = uuid.uuid4()
+
+        handler._routes.update([
+            _Route(connection_id=conn_id, host_id=changed_host, address="old.com", port=9042),
+            _Route(connection_id=conn_id, host_id=unrelated_host, address="keep.com", port=9042),
+        ])
+
+        mock_query.return_value = [
+            _Route(connection_id=conn_id, host_id=changed_host, address="new.com", port=9042),
+        ]
+
+        handler.handle_client_routes_change(
+            mock_conn, 5.0,
+            ClientRoutesChangeType.UPDATE_NODES,
+            connection_ids=[conn_id, "unrelated-conn-id"],
+            host_ids=[str(changed_host), str(unrelated_host)],
+        )
+
+        self.assertEqual(handler._routes.get_by_host_id(changed_host).address, "new.com")
+        self.assertEqual(handler._routes.get_by_host_id(unrelated_host).address, "keep.com")
+
     @patch.object(_ClientRoutesHandler, '_query_all_routes_for_connections')
     def test_handle_change_updates_when_no_host_ids(self, mock_query):
         """When no host_ids are provided, routes should be fully replaced."""
@@ -388,6 +417,42 @@ def test_resolve_host_missing_port_raises(self):
         with self.assertRaises(ValueError):
             self.handler.resolve_host(host_id)
 
+    def test_endpoint_identity_includes_original_port(self):
+        host_id = uuid.uuid4()
+        first = ClientRoutesEndPoint(
+            host_id=host_id,
+            handler=self.handler,
+            original_address="10.0.0.1",
+            original_port=9042,
+        )
+        second = ClientRoutesEndPoint(
+            host_id=host_id,
+            handler=self.handler,
+            original_address="10.0.0.1",
+            original_port=9142,
+        )
+
+        self.assertNotEqual(first, second)
+        self.assertEqual(len({first, second}), 2)
+
+    def test_endpoint_ordering_handles_missing_original_port(self):
+        host_id = uuid.uuid4()
+        without_port = ClientRoutesEndPoint(
+            host_id=host_id,
+            handler=self.handler,
+            original_address="10.0.0.1",
+            original_port=None,
+        )
+        with_port = ClientRoutesEndPoint(
+            host_id=host_id,
+            handler=self.handler,
+            original_address="10.0.0.1",
+            original_port=9042,
+        )
+
+        self.assertCountEqual(
+            sorted([without_port, with_port]), [without_port, with_port])
+
 
 class TestClientRoutesEndPointFactory(unittest.TestCase):