Merge pull request #81 from sontek/make_no_servers_graceful

hashclient ignores exceptions when no servers are found

2 files changed, 82 insertions, 4 deletions

diff --git a/pymemcache/client/hash.py b/pymemcache/client/hash.py
index 85d8a50..709ad31 100644
--- a/pymemcache/client/hash.py
+++ b/pymemcache/client/hash.py
@@ -130,6 +130,12 @@ class HashClient(object):
                         self._last_dead_check_time = current_time
 
         server = self.hasher.get_node(key)
+        # We've ran out of servers to try
+        if server is None:
+            if self.ignore_exc is True:
+                return
+            raise Exception('All servers seem to be down right now')
+
         client = self.clients[server]
         return client
 
@@ -211,6 +217,10 @@ class HashClient(object):
 
     def _run_cmd(self, cmd, key, default_val, *args, **kwargs):
         client = self._get_client(key)
+
+        if client is None:
+            return False
+
         func = getattr(client, cmd)
         args = list(args)
         args.insert(0, key)
@@ -232,15 +242,20 @@ class HashClient(object):
 
     def set_many(self, values, *args, **kwargs):
         client_batches = {}
+        end = []
+
         for key, value in values.items():
             client = self._get_client(key)
+
+            if client is None:
+                end.append(False)
+                continue
+
             if client.server not in client_batches:
                 client_batches[client.server] = {}
 
             client_batches[client.server][key] = value
 
-        end = []
-
         for server, values in client_batches.items():
             client = self.clients['%s:%s' % server]
             new_args = list(args)
@@ -257,15 +272,20 @@ class HashClient(object):
 
     def get_many(self, keys, *args, **kwargs):
         client_batches = {}
+        end = {}
+
         for key in keys:
             client = self._get_client(key)
+
+            if client is None:
+                end[key] = False
+                continue
+
             if client.server not in client_batches:
                 client_batches[client.server] = []
 
             client_batches[client.server].append(key)
 
-        end = {}
-
         for server, keys in client_batches.items():
             client = self.clients['%s:%s' % server]
             new_args = list(args)
diff --git a/pymemcache/test/test_client_hash.py b/pymemcache/test/test_client_hash.py
index 70b27c0..dc3dde0 100644
--- a/pymemcache/test/test_client_hash.py
+++ b/pymemcache/test/test_client_hash.py
@@ -119,4 +119,62 @@ class TestHashClient(ClientTestMixin, unittest.TestCase):
         client.set(b'key3', b'value2', noreply=False)
         result = client.get_many([b'key1', b'key3'])
         assert result == {}
+
+    def test_no_servers_left(self):
+        from pymemcache.client.hash import HashClient
+        client = HashClient(
+            [], use_pooling=True,
+            ignore_exc=True,
+            timeout=1, connect_timeout=1
+        )
+
+        hashed_client = client._get_client('foo')
+        assert hashed_client is None
+
+    def test_no_servers_left_raise_exception(self):
+        from pymemcache.client.hash import HashClient
+        client = HashClient(
+            [], use_pooling=True,
+            ignore_exc=False,
+            timeout=1, connect_timeout=1
+        )
+
+        with pytest.raises(Exception) as e:
+            client._get_client('foo')
+
+        assert str(e.value) == 'All servers seem to be down right now'
+
+    def test_no_servers_left_with_commands(self):
+        from pymemcache.client.hash import HashClient
+        client = HashClient(
+            [], use_pooling=True,
+            ignore_exc=True,
+            timeout=1, connect_timeout=1
+        )
+
+        result = client.get('foo')
+        assert result is False
+
+    def test_no_servers_left_with_set_many(self):
+        from pymemcache.client.hash import HashClient
+        client = HashClient(
+            [], use_pooling=True,
+            ignore_exc=True,
+            timeout=1, connect_timeout=1
+        )
+
+        result = client.set_many({'foo': 'bar'})
+        assert result is False
+
+    def test_no_servers_left_with_get_many(self):
+        from pymemcache.client.hash import HashClient
+        client = HashClient(
+            [], use_pooling=True,
+            ignore_exc=True,
+            timeout=1, connect_timeout=1
+        )
+
+        result = client.get_many(['foo', 'bar'])
+        assert result == {'foo': False, 'bar': False}
+
     # TODO: Test failover logic