schema-agreement: add session wait API

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,166 @@ 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', parallelism: int = 10):
+        """
+        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, parallelism)
+            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'], parallelism):
+        hosts = self._get_schema_agreement_hosts(scope)
+        versions = defaultdict(set)
+        errors = {}
+
+        if not hosts:
+            return {'unavailable': 'No local hosts available'}
+
+        cl = ConsistencyLevel.ONE
+        metadata_request_timeout = self.cluster.control_connection._metadata_request_timeout
+        query = QueryMessage(
+            query=maybe_add_timeout_to_query(ControlConnection._SELECT_SCHEMA_LOCAL, metadata_request_timeout),
+            consistency_level=cl)
+
+        parallelism = max(1, min(parallelism, len(hosts)))
+        for offset in range(0, len(hosts), parallelism):
+            remaining = deadline - time.time()
+            if remaining <= 0:
+                for host in hosts[offset:]:
+                    errors[host.endpoint] = "Timed out before querying host"
+                break
+
+            futures = {}
+            for host in hosts[offset:offset + parallelism]:
+                future = self.submit(self._query_local_schema_version, host, query, deadline)
+                if future is None:
+                    errors[host.endpoint] = "Schema agreement executor unavailable"
+                    continue
+                futures[future] = host
+
+            if not futures:
+                continue
+
+            remaining = deadline - time.time()
+            if remaining <= 0:
+                for future, host in futures.items():
+                    future.cancel()
+                    errors[host.endpoint] = "Timed out before querying host"
+                for host in hosts[offset + parallelism:]:
+                    errors[host.endpoint] = "Timed out before querying host"
+                break
+
+            done, not_done = wait_futures(tuple(futures), timeout=remaining)
+            for future in not_done:
+                future.cancel()
+                host = futures[future]
+                errors[host.endpoint] = "Timed out before querying host"
+
+            for future in done:
+                host = futures[future]
+                schema_version, error = future.result()
+                if error is not None:
+                    errors[host.endpoint] = error
+                    continue
+                versions[schema_version].add(host.endpoint)
+
+            if not_done:
+                for host in hosts[offset + parallelism:]:
+                    errors[host.endpoint] = "Timed out before querying host"
+                break
+
+        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"
+
+        pool = self._pools.get(host)
+        if not pool or pool.is_shutdown:
+            return None, "No active connection pool"
+
+        try:
+            connection = pool._get_connection_for_routing_key()
+            query_timeout = self._schema_agreement_query_timeout(remaining)
+            local_result = connection.wait_for_response(query, timeout=query_timeout)
+        except OperationTimedOut as timeout:
+            log.debug("[session] Timed out waiting for schema version from %s: %s", host, timeout)
+            return None, timeout
+        except (ConnectionException, NoConnectionsAvailable, ConnectionBusy) as exc:
+            log.debug("[session] Error querying schema version from %s: %s", host, exc)
+            return None, exc
+
+        rows = dict_factory(local_result.column_names, local_result.parsed_rows)
+        return (rows[0].get("schema_version") if rows 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 +4239,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

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import unittest
+from concurrent.futures import Future
 
 import logging
 import socket
@@ -23,6 +24,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 +249,65 @@ 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
+
     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_response(self, schema_version):
+        response = Mock()
+        response.column_names = ["schema_version"]
+        response.parsed_rows = [[schema_version]]
+        return response
+
+    def _new_schema_agreement_session(self, schema_versions, distances=None):
+        hosts = []
+        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.wait_for_response.return_value = self._mock_schema_response(schema_version)
+            session._pools[host] = self.MockPool(host, connection)
+
+        return session, hosts
+
     # 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 +395,106 @@ 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:
+            connection = session._pools[host].connection
+            connection.wait_for_response.assert_called_once()
+
+    @mock_session_pools
+    def test_wait_for_schema_agreement_retries_until_local_hosts_match(self, *_):
+        session, hosts = self._new_schema_agreement_session(["a", "b"])
+        clock = self.FakeTime()
+        second_connection = session._pools[hosts[1]].connection
+        second_connection.wait_for_response.side_effect = [
+            self._mock_schema_response("b"),
+            self._mock_schema_response("a")]
+
+        with patch('cassandra.cluster.time', new=clock):
+            assert session.wait_for_schema_agreement(wait_time=1)
+        assert second_connection.wait_for_response.call_count == 2
+        assert clock.clock == 0.2
+
+    @mock_session_pools
+    def test_wait_for_schema_agreement_retries_when_local_connection_is_busy(self, *_):
+        session, hosts = self._new_schema_agreement_session(["a", "a"])
+        clock = self.FakeTime()
+        busy_connection = session._pools[hosts[1]].connection
+        busy_connection.wait_for_response.side_effect = [
+            ConnectionBusy("connection overloaded"),
+            self._mock_schema_response("a")]
+
+        with patch('cassandra.cluster.time', new=clock):
+            assert session.wait_for_schema_agreement(wait_time=1)
+        assert busy_connection.wait_for_response.call_count == 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)
+        session._pools[hosts[0]].connection.wait_for_response.assert_called()
+
+    @mock_session_pools
+    @patch('cassandra.cluster.wait_futures')
+    def test_wait_for_schema_agreement_limits_parallel_queries_to_default(self, mocked_wait_futures, *_):
+        session, _ = self._new_schema_agreement_session(["a"] * 11)
+        batch_sizes = []
+
+        def submit(fn, host, query, deadline):
+            future = Future()
+            future.set_result(fn(host, query, deadline))
+            return future
+
+        def fake_wait(futures, timeout=None, return_when=None):
+            batch_sizes.append(len(futures))
+            return set(futures), set()
+
+        session.submit = Mock(side_effect=submit)
+        mocked_wait_futures.side_effect = fake_wait
+
+        assert session.wait_for_schema_agreement(wait_time=1)
+        assert batch_sizes == [10, 1]
+
+    @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')
+
+        session._pools[hosts[0]].connection.wait_for_response.assert_called_once()
+        session._pools[hosts[1]].connection.wait_for_response.assert_not_called()
+        session._pools[hosts[2]].connection.wait_for_response.assert_not_called()
+
+    @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:
+            session._pools[host].connection.wait_for_response.assert_called_once()
+
+    @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):