diff --git a/cassandra/cluster.py b/cassandra/cluster.py index f2e111894b..b3542f15ba 100644 --- a/cassandra/cluster.py +++ b/cassandra/cluster.py @@ -63,7 +63,8 @@ BatchMessage, RESULT_KIND_PREPARED, RESULT_KIND_SET_KEYSPACE, RESULT_KIND_ROWS, RESULT_KIND_SCHEMA_CHANGE, ProtocolHandler, - RESULT_KIND_VOID, ProtocolException) + RESULT_KIND_VOID, ProtocolException, + StartupMessage) from cassandra.metadata import Metadata, protect_name, murmur3, _NodeInfo from cassandra.policies import (TokenAwarePolicy, DCAwareRoundRobinPolicy, SimpleConvictionPolicy, ExponentialReconnectionPolicy, HostDistance, @@ -972,6 +973,21 @@ def default_retry_policy(self, policy): used for columns in this cluster. """ + extra_startup_options: dict[str, str] | None = None + """ + A dict of extra options sent in the STARTUP message when a connection is established. + + This is useful for sending custom startup options that are not supported by the driver. + For example, per `CASSANDRA-16378 `_, + custom application level options are exposed in client metrics:: + + Cluster(extra_startup_options={'APPLICATION_NAME': 'my-app'}) + + Options managed by the driver, such as ``CQL_VERSION``, ``COMPRESSION``, ``NO_COMPACT``, + ``DRIVER_NAME`` and ``DRIVER_VERSION``, always take precedence. Extra options that collide + with them are ignored. + """ + @property def schema_metadata_enabled(self): """ @@ -1069,7 +1085,8 @@ def __init__(self, ssl_context=None, endpoint_factory=None, cloud=None, - column_encryption_policy=None): + column_encryption_policy=None, + extra_startup_options=None): """ ``executor_threads`` defines the number of threads in a pool for handling asynchronous tasks such as establishing connection pools or refreshing metadata. @@ -1266,6 +1283,13 @@ def __init__(self, self.connect_timeout = connect_timeout self.prepare_on_all_hosts = prepare_on_all_hosts self.reprepare_on_up = reprepare_on_up + self.extra_startup_options = dict(extra_startup_options) if extra_startup_options else {} + driver_managed_options = self.extra_startup_options.keys() & StartupMessage.KNOWN_OPTION_KEYS + if driver_managed_options: + log.warning("Ignoring extra startup option(s) %s: they are managed by the driver " + "and cannot be overridden", ', '.join(sorted(driver_managed_options))) + for key in driver_managed_options: + del self.extra_startup_options[key] self._listeners = set() self._listener_lock = Lock() @@ -1556,6 +1580,7 @@ def _make_connection_kwargs(self, endpoint, kwargs_dict): kwargs_dict.setdefault('user_type_map', self._user_types) kwargs_dict.setdefault('allow_beta_protocol_version', self.allow_beta_protocol_version) kwargs_dict.setdefault('no_compact', self.no_compact) + kwargs_dict.setdefault('extra_startup_options', self.extra_startup_options) return kwargs_dict diff --git a/cassandra/connection.py b/cassandra/connection.py index d9bca62718..394a06dd3a 100644 --- a/cassandra/connection.py +++ b/cassandra/connection.py @@ -675,6 +675,7 @@ class Connection(object): cql_version = None no_compact = False + extra_startup_options = None protocol_version = ProtocolVersion.MAX_SUPPORTED keyspace = None @@ -758,7 +759,7 @@ def __init__(self, host='127.0.0.1', port=9042, authenticator=None, ssl_options=None, sockopts=None, compression=True, cql_version=None, protocol_version=ProtocolVersion.MAX_SUPPORTED, is_control_connection=False, user_type_map=None, connect_timeout=None, allow_beta_protocol_version=False, no_compact=False, - ssl_context=None, on_orphaned_stream_released=None): + ssl_context=None, on_orphaned_stream_released=None, extra_startup_options=None): # TODO next major rename host to endpoint and remove port kwarg. self.endpoint = host if isinstance(host, EndPoint) else DefaultEndPoint(host, port) @@ -782,6 +783,7 @@ def __init__(self, host='127.0.0.1', port=9042, authenticator=None, self._socket_writable = True self.orphaned_request_ids = set() self._on_orphaned_stream_released = on_orphaned_stream_released + self.extra_startup_options = extra_startup_options or {} if ssl_options: self.ssl_options.update(self.endpoint.ssl_options or {}) @@ -1403,7 +1405,7 @@ def _send_startup_message(self, compression=None, no_compact=False): opts['COMPRESSION'] = compression if no_compact: opts['NO_COMPACT'] = 'true' - sm = StartupMessage(cqlversion=self.cql_version, options=opts) + sm = StartupMessage(cqlversion=self.cql_version, options=opts, extra_options=self.extra_startup_options) self.send_msg(sm, self.get_request_id(), cb=self._handle_startup_response) log.debug("Sent StartupMessage on %s", self) diff --git a/cassandra/protocol.py b/cassandra/protocol.py index b1c4183cf4..24ec7f2f09 100644 --- a/cassandra/protocol.py +++ b/cassandra/protocol.py @@ -400,15 +400,19 @@ class StartupMessage(_MessageType): KNOWN_OPTION_KEYS = set(( 'CQL_VERSION', 'COMPRESSION', - 'NO_COMPACT' + 'NO_COMPACT', + 'DRIVER_NAME', + 'DRIVER_VERSION' )) - def __init__(self, cqlversion, options): + def __init__(self, cqlversion, options, extra_options=None): self.cqlversion = cqlversion self.options = options + self.extra_options = extra_options def send_body(self, f, protocol_version): - optmap = self.options.copy() + optmap: dict[str, str] = self.options.copy() + optmap.update(self.extra_options or {}) optmap['CQL_VERSION'] = self.cqlversion write_stringmap(f, optmap) diff --git a/docs/api/cassandra/cluster.rst b/docs/api/cassandra/cluster.rst index a9a9d378a4..f1a7ee7a84 100644 --- a/docs/api/cassandra/cluster.rst +++ b/docs/api/cassandra/cluster.rst @@ -74,6 +74,8 @@ .. autoattribute:: cloud + .. autoattribute:: extra_startup_options + .. automethod:: connect .. automethod:: shutdown diff --git a/tests/integration/standard/test_connection.py b/tests/integration/standard/test_connection.py index 88788a4ce2..d9d760ded5 100644 --- a/tests/integration/standard/test_connection.py +++ b/tests/integration/standard/test_connection.py @@ -23,16 +23,19 @@ import threading from threading import Thread, Event import time +import uuid from unittest import SkipTest from cassandra import ConsistencyLevel, OperationTimedOut, DependencyException from cassandra.cluster import NoHostAvailable, ConnectionShutdown, ExecutionProfile, EXEC_PROFILE_DEFAULT +from cassandra.connection import DRIVER_NAME, DRIVER_VERSION from cassandra.protocol import QueryMessage from cassandra.policies import HostFilterPolicy, RoundRobinPolicy, HostStateListener from cassandra.pool import HostConnectionPool from tests.integration import use_singledc, get_node, CASSANDRA_IP, local, \ - requiresmallclockgranularity, greaterthancass20, TestCluster + requiresmallclockgranularity, greaterthancass20, greaterthanorequalcass41, \ + requirecassandra, TestCluster try: import cassandra.io.asyncorereactor @@ -465,3 +468,57 @@ def setUp(self): def clean_global_loop(self): cassandra.io.libevreactor._global_loop._cleanup() cassandra.io.libevreactor._global_loop = None + + +@requirecassandra +@greaterthanorequalcass41 +class ExtraStartupOptionsTest(unittest.TestCase): + """ + Ensures the extra startup options configured on a Cluster reach the server. + + The options a client sent in its STARTUP message are exposed by the + ``system_views.clients`` virtual table since Cassandra 4.1. + """ + + def connect_and_get_client_options(self, **cluster_kwargs): + """ + Connects a cluster tagged with a unique APPLICATION_NAME and returns the + options the server recorded for it. + """ + application_name = f'app-{uuid.uuid4()}' + extra_startup_options = dict(cluster_kwargs.pop('extra_startup_options', {}), + APPLICATION_NAME=application_name) + + cluster = TestCluster(extra_startup_options=extra_startup_options, **cluster_kwargs) + session = cluster.connect(wait_for_all_pools=True) + self.addCleanup(cluster.shutdown) + + rows = session.execute("SELECT client_options FROM system_views.clients") + options = [row.client_options for row in rows + if row.client_options + and row.client_options.get('APPLICATION_NAME') == application_name] + + self.assertGreater(len(options, 0)) + return application_name, options + + def test_extra_startup_options_are_sent_to_server(self): + application_name, options = self.connect_and_get_client_options( + extra_startup_options={'APPLICATION_VERSION': '1.2.3'}) + + for client_options in options: + self.assertEqual(client_options['APPLICATION_NAME'], application_name) + self.assertEqual(client_options['APPLICATION_VERSION'], '1.2.3') + self.assertEqual(client_options['DRIVER_NAME'], DRIVER_NAME) + self.assertEqual(client_options['DRIVER_VERSION'], DRIVER_VERSION) + self.assertIn('CQL_VERSION', client_options) + + def test_extra_startup_options_do_not_override_driver_options(self): + _, options = self.connect_and_get_client_options( + extra_startup_options={'DRIVER_NAME': 'not the driver', + 'DRIVER_VERSION': '0.0.0', + 'CQL_VERSION': '2.0.0'}) + + for client_options in options: + self.assertEqual(client_options['DRIVER_NAME'], DRIVER_NAME) + self.assertEqual(client_options['DRIVER_VERSION'], DRIVER_VERSION) + self.assertNotEqual(client_options['CQL_VERSION'], '2.0.0') diff --git a/tests/unit/test_cluster.py b/tests/unit/test_cluster.py index 69a65855a0..72a941635c 100644 --- a/tests/unit/test_cluster.py +++ b/tests/unit/test_cluster.py @@ -25,7 +25,8 @@ InvalidRequest, Unauthorized, AuthenticationFailed, OperationTimedOut, UnsupportedOperation, RequestValidationException, ConfigurationException, ProtocolVersion from cassandra.cluster import _Scheduler, Session, Cluster, default_lbp_factory, \ ExecutionProfile, _ConfigMode, EXEC_PROFILE_DEFAULT -from cassandra.connection import SniEndPoint, SniEndPointFactory +from cassandra.connection import SniEndPoint, SniEndPointFactory, DefaultEndPoint +from cassandra.protocol import StartupMessage from cassandra.pool import Host from cassandra.policies import HostDistance, RetryPolicy, RoundRobinPolicy, DowngradingConsistencyRetryPolicy, SimpleConvictionPolicy from cassandra.query import SimpleStatement, named_tuple_factory, tuple_factory @@ -150,6 +151,46 @@ def _mocked_proxy_dns_resolution(self): # single SNI endpoint should be resolved to multiple unique IP addresses self.assertEqual(len(addrs), len(set(addrs))) + def test_extra_startup_options_passed_to_connections(self): + """ + Ensures the extra startup options configured on a cluster reach the + connections it creates. + """ + extra_startup_options = {'EXTRA_OPTION': 'option-1'} + cluster = Cluster(extra_startup_options=extra_startup_options) + + kwargs = cluster._make_connection_kwargs(DefaultEndPoint('127.0.0.1'), {}) + self.assertEqual(kwargs['extra_startup_options'], extra_startup_options) + + def test_extra_startup_options_default_to_empty(self): + kwargs = Cluster()._make_connection_kwargs(DefaultEndPoint('127.0.0.1'), {}) + self.assertEqual(kwargs['extra_startup_options'], {}) + + def test_driver_managed_extra_startup_options_are_removed(self): + """ + Ensures the options managed by the driver are dropped from the extra + startup options, so they cannot be overridden. + """ + extra_startup_options = dict.fromkeys(StartupMessage.KNOWN_OPTION_KEYS, 'overridden') + extra_startup_options['EXTRA_OPTION'] = 'option-1' + + with self.assertLogs('cassandra.cluster', level='WARNING') as logs: + cluster = Cluster(extra_startup_options=extra_startup_options) + + self.assertEqual(cluster.extra_startup_options, {'EXTRA_OPTION': 'option-1'}) + + warnings = [line for line in logs.output if 'Ignoring extra startup option' in line] + self.assertEqual(len(warnings), 1) + for key in StartupMessage.KNOWN_OPTION_KEYS: + self.assertIn(key, warnings[0]) + self.assertNotIn('EXTRA_OPTION', warnings[0]) + + def test_extra_startup_options_are_not_warned_about(self): + with self.assertNoLogs('cassandra.cluster', level='WARNING'): + cluster = Cluster(extra_startup_options={'EXTRA_OPTION': 'option-1'}) + + self.assertEqual(cluster.extra_startup_options, {'EXTRA_OPTION': 'option-1'}) + class SchedulerTest(unittest.TestCase): # TODO: this suite could be expanded; for now just adding a test covering a ticket diff --git a/tests/unit/test_connection.py b/tests/unit/test_connection.py index 3bca654c55..a55b7cd166 100644 --- a/tests/unit/test_connection.py +++ b/tests/unit/test_connection.py @@ -23,18 +23,19 @@ from cassandra.cluster import Cluster from cassandra.connection import (Connection, HEADER_DIRECTION_TO_CLIENT, ProtocolError, locally_supported_compressions, ConnectionHeartbeat, _Frame, Timer, TimerManager, - ConnectionException, DefaultEndPoint) + ConnectionException, DefaultEndPoint, DRIVER_NAME, DRIVER_VERSION) from cassandra.marshal import uint8_pack, uint32_pack, int32_pack from cassandra.protocol import (write_stringmultimap, write_int, write_string, - SupportedMessage, ProtocolHandler) + SupportedMessage, ProtocolHandler, StartupMessage, + read_stringmap) from tests.util import wait_until class ConnectionTest(unittest.TestCase): - def make_connection(self): - c = Connection(DefaultEndPoint('1.2.3.4')) + def make_connection(self, **kwargs): + c = Connection(DefaultEndPoint('1.2.3.4'), **kwargs) c._socket = Mock() c._socket.send.side_effect = lambda x: len(x) return c @@ -245,6 +246,52 @@ def test_disable_compression(self, *args): self.assertEqual(c.decompressor, None) + def send_startup_and_get_options(self, extra_startup_options, **startup_kwargs): + """ + Sends a StartupMessage on a connection configured with the given extra + options, and reads back the option map it would put on the wire. + """ + c = self.make_connection(cql_version='3.4.5', extra_startup_options=extra_startup_options) + c.send_msg = Mock() + + c._send_startup_message(**startup_kwargs) + + self.assertEqual(c.send_msg.call_count, 1) + message = c.send_msg.call_args[0][0] + self.assertIsInstance(message, StartupMessage) + + buf = BytesIO() + message.send_body(buf, c.protocol_version) + buf.seek(0) + return read_stringmap(buf) + + def test_extra_startup_options_are_sent(self): + """ + Ensures the extra options configured on a connection are sent in the + STARTUP message, alongside the options managed by the driver. + """ + options = self.send_startup_and_get_options({'EXTRA_OPTION_1': 'option-1', 'EXTRA_OPTION_2': 'option-2'}, + compression='lz4', no_compact=True) + + self.assertEqual(options, { + 'CQL_VERSION': '3.4.5', + 'DRIVER_NAME': DRIVER_NAME, + 'DRIVER_VERSION': DRIVER_VERSION, + 'COMPRESSION': 'lz4', + 'NO_COMPACT': 'true', + 'EXTRA_OPTION_1': 'option-1', + 'EXTRA_OPTION_2': 'option-2' + }) + + def test_no_extra_startup_options(self): + """ + Ensures a connection without extra options still sends a valid STARTUP message. + """ + expected = {'CQL_VERSION': '3.4.5', 'DRIVER_NAME': DRIVER_NAME, 'DRIVER_VERSION': DRIVER_VERSION} + for extra_startup_options in (None, {}): + with self.subTest(extra_startup_options=extra_startup_options): + self.assertEqual(self.send_startup_and_get_options(extra_startup_options), expected) + def test_not_implemented(self): """ Ensure the following methods throw NIE's. If not, come back and test them.