Added support for UNIX sockets in HashClient.

2 files changed, 68 insertions, 19 deletions

diff --git a/pymemcache/client/hash.py b/pymemcache/client/hash.py
index 536b915..4716549 100644
--- a/pymemcache/client/hash.py
+++ b/pymemcache/client/hash.py
@@ -46,7 +46,8 @@ class HashClient(object):
         Constructor.
 
         Args:
-          servers: list(tuple(hostname, port))
+          servers: list() of tuple(hostname, port) or string containing a UNIX
+                   socket path.
           hasher: optional class three functions ``get_node``, ``add_node``,
                   and ``remove_node``
                   defaults to Rendezvous (HRW) hash.
@@ -101,30 +102,51 @@ class HashClient(object):
                 'lock_generator': lock_generator
             })
 
-        for server, port in servers:
-            self.add_server(server, port)
+        for server in servers:
+            self.add_server(server)
         self.encoding = encoding
         self.tls_context = tls_context
 
-    def add_server(self, server, port):
-        key = '%s:%s' % (server, port)
+    def add_server(self, server, port=None):
+        # To maintain backward compatibility, if a port is provided, assume
+        # that server wasn't provided as a (host, port) tuple.
+        if port is not None:
+            if not isinstance(server, six.string_types):
+                raise TypeError('Server must be a string when passing port.')
+            server = (server, port)
+
+        if isinstance(server, six.string_types):
+            key = server
+        else:
+            key = '%s:%s' % server
 
         if self.use_pooling:
             client = PooledClient(
-                (server, port),
+                server,
                 **self.default_kwargs
             )
         else:
-            client = self.client_class((server, port), **self.default_kwargs)
+            client = self.client_class(server, **self.default_kwargs)
 
         self.clients[key] = client
         self.hasher.add_node(key)
 
-    def remove_server(self, server, port):
+    def remove_server(self, server, port=None):
+        # To maintain backward compatibility, if a port is provided, assume
+        # that server wasn't provided as a (host, port) tuple.
+        if port is not None:
+            if not isinstance(server, six.string_types):
+                raise TypeError('Server must be a string when passing port.')
+            server = (server, port)
+
+        if isinstance(server, six.string_types):
+            key = server
+        else:
+            key = '%s:%s' % server
+
         dead_time = time.time()
-        self._failed_clients.pop((server, port))
-        self._dead_clients[(server, port)] = dead_time
-        key = '%s:%s' % (server, port)
+        self._failed_clients.pop(server)
+        self._dead_clients[server] = dead_time
         self.hasher.remove_node(key)
 
     def _retry_dead(self):
@@ -141,7 +163,7 @@ class HashClient(object):
                     'bringing server back into rotation %s',
                     server
                 )
-                self.add_server(*server)
+                self.add_server(server)
                 del self._dead_clients[server]
             self._last_dead_check_time = current_time
 
@@ -185,7 +207,7 @@ class HashClient(object):
                     # We've reached our max retry attempts, we need to mark
                     # the sever as dead
                     logger.debug('marking server as dead: %s', client.server)
-                    self.remove_server(*client.server)
+                    self.remove_server(client.server)
 
             result = func(*args, **kwargs)
             return result
@@ -239,7 +261,7 @@ class HashClient(object):
                     # We've reached our max retry attempts, we need to mark
                     # the sever as dead
                     logger.debug('marking server as dead: %s', client.server)
-                    self.remove_server(*client.server)
+                    self.remove_server(client.server)
 
             succeeded, failed, err = self._set_many(
                 client, values, *args, **kwargs
@@ -289,7 +311,7 @@ class HashClient(object):
                 'attempts': 0,
             }
             logger.debug("marking server as dead %s", server)
-            self.remove_server(*server)
+            self.remove_server(server)
         # This client has failed previously, we need to update the metadata
         # to reflect that we have attempted it again
         else:
diff --git a/pymemcache/test/test_client_hash.py b/pymemcache/test/test_client_hash.py
index 322ce3d..f972f34 100644
--- a/pymemcache/test/test_client_hash.py
+++ b/pymemcache/test/test_client_hash.py
@@ -5,6 +5,7 @@ from pymemcache import pool
 
 from .test_client import ClientTestMixin, MockSocket
 import unittest
+import os
 import pytest
 import mock
 import socket
@@ -42,7 +43,7 @@ class TestHashClient(ClientTestMixin, unittest.TestCase):
         client_class = 'pymemcache.client.hash.HashClient.client_class'
         with mock.patch(client_class) as internal_client:
             client = HashClient([], timeout=999, key_prefix='foo_bar_baz')
-            client.add_server('127.0.0.1', '11211')
+            client.add_server(('127.0.0.1', '11211'))
 
         assert internal_client.call_args[0][0] == ('127.0.0.1', '11211')
         kwargs = internal_client.call_args[1]
@@ -313,7 +314,7 @@ class TestHashClient(ClientTestMixin, unittest.TestCase):
     @mock.patch("pymemcache.client.hash.HashClient.client_class")
     def test_dead_server_comes_back(self, client_patch):
         client = HashClient([], dead_timeout=0, retry_attempts=0)
-        client.add_server("127.0.0.1", 11211)
+        client.add_server(("127.0.0.1", 11211))
 
         test_client = client_patch.return_value
         test_client.server = ("127.0.0.1", 11211)
@@ -332,7 +333,7 @@ class TestHashClient(ClientTestMixin, unittest.TestCase):
     @mock.patch("pymemcache.client.hash.HashClient.client_class")
     def test_failed_is_retried(self, client_patch):
         client = HashClient([], retry_attempts=1, retry_timeout=0)
-        client.add_server("127.0.0.1", 11211)
+        client.add_server(("127.0.0.1", 11211))
 
         assert client_patch.call_count == 1
 
@@ -354,7 +355,33 @@ class TestHashClient(ClientTestMixin, unittest.TestCase):
 
         client = HashClient([])
         client.client_class = MyClient
-        client.add_server('host', 11211)
+        client.add_server(('host', 11211))
         assert isinstance(client.clients['host:11211'], MyClient)
 
+    def test_mixed_inet_and_unix_sockets(self):
+        servers = [
+            '/tmp/pymemcache.{pid}'.format(pid=os.getpid()),
+            ('127.0.0.1', 11211),
+        ]
+        client = HashClient(servers)
+        assert set(servers) == {c.server for c in client.clients.values()}
+
+    def test_legacy_add_remove_server_signature(self):
+        server = ('127.0.0.1', 11211)
+        client = HashClient([])
+        assert client.clients == {}
+        client.add_server(*server)  # Unpack (host, port) tuple.
+        assert ('%s:%s' % server) in client.clients
+        client._mark_failed_server(server)
+        assert server in client._failed_clients
+        client.remove_server(*server)  # Unpack (host, port) tuple.
+        assert server in client._dead_clients
+        assert server not in client._failed_clients
+
+        # Ensure that server is a string if passing port argument:
+        with pytest.raises(TypeError):
+            client.add_server(server, server[-1])
+        with pytest.raises(TypeError):
+            client.remove_server(server, server[-1])
+
     # TODO: Test failover logic