added AddressTranlator interface

aholmberg · aholmberg · commit 940166863af1 · 2016-04-08T13:11:41.000-05:00
PYTHON-69
diff --git a/cassandra/cluster.py b/cassandra/cluster.py
@@ -65,7 +65,7 @@
 from cassandra.metadata import Metadata, protect_name, murmur3
 from cassandra.policies import (TokenAwarePolicy, DCAwareRoundRobinPolicy, SimpleConvictionPolicy,
                                 ExponentialReconnectionPolicy, HostDistance,
-                                RetryPolicy)
+                                RetryPolicy, IdentityTranslator)
 from cassandra.pool import (Host, _ReconnectionHandler, _HostReconnectionHandler,
                             HostConnectionPool, HostConnection,
                             NoConnectionsAvailable)
@@ -347,6 +347,12 @@ def auth_provider(self, value):
     :class:`.policies.SimpleConvictionPolicy`.
     """
 
+    address_translator = IdentityTranslator()
+    """
+    :class:`.policies.AddressTranslator` instance to be used in translating server node addresses
+    to driver connection addresses.
+    """
+
     connect_to_remote_hosts = True
     """
     If left as :const:`True`, hosts that are considered :attr:`~.HostDistance.REMOTE`
@@ -553,7 +559,8 @@ def __init__(self,
                  topology_event_refresh_window=10,
                  connect_timeout=5,
                  schema_metadata_enabled=True,
-                 token_metadata_enabled=True):
+                 token_metadata_enabled=True,
+                 address_translator=None):
         """
         Any of the mutable Cluster attributes may be set as keyword arguments
         to the constructor.
@@ -572,28 +579,30 @@ def __init__(self,
         if load_balancing_policy is not None:
             if isinstance(load_balancing_policy, type):
                 raise TypeError("load_balancing_policy should not be a class, it should be an instance of that class")
-
             self.load_balancing_policy = load_balancing_policy
         else:
             self.load_balancing_policy = default_lbp_factory()
 
         if reconnection_policy is not None:
             if isinstance(reconnection_policy, type):
                 raise TypeError("reconnection_policy should not be a class, it should be an instance of that class")
-
             self.reconnection_policy = reconnection_policy
 
         if default_retry_policy is not None:
             if isinstance(default_retry_policy, type):
                 raise TypeError("default_retry_policy should not be a class, it should be an instance of that class")
-
             self.default_retry_policy = default_retry_policy
 
         if conviction_policy_factory is not None:
             if not callable(conviction_policy_factory):
                 raise ValueError("conviction_policy_factory must be callable")
             self.conviction_policy_factory = conviction_policy_factory
 
+        if address_translator is not None:
+            if isinstance(address_translator, type):
+                raise TypeError("address_translator should not be a class, it should be an instance of that class")
+            self.address_translator = address_translator
+
         if connection_class is not None:
             self.connection_class = connection_class
 
@@ -2348,10 +2357,7 @@ def _refresh_node_list_and_token_map(self, connection, preloaded_results=None,
         should_rebuild_token_map = force_token_rebuild or self._cluster.metadata.partitioner is None
         found_hosts = set()
         for row in peers_result:
-            addr = row.get("rpc_address")
-
-            if not addr or addr in ["0.0.0.0", "::"]:
-                addr = row.get("peer")
+            addr = self._rpc_from_peer_row(row)
 
             tokens = row.get("tokens", None)
             if 'tokens' in row and not tokens:  # it was selected, but empty
@@ -2413,7 +2419,7 @@ def _delay_for_event_type(self, event_type, delay_window):
 
     def _handle_topology_change(self, event):
         change_type = event["change_type"]
-        addr, port = event["address"]
+        addr = self._translate_address(event["address"][0])
         if change_type == "NEW_NODE" or change_type == "MOVED_NODE":
             if self._topology_event_refresh_window >= 0:
                 delay = self._delay_for_event_type('topology_change', self._topology_event_refresh_window)
@@ -2424,7 +2430,7 @@ def _handle_topology_change(self, event):
 
     def _handle_status_change(self, event):
         change_type = event["change_type"]
-        addr, port = event["address"]
+        addr = self._translate_address(event["address"][0])
         host = self._cluster.metadata.get_host(addr)
         if change_type == "UP":
             delay = 1 + self._delay_for_event_type('status_change', 0.5)  # randomness to avoid thundering herd problem on events
@@ -2442,6 +2448,9 @@ def _handle_status_change(self, event):
                 # this will be run by the scheduler
                 self._cluster.on_down(host, is_host_addition=False)
 
+    def _translate_address(self, addr):
+        return self._cluster.address_translator.translate(addr)
+
     def _handle_schema_change(self, event):
         if self._schema_event_refresh_window < 0:
             return
@@ -2523,11 +2532,7 @@ def _get_schema_mismatches(self, peers_result, local_result, local_address):
             schema_ver = row.get('schema_version')
             if not schema_ver:
                 continue
-
-            addr = row.get("rpc_address")
-            if not addr or addr in ["0.0.0.0", "::"]:
-                addr = row.get("peer")
-
+            addr = self._rpc_from_peer_row(row)
             peer = self._cluster.metadata.get_host(addr)
             if peer and peer.is_up:
                 versions[schema_ver].add(addr)
@@ -2538,6 +2543,12 @@ def _get_schema_mismatches(self, peers_result, local_result, local_address):
 
         return dict((version, list(nodes)) for version, nodes in six.iteritems(versions))
 
+    def _rpc_from_peer_row(self, row):
+        addr = row.get("rpc_address")
+        if not addr or addr in ["0.0.0.0", "::"]:
+            addr = row.get("peer")
+        return self._translate_address(addr)
+
     def _signal_error(self):
         with self._lock:
             if self._is_shutdown:
diff --git a/cassandra/policies.py b/cassandra/policies.py
@@ -16,12 +16,9 @@
 import logging
 from random import randint
 from threading import Lock
-import six
 
 from cassandra import ConsistencyLevel
 
-from six.moves import range
-
 log = logging.getLogger(__name__)
 
 
@@ -849,3 +846,31 @@ def on_unavailable(self, query, consistency, required_replicas, alive_replicas,
             return (self.RETHROW, None)
         else:
             return self._pick_consistency(alive_replicas)
+
+
+class AddressTranslator(object):
+    """
+    Interface for translating cluster-defined endpoints.
+
+    The driver discovers nodes using server metadata and topology change events. Normally,
+    the endpoint defined by the server is the right way to connect to a node. In some environments,
+    these addresses may not be reachable, or not preferred (public vs. private IPs in cloud environments,
+    suboptimal routing, etc). This interface allows for translating from server defined endpoints to
+    preferred addresses for driver connections.
+
+    *Note:* :attr:`~Cluster.contact_points` provided while creating the :class:`~.Cluster` instance are not
+    translated using this mechanism -- only addresses received from Cassandra nodes are.
+    """
+    def translate(self, addr):
+        """
+        Accepts the node ip address, and returns a translated address to be used connecting to this node.
+        """
+        raise NotImplementedError
+
+
+class IdentityTranslator(AddressTranslator):
+    """
+    Returns the endpoint with no translation
+    """
+    def translate(self, addr):
+        return addr
