Merge pull request #287 from sylwiaszunejko/white_list_with_unix_sockets

Add support for unix domain sockets to `WhiteListRoundRobinPolicy`

Author: avelanarius

cassandra/policies.py

@@ -19,6 +19,7 @@
 import socket
 import warnings
 from cassandra import WriteType as WT
+from cassandra.connection import UnixSocketEndPoint
 
 
 # This is done this way because WriteType was originally
@@ -436,8 +437,13 @@ def __init__(self, hosts):
         connections to.
         """
         self._allowed_hosts = tuple(hosts)
-        self._allowed_hosts_resolved = [endpoint[4][0] for a in self._allowed_hosts
-                                        for endpoint in socket.getaddrinfo(a, None, socket.AF_UNSPEC, socket.SOCK_STREAM)]
+        self._allowed_hosts_resolved = []
+        for h in self._allowed_hosts:
+            if isinstance(h, UnixSocketEndPoint):
+                self._allowed_hosts_resolved.append(h._unix_socket_path)
+            else:
+                self._allowed_hosts_resolved.extend([endpoint[4][0]
+                                        for endpoint in socket.getaddrinfo(h, None, socket.AF_UNSPEC, socket.SOCK_STREAM)])
 
         RoundRobinPolicy.__init__(self)
 

tests/unit/test_policies.py

@@ -34,7 +34,7 @@
                                 LoadBalancingPolicy, ConvictionPolicy, ReconnectionPolicy, FallthroughRetryPolicy,
                                 IdentityTranslator, EC2MultiRegionTranslator, HostFilterPolicy)
 from cassandra.pool import Host
-from cassandra.connection import DefaultEndPoint
+from cassandra.connection import DefaultEndPoint, UnixSocketEndPoint
 from cassandra.query import Statement
 
 from six.moves import xrange
@@ -1259,6 +1259,17 @@ def test_hosts_with_hostname(self):
         self.assertEqual(sorted(qplan), [host])
 
         self.assertEqual(policy.distance(host), HostDistance.LOCAL)
+    
+    def test_hosts_with_socket_hostname(self):
+        hosts = [UnixSocketEndPoint('/tmp/scylla-workdir/cql.m')]
+        policy = WhiteListRoundRobinPolicy(hosts)
+        host = Host(UnixSocketEndPoint('/tmp/scylla-workdir/cql.m'), SimpleConvictionPolicy)
+        policy.populate(None, [host])
+
+        qplan = list(policy.make_query_plan())
+        self.assertEqual(sorted(qplan), [host])
+
+        self.assertEqual(policy.distance(host), HostDistance.LOCAL)
 
 
 class AddressTranslatorTest(unittest.TestCase):