add HostFilterPolicy and tests

mambocab · mambocab · commit 391889b06791 · 2017-07-11T10:55:44.000-04:00
diff --git a/CHANGELOG.rst b/CHANGELOG.rst
@@ -4,6 +4,7 @@
 Features
 --------
 * Add idle_heartbeat_timeout cluster option to tune how long to wait for heartbeat responses. (PYTHON-762)
+* Add HostFilterPolicy (PYTHON-761)
 
 Bug Fixes
 ---------
@@ -19,6 +20,10 @@ Other
 * Bump Cython dependency version to 0.25.2 (PYTHON-754)
 * Fix DeprecationWarning when using lz4 (PYTHON-769)
 
+Other
+-----
+* Deprecate WhiteListRoundRobinPolicy (PYTHON-759)
+
 3.10.0
 ======
 May 24, 2017
diff --git a/cassandra/policies.py b/cassandra/policies.py
@@ -17,6 +17,7 @@
 from random import randint, shuffle
 from threading import Lock
 import socket
+from warnings import warn
 
 from cassandra import ConsistencyLevel, OperationTimedOut
 
@@ -396,6 +397,10 @@ def on_remove(self, *args, **kwargs):
 
 class WhiteListRoundRobinPolicy(RoundRobinPolicy):
     """
+    |wlrrp| **is deprecated. It will be removed in 4.0.** It can effectively be
+    reimplemented using :class:`.HostFilterPolicy`. For more information, see
+    PYTHON-758_.
+
     A subclass of :class:`.RoundRobinPolicy` which evenly
     distributes queries across all nodes in the cluster,
     regardless of what datacenter the nodes may be in, but
@@ -405,12 +410,25 @@ class WhiteListRoundRobinPolicy(RoundRobinPolicy):
     https://datastax-oss.atlassian.net/browse/JAVA-145
     Where connection errors occur when connection
     attempts are made to private IP addresses remotely
+
+    .. |wlrrp| raw:: html
+
+       <b><code>WhiteListRoundRobinPolicy</code></b>
+
+    .. _PYTHON-758: https://datastax-oss.atlassian.net/browse/PYTHON-758
+
     """
     def __init__(self, hosts):
         """
         The `hosts` parameter should be a sequence of hosts to permit
         connections to.
         """
