schema agreement: fix scope handling

Author: dkropachev

cassandra/cluster.py

@@ -20,6 +20,7 @@
 
 import atexit
 import datetime
+from enum import Enum
 from binascii import hexlify
 from collections import defaultdict
 from collections.abc import Mapping
@@ -29,7 +30,7 @@
 from itertools import groupby, count, chain
 import json
 import logging
-from typing import Any, Dict, Optional, Union, Literal
+from typing import Any, Dict, Optional, Union, Tuple
 from warnings import warn
 from random import random
 import re
@@ -214,6 +215,27 @@ def __init__(self, message, errors):
         self.errors = errors
 
 
+class SchemaAgreementScope(str, Enum):
+    """Scope selectors for :meth:`.Session.wait_for_schema_agreement`."""
+
+    RACK = 'rack'
+    DC = 'dc'
+    CLUSTER = 'cluster'
+
+
+_SCHEMA_AGREEMENT_SCOPE_LABELS = {
+    SchemaAgreementScope.RACK: 'local rack',
+    SchemaAgreementScope.DC: 'local datacenter',
+    SchemaAgreementScope.CLUSTER: 'cluster',
+}
+
+_SCHEMA_AGREEMENT_ALLOWED_DISTANCES = {
+    SchemaAgreementScope.RACK: (HostDistance.LOCAL_RACK,),
+    SchemaAgreementScope.DC: (HostDistance.LOCAL_RACK, HostDistance.LOCAL),
+    SchemaAgreementScope.CLUSTER: (HostDistance.LOCAL_RACK, HostDistance.LOCAL, HostDistance.REMOTE),
+}
+
+
 def _future_completed(future):
     """ Helper for run_in_executor() """
     exc = future.exception()
@@ -3374,7 +3396,8 @@ 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'):
+    def wait_for_schema_agreement(self, wait_time: Optional[float] = None,
+                                  scope: SchemaAgreementScope = SchemaAgreementScope.CLUSTER) -> bool:
         """
         Wait for connected hosts in the selected scope to report the same
         schema version from ``system.local``.
@@ -3388,11 +3411,13 @@ def wait_for_schema_agreement(self, wait_time=None, scope: Literal['rack', 'dc',
         must be greater than 0.
 
         ``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.
+        Pass :attr:`SchemaAgreementScope.RACK`, :attr:`SchemaAgreementScope.DC`,
+        or :attr:`SchemaAgreementScope.CLUSTER`. String values ``'rack'``,
+        ``'dc'``, and ``'cluster'`` are accepted for backward compatibility.
+        The default is :attr:`SchemaAgreementScope.CLUSTER`. ``RACK`` narrows
+        the check to connected hosts in the local rack only. ``DC`` checks
+        connected hosts in the local datacenter. ``CLUSTER`` queries every
+        connected host across all datacenters.
 
         :param wait_time: Override for
             :attr:`~.Cluster.max_schema_agreement_wait`.
@@ -3402,10 +3427,12 @@ def wait_for_schema_agreement(self, wait_time=None, scope: Literal['rack', 'dc',
             otherwise ``False``.
         :raises ValueError: If ``wait_time`` is provided and is not greater
             than 0.
-        :raises ValueError: If ``scope`` is not one of ``'rack'``, ``'dc'``,
-            or ``'cluster'``.
+        :raises ValueError: If ``scope`` is not one of the schema agreement
+            scope values.
         """
-        if scope not in ('rack', 'dc', 'cluster'):
+        try:
+            scope = SchemaAgreementScope(scope)
+        except ValueError:
             raise ValueError("Invalid schema agreement scope: %s" % (scope,))
 
         if wait_time is not None and wait_time <= 0:
@@ -3417,43 +3444,50 @@ def wait_for_schema_agreement(self, wait_time=None, scope: Literal['rack', 'dc',
 
         deadline = time.time() + total_timeout
         schema_mismatches = None
+        scope_label = _SCHEMA_AGREEMENT_SCOPE_LABELS[scope]
 
         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")
+            log.debug("[session] Connected hosts in the %s still disagree on schema, trying again", scope_label)
             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)
+        log.warning("[session] Connected hosts in the %s are reporting a schema disagreement: %s",
+                    scope_label, schema_mismatches)
         return False
 
-    def _get_schema_mismatches_for_scope(self, deadline, scope: Literal['rack', 'dc', 'cluster']):
+    def _get_schema_mismatches_for_scope(self, deadline: float,
+                                         scope: SchemaAgreementScope) -> Optional[Dict[Any, Any]]:
         hosts = self._get_schema_agreement_hosts(scope)
-        versions = defaultdict(set)
+        mismatches = defaultdict(list)
         errors = {}
 
         if not hosts:
-            return {'unavailable': 'No local hosts available'}
+            errors[scope.value] = ConnectionException(
+                "No connected hosts available in the %s" % (_SCHEMA_AGREEMENT_SCOPE_LABELS[scope],)
+            )
+            return {'unavailable': errors}
 
         metadata_request_timeout = self.cluster.control_connection._metadata_request_timeout
         query = maybe_add_timeout_to_query(ControlConnection._SELECT_SCHEMA_LOCAL, metadata_request_timeout)
 
         schema_version_futures = []
         for host in hosts:
-            schema_version, error = self._query_local_schema_version(host, query, deadline)
-            if error is not None:
-                errors[host.endpoint] = error
+            try:
+                schema_version_future = self._query_local_schema_version(host, query, deadline)
+            except Exception as exc:
+                errors[host.endpoint] = exc
                 continue
 
-            schema_version_futures.append((host, schema_version))
+            schema_version_futures.append((host, schema_version_future))
 
         if schema_version_futures:
             # Start all host queries first, then wait for the whole batch.
-            remaining = deadline - time.time()
+            remaining = max(0.0, deadline - time.time())
             if remaining > 0:
                 wait_futures([future for _, future in schema_version_futures], timeout=remaining)
 
@@ -3467,35 +3501,27 @@ def _get_schema_mismatches_for_scope(self, deadline, scope: Literal['rack', 'dc'
 
                     row = rows.one()
                     schema_version = getattr(row, "schema_version", None) if row is not None else None
-                    versions[schema_version].add(host.endpoint)
+                    mismatches[schema_version].append(host.endpoint)
                 else:
-                    errors[host.endpoint] = "Timed out before querying host"
+                    errors[host.endpoint] = OperationTimedOut(last_host=host, timeout=max(0.0, deadline - time.time()))
 
-        if len(versions) == 1 and None not in versions and not errors:
-            log.debug("[session] Local schemas match")
+        if len(mismatches) == 1 and None not in mismatches and not errors:
+            log.debug("[session] Connected hosts in the %s agree on schema", _SCHEMA_AGREEMENT_SCOPE_LABELS[scope])
             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
+            mismatches['unavailable'] = errors
+        return dict(mismatches)
 
-    def _get_schema_agreement_hosts(self, scope):
-        allowed_distances = {
-            'rack': (HostDistance.LOCAL_RACK,),
-            'dc': (HostDistance.LOCAL_RACK, HostDistance.LOCAL),
-        }
+    def _get_schema_agreement_hosts(self, scope: SchemaAgreementScope) -> Tuple[Host, ...]:
         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"
+            and self._profile_manager.distance(host) in _SCHEMA_AGREEMENT_ALLOWED_DISTANCES[scope])
 
+    def _query_local_schema_version(self, host: Host, query: str, deadline: float) -> Future:
+        remaining = max(0.0, deadline - time.time())
         try:
             response_future = self.execute_async(
                 query,
@@ -3504,10 +3530,10 @@ def _query_local_schema_version(self, host, query, deadline):
             )
         except OperationTimedOut as timeout:
             log.debug("[session] Timed out waiting for schema version from %s: %s", host, timeout)
-            return None, timeout
+            raise
         except Exception as exc:
             log.debug("[session] Error querying schema version from %s: %s", host, exc)
-            return None, exc
+            raise
 
         schema_version_future = Future()
 
@@ -3528,13 +3554,15 @@ def _set_exception(exc, result_future=schema_version_future):
             response_future.add_callbacks(_set_result, _set_exception)
         except Exception as exc:
             log.debug("[session] Error registering schema version callback from %s: %s", host, exc)
-            return None, exc
+            raise
 
-        return schema_version_future, None
+        return schema_version_future
 
-    def _schema_agreement_query_timeout(self, remaining):
+    def _schema_agreement_query_timeout(self, remaining: float) -> float:
         control_timeout = self.cluster.control_connection._timeout
-        return min(control_timeout, remaining) if control_timeout is not None else remaining
+        if control_timeout is None:
+            return max(0.0, remaining)
+        return max(0.0, min(control_timeout, remaining))
 
     def user_type_registered(self, keyspace, user_type, klass):
         """

tests/unit/test_cluster.py

@@ -22,7 +22,7 @@
 
 from cassandra import ConsistencyLevel, DriverException, Timeout, Unavailable, RequestExecutionException, ReadTimeout, WriteTimeout, CoordinationFailure, ReadFailure, WriteFailure, FunctionFailure, AlreadyExists,\
     InvalidRequest, Unauthorized, AuthenticationFailed, OperationTimedOut, UnsupportedOperation, RequestValidationException, ConfigurationException, ProtocolVersion
-from cassandra.cluster import _Scheduler, Session, Cluster, ResultSet, default_lbp_factory, \
+from cassandra.cluster import _Scheduler, Session, Cluster, ResultSet, SchemaAgreementScope, default_lbp_factory, \
     ExecutionProfile, _ConfigMode, EXEC_PROFILE_DEFAULT
 from cassandra.connection import ConnectionBusy
 from cassandra.pool import Host
@@ -454,8 +454,10 @@ def test_set_keyspace_escapes_quotes(self, *_):
             "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"])
+    def test_wait_for_schema_agreement_default_scope_queries_all_connected_hosts(self, *_):
+        session, hosts, _ = self._new_schema_agreement_session(
+            ["a", "a"],
+            distances=[HostDistance.LOCAL_RACK, HostDistance.REMOTE])
 
         assert session.wait_for_schema_agreement(wait_time=1)
 
@@ -512,22 +514,22 @@ def test_wait_for_schema_agreement_rack_scope_only_queries_local_rack_connection
             ["a", "a", "a"],
             distances=[HostDistance.LOCAL_RACK, HostDistance.LOCAL, HostDistance.REMOTE])
 
-        assert session.wait_for_schema_agreement(wait_time=1, scope='rack')
+        assert session.wait_for_schema_agreement(wait_time=1, scope=SchemaAgreementScope.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, *_):
+    def test_wait_for_schema_agreement_cluster_scope_skips_ignored_hosts(self, *_):
         session, hosts, _ = self._new_schema_agreement_session(
-            ["a", "a", "a"],
-            distances=[HostDistance.LOCAL_RACK, HostDistance.LOCAL, HostDistance.REMOTE])
+            ["a", "a"],
+            distances=[HostDistance.IGNORED, HostDistance.LOCAL])
 
-        assert session.wait_for_schema_agreement(wait_time=1, scope='cluster')
+        assert session.wait_for_schema_agreement(wait_time=1, scope=SchemaAgreementScope.CLUSTER)
 
-        for host in hosts:
-            assert self._host_query_count(session, host) == 1
+        assert self._host_query_count(session, hosts[0]) == 0
+        assert self._host_query_count(session, hosts[1]) == 1
 
     @mock_session_pools
     def test_wait_for_schema_agreement_rejects_unknown_scope(self, *_):

tests/unit/test_session_schema_agreement.py

@@ -157,6 +157,16 @@ def test_wait_for_schema_agreement_uses_cluster_timeout_when_wait_time_is_missin
     assert session.execute_async.call_count == 0
 
 
+def test_wait_for_schema_agreement_returns_false_when_no_hosts_match_scope(monkeypatch):
+    session, _, _ = _new_session(["a"], distances=[HostDistance.IGNORED])
+    clock = FakeTime()
+    monkeypatch.setattr(cluster_module, "time", clock)
+
+    assert session.wait_for_schema_agreement(wait_time=1) is False
+    assert session.execute_async.call_count == 0
+    assert clock.clock == pytest.approx(1.0)
+
+
 def test_wait_for_schema_agreement_uses_host_targeted_session_queries():
     session, hosts, _ = _new_session(["a", "a"])