client_routes: separate API methods for complete refresh vs change events

Split _query_routes into two specialized methods:

1. _query_all_routes_for_connections: Used during initialization and
   reconnection for complete refresh of all routes for configured
   connection IDs

2. _query_routes_for_change_event: Used for CLIENT_ROUTES_CHANGE event
   handling with proper AND logic for related connection_ids/host_ids

3. _execute_routes_query: Common helper for query execution and result
   parsing with proxy address override support

This separation clarifies the different use cases and semantics:
- Complete refresh: loads all routes for our connection IDs
- Change events: queries specific routes affected by the event

Updated unit tests to use the appropriate method mocks.
All tests pass.

Author: dkropachev

cassandra/client_routes.py

@@ -242,8 +242,8 @@ def initialize(self, control_connection: 'ControlConnection') -> None:
         log.info("[client routes] Initializing with %d proxies", len(self.config.proxies))
 
         try:
-            connection_ids = self._connection_ids
-            routes = self._query_routes(control_connection, connection_ids=connection_ids)
+            connection_ids = list(self._connection_ids)
+            routes = self._query_all_routes_for_connections(control_connection, connection_ids)
 
             self._routes.update(routes)
         except Exception as e:
@@ -264,30 +264,36 @@ def handle_client_routes_change(self, control_connection: 'ControlConnection',
         :param connection_ids: Affected connection ID strings; empty means all.
         :param host_ids: Affected host ID strings; empty means all.
         """
-        filtered_conn_ids = None
-        if connection_ids:
-            configured_ids = self._connection_ids
-            filtered = [cid for cid in connection_ids if cid in configured_ids]
-            if not filtered:
+
+        try:
+            # Both arrays must be present and same length, or both empty
+
+            full_refresh = False
+            if not connection_ids and not host_ids:
+                log.error(
+                    "[client routes] CLIENT_ROUTES_CHANGE has no connection_ids and host_ids, doing full refresh")
+                full_refresh = True
+            elif len(connection_ids) != len(host_ids):
+                log.error("[client routes] CLIENT_ROUTES_CHANGE has mismatched lengths (conn: %d, host: %d), doing full refresh",
+                         len(connection_ids), len(host_ids))
+                full_refresh = True
+
+            if full_refresh:
+                routes = self._query_all_routes_for_connections(control_connection, list(self._connection_ids))
+                self._routes.update(routes)
                 return
-            filtered_conn_ids = filtered
 
-        host_uuids = [uuid.UUID(hid) for hid in host_ids] if host_ids else None
+            host_uuids = [uuid.UUID(hid) for hid in host_ids]
+            pairs = [(cid, hid) for cid, hid in zip(connection_ids, host_uuids)
+                     if cid in self._connection_ids]
 
-        try:
-            routes = self._query_routes(
-                control_connection,
-                connection_ids=filtered_conn_ids,
-                host_ids=host_uuids
-            )
-        except Exception as e:
-            log.warning("[client routes] Failed to query routes for CLIENT_ROUTES_CHANGE: %s", e, exc_info=True)
-            return
+            if not pairs:
+                return  # No relevant connection IDs
 
-        if host_uuids:
+            routes = self._query_routes_for_change_event(control_connection, pairs)
             self._routes.merge(routes, affected_host_ids=set(host_uuids))
-        else:
-            self._routes.update(routes)
+        except Exception as e:
+            log.warning("[client routes] Failed to handle CLIENT_ROUTES_CHANGE: %s", e, exc_info=True)
 
     def handle_control_connection_reconnect(self, control_connection: 'ControlConnection') -> None:
         """
@@ -307,57 +313,80 @@ def handle_control_connection_reconnect(self, control_connection: 'ControlConnec
         """
         log.info("[client routes] Control connection reconnected, reloading all routes")
 
-        connection_ids = self._connection_ids
-        routes = self._query_routes(control_connection, connection_ids=connection_ids)
+        connection_ids = list(self._connection_ids)
+        routes = self._query_all_routes_for_connections(control_connection, connection_ids)
         self._routes.update(routes)
 
-    def _query_routes(self, control_connection: 'ControlConnection', connection_ids: Optional[List[str]] = None,
-                     host_ids: Optional[List[uuid.UUID]] = None) -> List[_Route]:
+    def _query_all_routes_for_connections(self, control_connection: 'ControlConnection', 
+                                          connection_ids: List[str]) -> List[_Route]:
         """
-        Query system.client_routes table.
-
-        When both connection_ids and host_ids are provided, uses AND logic:
-        returns routes where connection_id IN (...) AND host_id IN (...).
-        This matches the semantics of CLIENT_ROUTES_CHANGE events which contain
-        related connection_ids and host_ids from the same operation.
-
+        Query all routes for the given connection IDs (complete refresh).
+        
+        Used when control connection reconnects or as a fallback when
+        CLIENT_ROUTES_CHANGE event has malformed data.
+        
         :param control_connection: ControlConnection to execute query
-        :param connection_ids: Optional list of connection ID strings to filter by
-        :param host_ids: Optional list of host UUIDs to filter by
+        :param connection_ids: List of connection ID strings
         :return: List of _Route
         """
-        query = "SELECT connection_id, host_id, address, port, tls_port FROM system.client_routes"
-        where_clauses = []
+        if not connection_ids:
+            return []
+            
+        # Deduplicate connection_ids to avoid redundant parameters
+        unique_connection_ids = list(dict.fromkeys(connection_ids))
+        placeholders = ', '.join('?' for _ in unique_connection_ids)
+        query = f"SELECT connection_id, host_id, address, port, tls_port FROM system.client_routes WHERE connection_id IN ({placeholders})"
+        params = [cid.encode('utf-8') for cid in unique_connection_ids]
+        
+        log.debug("[client routes] Querying all routes for connection_ids=%s", unique_connection_ids)
+        return self._execute_routes_query(control_connection, query, params)
+    
+    def _query_routes_for_change_event(self, control_connection: 'ControlConnection', 
+                                       route_pairs: List[Tuple[str, uuid.UUID]]) -> List[_Route]:
+        """
+        Query specific routes affected by a CLIENT_ROUTES_CHANGE event.
+        
+        Takes a list of (connection_id, host_id) pairs that represent the exact
+        routes affected by an operation. This provides precise updates without
+        fetching unrelated routes.
+        
+        If the pairs list is empty or None, falls back to a complete refresh
+        of all routes for safety.
+        
+        :param control_connection: ControlConnection to execute query
+        :param route_pairs: List of (connection_id, host_id) tuples
+        :return: List of _Route
+        """
+        # Deduplicate pairs while preserving order
+        unique_pairs = list(dict.fromkeys(route_pairs))
+        
+        # Build query with compound WHERE clause: WHERE (connection_id, host_id) IN ((?, ?), (?, ?), ...)
+        placeholders = ', '.join('(?, ?)' for _ in unique_pairs)
+        query = f"SELECT connection_id, host_id, address, port, tls_port FROM system.client_routes WHERE (connection_id, host_id) IN ({placeholders})"
+        
+        # Flatten pairs into params list: [conn_id1, host_id1, conn_id2, host_id2, ...]
         params = []
-
-        if connection_ids:
-            # Deduplicate connection_ids to avoid redundant parameters
-            unique_connection_ids = list(dict.fromkeys(connection_ids))
-            placeholders = ', '.join('?' for _ in unique_connection_ids)
-            where_clauses.append("connection_id IN (%s)" % placeholders)
-            params.extend(cid.encode('utf-8') for cid in unique_connection_ids)
-
-        if host_ids:
-            # Deduplicate host_ids to avoid redundant parameters  
-            unique_host_ids = list(dict.fromkeys(host_ids))
-            placeholders = ', '.join('?' for _ in unique_host_ids)
-            where_clauses.append("host_id IN (%s)" % placeholders)
-            params.extend(hid.bytes for hid in unique_host_ids)
-
-        if where_clauses:
-            query += " WHERE " + " AND ".join(where_clauses)
-        else:
-            query += " ALLOW FILTERING"
-
-        # Log the query for debugging CLIENT_ROUTES_CHANGE event handling
-        if connection_ids and host_ids:
-            log.debug("[client routes] Querying with AND logic: connection_ids=%s, host_ids=%s", 
-                     connection_ids, host_ids)
-        elif connection_ids:
-            log.debug("[client routes] Querying by connection_ids=%s", connection_ids)
-        elif host_ids:
-            log.debug("[client routes] Querying by host_ids=%s", host_ids)
-
+        for conn_id, host_id in unique_pairs:
+            params.append(conn_id.encode('utf-8'))
+            params.append(host_id.bytes)
+        
+        log.debug("[client routes] Querying specific route pairs from CLIENT_ROUTES_CHANGE: %s", unique_pairs[:5])  # Log first 5 pairs for debugging
+        return self._execute_routes_query(control_connection, query, params)
+    
+    def _execute_routes_query(self, control_connection: 'ControlConnection', 
+                             query: str, params: List) -> List[_Route]:
+        """
+        Execute a routes query and parse results.
+        
+        Common helper for both complete refresh and change event queries.
+        
+        :param control_connection: ControlConnection to execute query
+        :param query: CQL query string
+        :param params: Query parameters
+        :return: List of _Route
+        """
+        log.debug("[client routes] Executing query: %s with %d parameters", query, len(params))
+        
         query_msg = QueryMessage(query=query, consistency_level=ConsistencyLevel.ONE,
                                  query_params=params if params else None)
         result = control_connection._connection.wait_for_response(

tests/unit/test_client_routes.py

@@ -165,7 +165,7 @@ def test_handler_initialization(self):
         self.assertIsNotNone(handler)
         self.assertEqual(handler.ssl_enabled, False)
 
-    @patch.object(_ClientRoutesHandler, '_query_routes')
+    @patch.object(_ClientRoutesHandler, '_query_all_routes_for_connections')
     def test_initialize(self, mock_query):
         host_id = uuid.uuid4()
         mock_query.return_value = [
@@ -187,7 +187,7 @@ def test_initialize(self, mock_query):
         self.assertIsNotNone(route)
         self.assertEqual(route.address, "node1.example.com")
 
-    @patch.object(_ClientRoutesHandler, '_query_routes')
+    @patch.object(_ClientRoutesHandler, '_query_routes_for_change_event')
     def test_handle_change_filters_by_configured_connection_ids(self, mock_query):
         """Events with unrelated connection_ids should be ignored."""
         handler = _ClientRoutesHandler(self.config)
@@ -203,7 +203,7 @@ def test_handle_change_filters_by_configured_connection_ids(self, mock_query):
         )
         mock_query.assert_not_called()
 
-    @patch.object(_ClientRoutesHandler, '_query_routes')
+    @patch.object(_ClientRoutesHandler, '_query_routes_for_change_event')
     def test_handle_change_merges_when_host_ids_present(self, mock_query):
         """When host_ids are provided, routes should be merged (not full replace)."""
         handler = _ClientRoutesHandler(self.config)
@@ -233,7 +233,7 @@ 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')
+    @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."""
         handler = _ClientRoutesHandler(self.config)
@@ -261,7 +261,7 @@ def test_handle_change_updates_when_no_host_ids(self, mock_query):
         self.assertIsNone(handler._routes.get_by_host_id(old_host))
         self.assertIsNotNone(handler._routes.get_by_host_id(new_host))
 
-    @patch.object(_ClientRoutesHandler, '_query_routes')
+    @patch.object(_ClientRoutesHandler, '_query_routes_for_change_event')
     def test_handle_change_swallows_query_failure(self, mock_query):
         """If _query_routes raises, handle_client_routes_change should not propagate."""
         handler = _ClientRoutesHandler(self.config)
@@ -276,7 +276,7 @@ def test_handle_change_swallows_query_failure(self, mock_query):
             host_ids=[str(uuid.uuid4())],
         )
 
-    @patch.object(_ClientRoutesHandler, '_query_routes')
+    @patch.object(_ClientRoutesHandler, '_query_all_routes_for_connections')
     def test_reconnect_propagates_exception_on_failure(self, mock_query):
         """handle_control_connection_reconnect should propagate exceptions to caller."""
         handler = _ClientRoutesHandler(self.config)
@@ -288,7 +288,7 @@ def test_reconnect_propagates_exception_on_failure(self, mock_query):
         self.assertIn("query failed", str(ctx.exception))
         self.assertEqual(mock_query.call_count, 1)
 
-    @patch.object(_ClientRoutesHandler, '_query_routes')
+    @patch.object(_ClientRoutesHandler, '_query_all_routes_for_connections')
     def test_reconnect_keeps_old_routes_on_failure(self, mock_query):
         """On failure, existing routes must be preserved (critical for PL clusters)."""
         handler = _ClientRoutesHandler(self.config)
@@ -307,7 +307,7 @@ def test_reconnect_keeps_old_routes_on_failure(self, mock_query):
         # Old route must still be there
         self.assertIsNotNone(handler._routes.get_by_host_id(host_id))
 
-    @patch.object(_ClientRoutesHandler, '_query_routes')
+    @patch.object(_ClientRoutesHandler, '_query_all_routes_for_connections')
     def test_reconnect_updates_routes_on_success(self, mock_query):
         """handle_control_connection_reconnect should update routes on success."""
         handler = _ClientRoutesHandler(self.config)