+        msg = ('WhiteListRoundRobinPolicy is deprecated. '
+               'It will be removed in 4.0. '
+               'It can effectively be reimplemented using HostFilterPolicy.')
+        warn(msg, DeprecationWarning)
+        # DeprecationWarnings are silent by default so we also log the message
+        log.warning(msg)
 
         self._allowed_hosts = hosts
         self._allowed_hosts_resolved = [endpoint[4][0] for a in self._allowed_hosts
@@ -441,6 +459,116 @@ def on_add(self, host):
             RoundRobinPolicy.on_add(self, host)
 
 
+class HostFilterPolicy(LoadBalancingPolicy):
+    """
+    A :class:`.LoadBalancingPolicy` subclass configured with a child policy,
+    and a single-argument predicate. This policy defers to the child policy for
+    hosts where ``predicate(host)`` is truthy. Hosts for which
+    ``predicate(host)`` is falsey will be considered :attr:`.IGNORED`, and will
+    not be used in a query plan.
+
+    This can be used in the cases where you need a whitelist or blacklist
+    policy, e.g. to prepare for decommissioning nodes or for testing:
+
+    .. code-block:: python
+
+        def address_is_ignored(host):
+            return host.address in [ignored_address0, ignored_address1]
+
+        blacklist_filter_policy = HostFilterPolicy(
+            child_policy=RoundRobinPolicy(),
+            predicate=address_is_ignored
+        )
+
+        cluster = Cluster(
+            primary_host,
+            load_balancing_policy=blacklist_filter_policy,
+        )
+
+    Please note that whitelist and blacklist policies are not recommended for
+    general, day-to-day use. You probably want something like
+    :class:`.DCAwareRoundRobinPolicy`, which prefers a local DC but has
+    fallbacks, over a brute-force method like whitelisting or blacklisting.
+    """
+
+    def __init__(self, child_policy, predicate):
+        """
+        :param child_policy: an instantiated :class:`.LoadBalancingPolicy`
+                             that this one will defer to.
+        :param predicate: a one-parameter function that takes a :class:`.Host`.
+                          If it returns a falsey value, the :class:`.Host` will
+                          be :attr:`.IGNORED` and not returned in query plans.
+        """
+        super(HostFilterPolicy, self).__init__()
+        self._child_policy = child_policy
+        self._predicate = predicate
+
+    def on_up(self, host, *args, **kwargs):
+        if self.predicate(host):
+            return self._child_policy.on_up(host, *args, **kwargs)
+
+    def on_down(self, host, *args, **kwargs):
+        if self.predicate(host):
+            return self._child_policy.on_down(host, *args, **kwargs)
+
+    def on_add(self, host, *args, **kwargs):
+        if self.predicate(host):
+            return self._child_policy.on_add(host, *args, **kwargs)
+
+    def on_remove(self, host, *args, **kwargs):
+        if self.predicate(host):
+            return self._child_policy.on_remove(host, *args, **kwargs)
+
+    @property
+    def predicate(self):
+        """
+        A predicate, set on object initialization, that takes a :class:`.Host`
+        and returns a value. If the value is falsy, the :class:`.Host` is
+        :class:`~HostDistance.IGNORED`. If the value is truthy,
+        :class:`.HostFilterPolicy` defers to the child policy to determine the
+        host's distance.
+
+        This is a read-only value set in ``__init__``, implemented as a
+        ``property``.
+        """
+        return self._predicate
+
+    def distance(self, host):
+        """
+        Checks if ``predicate(host)``, then returns
+        :attr:`~HostDistance.IGNORED` if falsey, and defers to the child policy
+        otherwise.
+        """
+        if self.predicate(host):
+            return self._child_policy.distance(host)
+        else:
+            return HostDistance.IGNORED
+
+    def populate(self, cluster, hosts):
+        self._child_policy.populate(
+            cluster=cluster,
+            hosts=[h for h in hosts if self.predicate(h)]
+        )
+
+    def make_query_plan(self, working_keyspace=None, query=None):
+        """
+        Defers to the child policy's
+        :meth:`.LoadBalancingPolicy.make_query_plan`. Since host changes (up,
+        down, addition, and removal) have not been propagated to the child
+        policy, the child policy will only ever return policies for which
+        :meth:`.predicate(host)` was truthy when that change occurred.
+        """
+        child_qp = self._child_policy.make_query_plan(
+            working_keyspace=working_keyspace, query=query
+        )
+        for host in child_qp:
+            if self.predicate(host):
+                yield host
+
+    def check_supported(self):
+        return self._child_policy.check_supported()
+
+
 class ConvictionPolicy(object):
     """
     A policy which decides when hosts should be considered down
diff --git a/docs/api/cassandra/policies.rst b/docs/api/cassandra/policies.rst
@@ -24,6 +24,14 @@ Load Balancing
 .. autoclass:: TokenAwarePolicy
    :members:
 
+.. autoclass:: HostFilterPolicy
+
+   # we document these methods manually so we can specify a param to predicate
+
+   .. automethod:: predicate(host)
+   .. automethod:: distance
+   .. automethod:: make_query_plan
+
 Translating Server Node Addresses
 ---------------------------------
 
diff --git a/tests/integration/long/test_loadbalancingpolicies.py b/tests/integration/long/test_loadbalancingpolicies.py
@@ -23,7 +23,8 @@
 from cassandra.concurrent import execute_concurrent_with_args
 from cassandra.metadata import murmur3
 from cassandra.policies import (RoundRobinPolicy, DCAwareRoundRobinPolicy,
-                                TokenAwarePolicy, WhiteListRoundRobinPolicy)
+                                TokenAwarePolicy, WhiteListRoundRobinPolicy,
+                                HostFilterPolicy)
 from cassandra.query import SimpleStatement
 
 from tests.integration import use_singledc, use_multidc, remove_cluster, PROTOCOL_VERSION
@@ -665,3 +666,37 @@ def test_white_list(self):
             pass
         finally:
             cluster.shutdown()
+
+    def test_black_list_with_host_filter_policy(self):
+        use_singledc()
+        keyspace = 'test_black_list_with_hfp'
+        ignored_address = (IP_FORMAT % 2)
+        hfp = HostFilterPolicy(
+            child_policy=RoundRobinPolicy(),
+            predicate=lambda host: host.address != ignored_address
+        )
+        cluster = Cluster(
+            (IP_FORMAT % 1,),
+            load_balancing_policy=hfp,
+            protocol_version=PROTOCOL_VERSION,
+            topology_event_refresh_window=0,
+            status_event_refresh_window=0
+        )
+        self.addCleanup(cluster.shutdown)
+        session = cluster.connect()
+        self._wait_for_nodes_up([1, 2, 3])
+
+        self.assertNotIn(ignored_address, [h.address for h in hfp.make_query_plan()])
+
+        create_schema(cluster, session, keyspace)
+        self._insert(session, keyspace)
+        self._query(session, keyspace)
+
+        self.coordinator_stats.assert_query_count_equals(self, 1, 6)
+        self.coordinator_stats.assert_query_count_equals(self, 2, 0)
+        self.coordinator_stats.assert_query_count_equals(self, 3, 6)
+
+        # policy should not allow reconnecting to ignored host
+        force_stop(2)
+        self._wait_for_nodes_down([2])
+        self.assertFalse(cluster.metadata._hosts[ignored_address].is_currently_reconnecting())
diff --git a/tests/unit/test_policies.py b/tests/unit/test_policies.py
