cluster: add Session.wait_for_schema_agreement

Add Session.wait_for_schema_agreement() as a session-scoped schema agreement check. The new API queries schema_version from system.local on the connected hosts selected by the requested rack, dc, or cluster scope, respects Cluster.max_schema_agreement_wait and the control-connection metadata timeouts, and bounds the fan-out with configurable parallelism.

Update the public Session docs and switch the integration callers that were explicitly waiting on schema agreement to use the session API. Add unit coverage for agreement, retries, busy connections, missing pools, batching, scope filtering, and invalid scope handling.

Author: dkropachev

cassandra/cluster.py

@@ -29,7 +29,7 @@
 from itertools import groupby, count, chain
 import json
 import logging
-from typing import Any, Dict, Optional, Union
+from typing import Any, Dict, Optional, Union, Literal
 from warnings import warn
 from random import random
 import re
@@ -3374,6 +3374,125 @@ def pool_finished_setting_keyspace(pool, host_errors):
         for pool in tuple(self._pools.values()):
             pool._set_keyspace_for_all_conns(keyspace, pool_finished_setting_keyspace)
 
+    def wait_for_schema_agreement(self, wait_time=None, scope: Literal['rack', 'dc', 'cluster']='dc'):
+        """
+        Wait for connected hosts in the selected scope to report the same
+        schema version from ``system.local``.
+
+        By default, the timeout for this operation is governed by
+        :attr:`~.Cluster.max_schema_agreement_wait` and
+        :attr:`~.Cluster.control_connection_timeout`.
+
+        Passing ``wait_time`` here overrides
+        :attr:`~.Cluster.max_schema_agreement_wait`. Setting ``wait_time <= 0``
+        will bypass schema agreement waits.
+
+        ``scope`` determines which connected hosts participate in the check.
+        Accepted values are ``'rack'``, ``'dc'``, and ``'cluster'``. The
+        default ``'dc'`` scope queries connected hosts in the local rack and
+        local datacenter. ``'rack'`` narrows the check to connected hosts in
+        the local rack only. ``'cluster'`` queries every host this session has
+        a live connection pool for, across all datacenters.
+
+        :param wait_time: Override for
+            :attr:`~.Cluster.max_schema_agreement_wait`.
+        :param scope: Restricts the check to connected hosts in the local rack,
+            local datacenter, or whole connected cluster.
+        :returns: ``True`` when the selected connected hosts agree on schema,
+            otherwise ``False``.
+        :raises ValueError: If ``scope`` is not one of ``'rack'``, ``'dc'``,
+            or ``'cluster'``.
+        """
+        if scope not in ('rack', 'dc', 'cluster'):
+            raise ValueError("Invalid schema agreement scope: %s" % (scope,))
+
+        total_timeout = wait_time if wait_time is not None else self.cluster.max_schema_agreement_wait
+        if total_timeout <= 0:
+            return True
+
+        deadline = time.time() + total_timeout
+        schema_mismatches = None
+
+        while time.time() < deadline:
+            schema_mismatches = self._get_schema_mismatches_for_scope(deadline, scope)
+            if schema_mismatches is None:
+                return True
+
+            log.debug("[session] Local schemas mismatched, trying again")
+            remaining = deadline - time.time()
+            if remaining > 0:
+                time.sleep(min(0.2, remaining))
+
+        log.warning("Local nodes are reporting a schema disagreement: %s", schema_mismatches)
+        return False
+
+    def _get_schema_mismatches_for_scope(self, deadline, scope: Literal['rack', 'dc', 'cluster']):
+        hosts = self._get_schema_agreement_hosts(scope)
+        versions = defaultdict(set)
+        errors = {}
+
+        if not hosts:
+            return {'unavailable': 'No local hosts available'}
+
+        metadata_request_timeout = self.cluster.control_connection._metadata_request_timeout
+        query = maybe_add_timeout_to_query(ControlConnection._SELECT_SCHEMA_LOCAL, metadata_request_timeout)
+
+        for index, host in enumerate(hosts):
+            schema_version, error = self._query_local_schema_version(host, query, deadline)
+            if error is not None:
+                errors[host.endpoint] = error
+                if error == "Timed out before querying host":
+                    for timed_out_host in hosts[index + 1:]:
+                        errors[timed_out_host.endpoint] = error
+                    break
+                continue
+            versions[schema_version].add(host.endpoint)
+
+        if len(versions) == 1 and None not in versions and not errors:
+            log.debug("[session] Local schemas match")
+            return None
+
+        mismatches = dict((version, list(nodes)) for version, nodes in versions.items())
+        if errors:
+            mismatches['unavailable'] = dict((endpoint, str(error)) for endpoint, error in errors.items())
+        return mismatches
+
+    def _get_schema_agreement_hosts(self, scope):
+        allowed_distances = {
+            'rack': (HostDistance.LOCAL_RACK,),
+            'dc': (HostDistance.LOCAL_RACK, HostDistance.LOCAL),
+        }
+        return tuple(
+            host for host, pool in tuple(self._pools.items())
+            if host.is_up is not False
+            and not pool.is_shutdown
+            and (scope == 'cluster' or self._profile_manager.distance(host) in allowed_distances[scope]))
+
+    def _query_local_schema_version(self, host, query, deadline):
+        remaining = deadline - time.time()
+        if remaining <= 0:
+            return None, "Timed out before querying host"
+
+        try:
+            rows = self.execute(
+                query,
+                timeout=self._schema_agreement_query_timeout(remaining),
+                host=host,
+            )
+        except OperationTimedOut as timeout:
+            log.debug("[session] Timed out waiting for schema version from %s: %s", host, timeout)
+            return None, timeout
+        except Exception as exc:
+            log.debug("[session] Error querying schema version from %s: %s", host, exc)
+            return None, exc
+
+        row = rows.one()
+        return (getattr(row, "schema_version", None) if row is not None else None), None
+
+    def _schema_agreement_query_timeout(self, remaining):
+        control_timeout = self.cluster.control_connection._timeout
+        return min(control_timeout, remaining) if control_timeout is not None else remaining
+
     def user_type_registered(self, keyspace, user_type, klass):
         """
         Called by the parent Cluster instance when the user registers a new
@@ -4079,7 +4198,6 @@ def _handle_schema_change(self, event):
         self._cluster.scheduler.schedule_unique(delay, self.refresh_schema, **event)
 
     def wait_for_schema_agreement(self, connection=None, preloaded_results=None, wait_time=None):
-
         total_timeout = wait_time if wait_time is not None else self._cluster.max_schema_agreement_wait
         if total_timeout <= 0:
             return True

docs/api/cassandra/cluster.rst

@@ -169,6 +169,8 @@ Clusters and Sessions
 
    .. automethod:: set_keyspace(keyspace)
 
+   .. automethod:: wait_for_schema_agreement
+
    .. automethod:: get_execution_profile
 
    .. automethod:: execution_profile_clone_update

tests/integration/long/test_schema.py

@@ -158,4 +158,4 @@ def check_and_wait_for_agreement(self, session, rs, exepected):
         time.sleep(1)
         assert rs.response_future.is_schema_agreed == exepected
         if not rs.response_future.is_schema_agreed:
-            session.cluster.control_connection.wait_for_schema_agreement(wait_time=1000)
+            session.wait_for_schema_agreement(wait_time=1000)

tests/integration/standard/test_udts.py

@@ -147,7 +147,7 @@ def test_can_register_udt_before_connecting(self):
         c.register_user_type("udt_test_register_before_connecting2", "user", User2)
 
         s = c.connect(wait_for_all_pools=True)
-        c.control_connection.wait_for_schema_agreement()
+        s.wait_for_schema_agreement()
 
         s.execute("INSERT INTO udt_test_register_before_connecting.mytable (a, b) VALUES (%s, %s)", (0, User1(42, 'bob')))
         result = s.execute("SELECT b FROM udt_test_register_before_connecting.mytable WHERE a=0")

tests/unit/test_cluster.py

@@ -23,6 +23,7 @@
     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 ConnectionBusy
 from cassandra.pool import Host
 from cassandra.policies import HostDistance, RetryPolicy, RoundRobinPolicy, DowngradingConsistencyRetryPolicy, SimpleConvictionPolicy
 from cassandra.query import SimpleStatement, named_tuple_factory, tuple_factory
@@ -247,11 +248,90 @@ def test_event_delay_timing(self, *_):
 
 
 class SessionTest(unittest.TestCase):
+    class FakeTime(object):
+
+        def __init__(self):
+            self.clock = 0
+
+        def time(self):
+            return self.clock
+
+        def sleep(self, amount):
+            self.clock += amount
+
+    class MockPool(object):
+
+        def __init__(self, host, connection):
+            self.host = host
+            self.host_distance = HostDistance.LOCAL
+            self.is_shutdown = False
+            self.connection = connection
+
+        def _get_connection_for_routing_key(self):
+            return self.connection
+
+    class MockResultSet(object):
+
+        def __init__(self, schema_version):
+            self._schema_version = schema_version
+
+        def one(self):
+            return Mock(schema_version=self._schema_version)
+
     def setUp(self):
         if connection_class is None:
             raise unittest.SkipTest('libev does not appear to be installed correctly')
         connection_class.initialize_reactor()
 
+    def _mock_schema_future(self, outcome):
+        future = Mock()
+        if isinstance(outcome, Exception):
+            future.result.side_effect = outcome
+        else:
+            future.result.return_value = self.MockResultSet(outcome)
+        return future
+
+    def _host_query_count(self, session, target_host):
+        return sum(1 for call in session.execute.call_args_list if call.kwargs.get('host') is target_host)
+
+    def _new_schema_agreement_session(self, schema_versions, distances=None):
+        hosts = []
+        connections = {}
+        distance_map = {}
+        if distances is None:
+            distances = [HostDistance.LOCAL] * len(schema_versions)
+
+        for index, schema_version in enumerate(schema_versions):
+            host = Host("127.0.0.%d" % (index + 1), SimpleConvictionPolicy, host_id=uuid.uuid4())
+            host.set_up()
+            hosts.append(host)
+            distance_map[host] = distances[index]
+
+        cluster = Cluster(protocol_version=4)
+        for host in hosts:
+            cluster.metadata.add_or_return_host(host)
+
+        session = Session(cluster, hosts)
+        session._profile_manager.distance = Mock(side_effect=lambda host: distance_map.get(host, HostDistance.LOCAL))
+        session._pools = {}
+        for host, schema_version in zip(hosts, schema_versions):
+            connection = Mock(endpoint=host.endpoint)
+            connection.future_outcomes = [schema_version]
+            session._pools[host] = self.MockPool(host, connection)
+            connections[host] = connection
+
+        def execute(query, parameters=None, timeout=None, trace=False,
+                    custom_payload=None, execution_profile=None,
+                    paging_state=None, host=None, execute_as=None):
+            connection = connections[host]
+            outcome = connection.future_outcomes.pop(0) if len(connection.future_outcomes) > 1 else connection.future_outcomes[0]
+            future = self._mock_schema_future(outcome)
+            return future.result()
+
+        session.execute = Mock(side_effect=execute)
+
+        return session, hosts, connections
+
     # TODO: this suite could be expanded; for now just adding a test covering a PR
     @mock_session_pools
     def test_default_serial_consistency_level_ep(self, *_):
@@ -339,6 +419,89 @@ def test_set_keyspace_escapes_quotes(self, *_):
         assert query == 'USE simple_ks', (
             "Simple keyspace names should not be quoted, got: %r" % query)
 
+    @mock_session_pools
+    def test_wait_for_schema_agreement_queries_all_local_hosts(self, *_):
+        session, hosts, _ = self._new_schema_agreement_session(["a", "a"])
+
+        assert session.wait_for_schema_agreement(wait_time=1)
+
+        for host in hosts:
+            assert self._host_query_count(session, host) == 1
+
+    @mock_session_pools
+    def test_wait_for_schema_agreement_retries_until_local_hosts_match(self, *_):
+        session, hosts, connections = self._new_schema_agreement_session(["a", "b"])
+        clock = self.FakeTime()
+        connections[hosts[1]].future_outcomes = ["b", "a"]
+
+        with patch('cassandra.cluster.time', new=clock):
+            assert session.wait_for_schema_agreement(wait_time=1)
+        for host in hosts:
+            assert self._host_query_count(session, host) == 2
+        assert clock.clock == 0.2
+
+    @mock_session_pools
+    def test_wait_for_schema_agreement_retries_when_local_connection_is_busy(self, *_):
+        session, hosts, connections = self._new_schema_agreement_session(["a", "a"])
+        clock = self.FakeTime()
+        connections[hosts[1]].future_outcomes = [
+            ConnectionBusy("connection overloaded"),
+            "a"]
+
+        with patch('cassandra.cluster.time', new=clock):
+            assert session.wait_for_schema_agreement(wait_time=1)
+        for host in hosts:
+            assert self._host_query_count(session, host) == 2
+        assert clock.clock == 0.2
+
+    @mock_session_pools
+    def test_wait_for_schema_agreement_ignores_local_hosts_without_session_pool(self, *_):
+        session, hosts, _ = self._new_schema_agreement_session(["a"])
+
+        unconnected_host = Host("127.0.0.2", SimpleConvictionPolicy, host_id=uuid.uuid4())
+        unconnected_host.set_up()
+        session.cluster.metadata.add_or_return_host(unconnected_host)
+
+        assert session.wait_for_schema_agreement(wait_time=1)
+        assert self._host_query_count(session, hosts[0]) == 1
+
+    @mock_session_pools
+    def test_wait_for_schema_agreement_queries_hosts_in_order(self, *_):
+        session, hosts, _ = self._new_schema_agreement_session(["a"] * 11)
+
+        assert session.wait_for_schema_agreement(wait_time=1)
+        assert [call.kwargs['host'] for call in session.execute.call_args_list] == list(hosts)
+
+    @mock_session_pools
+    def test_wait_for_schema_agreement_rack_scope_only_queries_local_rack_connections(self, *_):
+        session, hosts, _ = self._new_schema_agreement_session(
+            ["a", "a", "a"],
+            distances=[HostDistance.LOCAL_RACK, HostDistance.LOCAL, HostDistance.REMOTE])
+
+        assert session.wait_for_schema_agreement(wait_time=1, scope='rack')
+
+        assert self._host_query_count(session, hosts[0]) == 1
+        assert self._host_query_count(session, hosts[1]) == 0
+        assert self._host_query_count(session, hosts[2]) == 0
+
+    @mock_session_pools
+    def test_wait_for_schema_agreement_cluster_scope_queries_all_connected_hosts(self, *_):
+        session, hosts, _ = self._new_schema_agreement_session(
+            ["a", "a", "a"],
+            distances=[HostDistance.LOCAL_RACK, HostDistance.LOCAL, HostDistance.REMOTE])
+
+        assert session.wait_for_schema_agreement(wait_time=1, scope='cluster')
+
+        for host in hosts:
+            assert self._host_query_count(session, host) == 1
+
+    @mock_session_pools
+    def test_wait_for_schema_agreement_rejects_unknown_scope(self, *_):
+        session, _, _ = self._new_schema_agreement_session(["a"])
+
+        with pytest.raises(ValueError):
+            session.wait_for_schema_agreement(wait_time=1, scope='planet')
+
 class ProtocolVersionTests(unittest.TestCase):
 
     def test_protocol_downgrade_test(self):