|
29 | 29 | from itertools import groupby, count, chain |
30 | 30 | import json |
31 | 31 | import logging |
32 | | -from typing import Optional, Union |
| 32 | +from typing import Any, Dict, Optional, Union |
33 | 33 | from warnings import warn |
34 | 34 | from random import random |
35 | 35 | import re |
|
48 | 48 | SchemaTargetType, DriverException, ProtocolVersion, |
49 | 49 | UnresolvableContactPoints, DependencyException) |
50 | 50 | from cassandra.auth import _proxy_execute_key, PlainTextAuthProvider |
51 | | -from cassandra.connection import (ConnectionException, ConnectionShutdown, |
| 51 | +from cassandra.client_routes import ClientRoutesChangeType, ClientRoutesConfig, _ClientRoutesHandler |
| 52 | +from cassandra.connection import (ClientRoutesEndPointFactory, ConnectionException, ConnectionShutdown, |
52 | 53 | ConnectionHeartbeat, ProtocolVersionUnsupported, |
53 | 54 | EndPoint, DefaultEndPoint, DefaultEndPointFactory, |
54 | 55 | SniEndPointFactory, ConnectionBusy, locally_supported_compressions) |
@@ -1215,7 +1216,8 @@ def __init__(self, |
1215 | 1216 | shard_aware_options=None, |
1216 | 1217 | metadata_request_timeout: Optional[float] = None, |
1217 | 1218 | column_encryption_policy=None, |
1218 | | - application_info:Optional[ApplicationInfoBase]=None |
| 1219 | + application_info:Optional[ApplicationInfoBase]=None, |
| 1220 | + client_routes_config:Optional[ClientRoutesConfig]=None |
1219 | 1221 | ): |
1220 | 1222 | """ |
1221 | 1223 | ``executor_threads`` defines the number of threads in a pool for handling asynchronous tasks such as |
@@ -1280,6 +1282,45 @@ def __init__(self, |
1280 | 1282 | if column_encryption_policy is not None: |
1281 | 1283 | self.column_encryption_policy = column_encryption_policy |
1282 | 1284 |
|
| 1285 | + if client_routes_config is not None and endpoint_factory is not None: |
| 1286 | + raise ValueError("client_routes_config and endpoint_factory are mutually exclusive") |
| 1287 | + |
| 1288 | + self._client_routes_handler = None |
| 1289 | + if client_routes_config is not None: |
| 1290 | + if not isinstance(client_routes_config, ClientRoutesConfig): |
| 1291 | + raise TypeError("client_routes_config must be a ClientRoutesConfig instance") |
| 1292 | + |
| 1293 | + # SSL hostname verification is incompatible with client routes: |
| 1294 | + # connections go through NLB proxies whose addresses won't match |
| 1295 | + # server certificates. |
| 1296 | + _check_hostname_enabled = False |
| 1297 | + if ssl_context is not None and ssl_context.check_hostname: |
| 1298 | + _check_hostname_enabled = True |
| 1299 | + if ssl_options is not None and ssl_options.get('check_hostname', False): |
| 1300 | + _check_hostname_enabled = True |
| 1301 | + if _check_hostname_enabled: |
| 1302 | + raise ValueError( |
| 1303 | + "SSL hostname verification (check_hostname=True) is currently incompatible " |
| 1304 | + "with client_routes_config. When using client routes, connections " |
| 1305 | + "go through NLB proxies whose addresses won't match server " |
| 1306 | + "certificates. Disable hostname verification by setting " |
| 1307 | + "ssl_context.check_hostname = False." |
| 1308 | + ) |
| 1309 | + |
| 1310 | + ssl_enabled = ssl_context is not None or ssl_options is not None |
| 1311 | + self._client_routes_handler = _ClientRoutesHandler(client_routes_config, ssl_enabled=ssl_enabled) |
| 1312 | + |
| 1313 | + if contact_points is _NOT_SET or not self._contact_points_explicit: |
| 1314 | + seed_addrs = [dep.connection_addr for dep in client_routes_config.proxies |
| 1315 | + if dep.connection_addr] |
| 1316 | + if seed_addrs: |
| 1317 | + self.contact_points = seed_addrs |
| 1318 | + self._contact_points_explicit = True |
| 1319 | + log.info("[client routes] Using %d deployment connection addresses as contact points", |
| 1320 | + len(seed_addrs)) |
| 1321 | + |
| 1322 | + if self._client_routes_handler is not None: |
| 1323 | + endpoint_factory = ClientRoutesEndPointFactory(self._client_routes_handler, self.port) |
1283 | 1324 | self.endpoint_factory = endpoint_factory or DefaultEndPointFactory(port=self.port) |
1284 | 1325 | self.endpoint_factory.configure(self) |
1285 | 1326 |
|
@@ -3612,11 +3653,24 @@ def _try_connect(self, endpoint): |
3612 | 3653 | # this object (after a dereferencing a weakref) |
3613 | 3654 | self_weakref = weakref.ref(self, partial(_clear_watcher, weakref.proxy(connection))) |
3614 | 3655 | try: |
3615 | | - connection.register_watchers({ |
| 3656 | + watchers = { |
3616 | 3657 | "TOPOLOGY_CHANGE": partial(_watch_callback, self_weakref, '_handle_topology_change'), |
3617 | 3658 | "STATUS_CHANGE": partial(_watch_callback, self_weakref, '_handle_status_change'), |
3618 | 3659 | "SCHEMA_CHANGE": partial(_watch_callback, self_weakref, '_handle_schema_change') |
3619 | | - }, register_timeout=self._timeout) |
| 3660 | + } |
| 3661 | + |
| 3662 | + if self._cluster._client_routes_handler is not None: |
| 3663 | + watchers["CLIENT_ROUTES_CHANGE"] = partial(_watch_callback, self_weakref, '_handle_client_routes_change') |
| 3664 | + try: |
| 3665 | + self._cluster._client_routes_handler.initialize( |
| 3666 | + connection, |
| 3667 | + self._timeout) |
| 3668 | + except Exception as e: |
| 3669 | + log.error("[control connection] Failed to initialize client routes handler: %s", e, exc_info=True) |
| 3670 | + connection.close() |
| 3671 | + raise |
| 3672 | + |
| 3673 | + connection.register_watchers(watchers, register_timeout=self._timeout) |
3620 | 3674 |
|
3621 | 3675 | sel_peers = self._get_peers_query(self.PeersQueryType.PEERS, connection) |
3622 | 3676 | sel_local = self._SELECT_LOCAL if self._token_meta_enabled else self._SELECT_LOCAL_NO_TOKENS |
@@ -3658,6 +3712,13 @@ def _reconnect(self): |
3658 | 3712 | log.debug("[control connection] Attempting to reconnect") |
3659 | 3713 | try: |
3660 | 3714 | self._set_new_connection(self._reconnect_internal()) |
| 3715 | + |
| 3716 | + if self._cluster._client_routes_handler is not None: |
| 3717 | + try: |
| 3718 | + self._cluster._client_routes_handler.initialize( |
| 3719 | + self._connection, self._timeout) |
| 3720 | + except Exception as e: |
| 3721 | + log.warning("[control connection] Failed to notify client routes handler of reconnection: %s", e) |
3661 | 3722 | except NoHostAvailable: |
3662 | 3723 | # make a retry schedule (which includes backoff) |
3663 | 3724 | schedule = self._cluster.reconnection_policy.new_schedule() |
@@ -3979,6 +4040,34 @@ def _handle_status_change(self, event): |
3979 | 4040 | # this will be run by the scheduler |
3980 | 4041 | self._cluster.on_down(host, is_host_addition=False) |
3981 | 4042 |
|
| 4043 | + def _handle_client_routes_change(self, event: Dict[str, Any]) -> None: |
| 4044 | + """ |
| 4045 | + Handle CLIENT_ROUTES_CHANGE event from the server. |
| 4046 | +
|
| 4047 | + This event indicates that the system.client_routes table has been updated |
| 4048 | + and we need to refresh our route mappings. |
| 4049 | + """ |
| 4050 | + if self._cluster._client_routes_handler is None: |
| 4051 | + log.warning("[control connection] Received CLIENT_ROUTES_CHANGE but no handler configured") |
| 4052 | + return |
| 4053 | + |
| 4054 | + raw_change_type = event.get("change_type") |
| 4055 | + try: |
| 4056 | + change_type = ClientRoutesChangeType(raw_change_type) |
| 4057 | + except ValueError: |
| 4058 | + log.warning("[control connection] Unknown CLIENT_ROUTES_CHANGE type: %s", raw_change_type) |
| 4059 | + return |
| 4060 | + |
| 4061 | + connection_ids = tuple(event.get("connection_ids", [])) |
| 4062 | + host_ids = tuple(event.get("host_ids", [])) |
| 4063 | + |
| 4064 | + # Handle the event asynchronously |
| 4065 | + self._cluster.scheduler.schedule_unique( |
| 4066 | + 0, |
| 4067 | + self._cluster._client_routes_handler.handle_client_routes_change, |
| 4068 | + self._connection, self._timeout, change_type, connection_ids, host_ids |
| 4069 | + ) |
| 4070 | + |
3982 | 4071 | def _handle_schema_change(self, event): |
3983 | 4072 | if self._schema_event_refresh_window < 0: |
3984 | 4073 | return |
|
0 commit comments