Integrate client routes handler into Cluster and ControlConnection

Cluster:
- Add client_routes_config parameter with mutual exclusivity check
  against endpoint_factory
- Create _ClientRoutesHandler and ClientRoutesEndPointFactory when
  client_routes_config is provided

ControlConnection:
- Register CLIENT_ROUTES_CHANGE event watcher when handler is present
- Forward events to handler via _handle_client_routes_change
- Trigger full route re-read on control connection reconnection

Author: sylwiaszunejko

cassandra/cluster.py

@@ -29,7 +29,7 @@
 from itertools import groupby, count, chain
 import json
 import logging
-from typing import Optional, Union
+from typing import Any, Dict, Optional, Union
 from warnings import warn
 from random import random
 import re
@@ -48,7 +48,8 @@
                        SchemaTargetType, DriverException, ProtocolVersion,
                        UnresolvableContactPoints, DependencyException)
 from cassandra.auth import _proxy_execute_key, PlainTextAuthProvider
-from cassandra.connection import (ConnectionException, ConnectionShutdown,
+from cassandra.client_routes import ClientRoutesChangeType, ClientRoutesConfig, _ClientRoutesHandler
+from cassandra.connection import (ClientRoutesEndPointFactory, ConnectionException, ConnectionShutdown,
                                   ConnectionHeartbeat, ProtocolVersionUnsupported,
                                   EndPoint, DefaultEndPoint, DefaultEndPointFactory,
                                   SniEndPointFactory, ConnectionBusy, locally_supported_compressions)
@@ -1215,7 +1216,8 @@ def __init__(self,
                  shard_aware_options=None,
                  metadata_request_timeout: Optional[float] = None,
                  column_encryption_policy=None,
-                 application_info:Optional[ApplicationInfoBase]=None
+                 application_info:Optional[ApplicationInfoBase]=None,
+                 client_routes_config:Optional[ClientRoutesConfig]=None
                  ):
         """
         ``executor_threads`` defines the number of threads in a pool for handling asynchronous tasks such as
@@ -1280,6 +1282,45 @@ def __init__(self,
         if column_encryption_policy is not None:
             self.column_encryption_policy = column_encryption_policy
 
+        if client_routes_config is not None and endpoint_factory is not None:
+            raise ValueError("client_routes_config and endpoint_factory are mutually exclusive")
+
+        self._client_routes_handler = None
+        if client_routes_config is not None:
+            if not isinstance(client_routes_config, ClientRoutesConfig):
+                raise TypeError("client_routes_config must be a ClientRoutesConfig instance")
+
+            # SSL hostname verification is incompatible with client routes:
+            # connections go through NLB proxies whose addresses won't match
+            # server certificates.
+            _check_hostname_enabled = False
+            if ssl_context is not None and ssl_context.check_hostname:
+                _check_hostname_enabled = True
+            if ssl_options is not None and ssl_options.get('check_hostname', False):
+                _check_hostname_enabled = True
+            if _check_hostname_enabled:
+                raise ValueError(
+                    "SSL hostname verification (check_hostname=True) is currently incompatible "
+                    "with client_routes_config. When using client routes, connections "
+                    "go through NLB proxies whose addresses won't match server "
+                    "certificates. Disable hostname verification by setting "
+                    "ssl_context.check_hostname = False."
+                )
+
+            ssl_enabled = ssl_context is not None or ssl_options is not None
+            self._client_routes_handler = _ClientRoutesHandler(client_routes_config, ssl_enabled=ssl_enabled)
+
+            if contact_points is _NOT_SET or not self._contact_points_explicit:
+                seed_addrs = [dep.connection_addr_override for dep in client_routes_config.proxies
+                             if dep.connection_addr_override]
+                if seed_addrs:
+                    self.contact_points = seed_addrs
+                    self._contact_points_explicit = True
+                    log.info("[client routes] Using %d deployment connection addresses as contact points",
+                            len(seed_addrs))
+
+        if self._client_routes_handler is not None:
+            endpoint_factory = ClientRoutesEndPointFactory(self._client_routes_handler, self.port)
         self.endpoint_factory = endpoint_factory or DefaultEndPointFactory(port=self.port)
         self.endpoint_factory.configure(self)
 
@@ -1437,6 +1478,10 @@ def __init__(self,
         self.monitor_reporting_interval = monitor_reporting_interval
         self.shard_aware_options = ShardAwareOptions(opts=shard_aware_options)
 
+        if (client_routes_config is not None
+                and not client_routes_config.advanced_shard_awareness):
+            self.shard_aware_options.disable_shardaware_port = True
+
         self._listeners = set()
         self._listener_lock = Lock()
 
@@ -3612,11 +3657,21 @@ def _try_connect(self, endpoint):
         # this object (after a dereferencing a weakref)
         self_weakref = weakref.ref(self, partial(_clear_watcher, weakref.proxy(connection)))
         try:
-            connection.register_watchers({
+            watchers = {
                 "TOPOLOGY_CHANGE": partial(_watch_callback, self_weakref, '_handle_topology_change'),
                 "STATUS_CHANGE": partial(_watch_callback, self_weakref, '_handle_status_change'),
                 "SCHEMA_CHANGE": partial(_watch_callback, self_weakref, '_handle_schema_change')
-            }, register_timeout=self._timeout)
+            }
+
+            if self._cluster._client_routes_handler is not None:
+                watchers["CLIENT_ROUTES_CHANGE"] = partial(_watch_callback, self_weakref, '_handle_client_routes_change')
+
+            connection.register_watchers(watchers, register_timeout=self._timeout)
+
+            if self._cluster._client_routes_handler is not None:
+                self._cluster._client_routes_handler.initialize(
+                    connection,
+                    self._timeout)
 
             sel_peers = self._get_peers_query(self.PeersQueryType.PEERS, connection)
             sel_local = self._SELECT_LOCAL if self._token_meta_enabled else self._SELECT_LOCAL_NO_TOKENS
@@ -3979,6 +4034,44 @@ def _handle_status_change(self, event):
                 # this will be run by the scheduler
                 self._cluster.on_down(host, is_host_addition=False)
 
+    def _handle_client_routes_change(self, event: Dict[str, Any]) -> None:
+        """
+        Handle CLIENT_ROUTES_CHANGE event from the server.
+
+        This event indicates that the system.client_routes table has been updated
+        and we need to refresh our route mappings.
+        """
+        if self._cluster._client_routes_handler is None:
+            log.warning("[control connection] Received CLIENT_ROUTES_CHANGE but no handler configured")
+            return
+
+        raw_change_type = event.get("change_type")
+        try:
+            change_type = ClientRoutesChangeType(raw_change_type)
+        except ValueError:
+            log.warning("[control connection] Unknown CLIENT_ROUTES_CHANGE type: %s", raw_change_type)
+            return
+
+        connection_ids = tuple(event.get("connection_ids", []))
+        host_ids = tuple(event.get("host_ids", []))
+
+        self._cluster.scheduler.schedule_unique(
+            0,
+            self._handle_client_routes_refresh,
+            self._connection, self._timeout, change_type, connection_ids, host_ids
+        )
+
+    def _handle_client_routes_refresh(self, connection, timeout,
+                                      change_type, connection_ids, host_ids):
+        try:
+            self._cluster._client_routes_handler.handle_client_routes_change(
+                connection, timeout, change_type, connection_ids, host_ids)
+        except ReferenceError:
+            pass  # our weak reference to the Cluster is no good
+        except Exception:
+            log.debug("[control connection] Error handling CLIENT_ROUTES_CHANGE", exc_info=True)
+            self._signal_error()
+
     def _handle_schema_change(self, event):
         if self._schema_event_refresh_window < 0:
             return