From 45b612073c05e871b4d808c333e34e17b1d3bbaf Mon Sep 17 00:00:00 2001 From: Bohdan Siryk Date: Mon, 7 Sep 2026 10:44:12 +0300 Subject: [PATCH] Allow user to specify arbitrary parameters for STARTUP message Patch enables users to specify arbitrary parametrs for STARTUP messages. This is useful for observability since Cassandra 4.1 STARTUP options are exposed in system_views.clients virtual table. Patch by Bohdan Siryk; reviewed by TBD for CASSPYTHON-29 --- cassandra/cluster.py | 29 ++++++++- cassandra/connection.py | 6 +- cassandra/protocol.py | 10 +++- docs/api/cassandra/cluster.rst | 2 + tests/integration/standard/test_connection.py | 59 ++++++++++++++++++- tests/unit/test_cluster.py | 43 +++++++++++++- tests/unit/test_connection.py | 55 +++++++++++++++-- 7 files changed, 191 insertions(+), 13 deletions(-) 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.