Merge pull request #1098 from theodesp/client-kill-filter

Add client kill with filter

3 files changed, 104 insertions, 1 deletions

diff --git a/redis/client.py b/redis/client.py
index ec87da8..df02f4e 100755
--- a/redis/client.py
+++ b/redis/client.py
@@ -413,6 +413,12 @@ def parse_pubsub_numsub(response, **options):
     return list(zip(response[0::2], response[1::2]))
 
 
+def parse_client_kill(response, **options):
+    if isinstance(response, (long, int)):
+        return int(response)
+    return nativestr(response) == 'OK'
+
+
 class Redis(object):
     """
     Implementation of the Redis protocol.
@@ -471,7 +477,7 @@ class Redis(object):
         {
             'CLIENT GETNAME': lambda r: r and nativestr(r),
             'CLIENT ID': int,
-            'CLIENT KILL': bool_ok,
+            'CLIENT KILL': parse_client_kill,
             'CLIENT LIST': parse_client_list,
             'CLIENT SETNAME': bool_ok,
             'CLIENT UNBLOCK': lambda r: r and int(r) == 1 or False,
@@ -790,6 +796,42 @@ class Redis(object):
         "Disconnects the client at ``address`` (ip:port)"
         return self.execute_command('CLIENT KILL', address)
 
+    def client_kill_filter(self, _id=None, _type=None, addr=None, skipme=None):
+        """
+        Disconnects client(s) using a variety of filter options
+        :param id: Kills a client by its unique ID field
+        :param type: Kills a client by type where type is one of 'normal',
+        'master', 'slave' or 'pubsub'
+        :param addr: Kills a client by its 'address:port'
+        :param skipme: If True, then the client calling the command
+        will not get killed even if it is identified by one of the filter
+        options. If skipme is not provided, the server defaults to skipme=True
+        """
+        args = []
+        if _type is not None:
+            client_types = ('normal', 'master', 'slave', 'pubsub')
+            if str(_type).lower() not in client_types:
+                raise DataError("CLIENT KILL type must be one of %r" % (
+                                client_types,))
+            args.extend((Token.get_token('TYPE'), _type))
+        if skipme is not None:
+            if not isinstance(skipme, bool):
+                raise DataError("CLIENT KILL skipme must be a bool")
+            if skipme:
+                args.extend((Token.get_token('SKIPME'),
+                             Token.get_token('YES')))
+            else:
+                args.extend((Token.get_token('SKIPME'),
+                             Token.get_token('NO')))
+        if _id is not None:
+            args.extend((Token.get_token('ID'), _id))
+        if addr is not None:
+            args.extend((Token.get_token('ADDR'), addr))
+        if not args:
+            raise DataError("CLIENT KILL <filter> <value> ... ... <filter> "
+                            "<value> must specify at least one filter")
+        return self.execute_command('CLIENT KILL', *args)
+
     def client_list(self, _type=None):
         """
         Returns a list of currently connected clients.
diff --git a/tests/conftest.py b/tests/conftest.py
index 5a43968..6e046f2 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -47,6 +47,14 @@ def r(request, **kwargs):
     return _get_client(redis.Redis, request, **kwargs)
 
 
+@pytest.fixture()
+def r2(request, **kwargs):
+    return [
+        _get_client(redis.Redis, request, **kwargs),
+        _get_client(redis.Redis, request, **kwargs),
+    ]
+
+
 def _gen_cluster_mock_resp(r, response):
     mock_connection_pool = Mock()
     connection = Mock()
diff --git a/tests/test_commands.py b/tests/test_commands.py
index 733ec4f..144b94c 100644
--- a/tests/test_commands.py
+++ b/tests/test_commands.py
@@ -96,6 +96,59 @@ class TestRedisCommands(object):
         assert r.client_getname() == 'redis_py_test'
 
     @skip_if_server_version_lt('2.6.9')
+    def test_client_kill(self, r, r2):
+        r.client_setname('redis-py-c1')
+        r2[0].client_setname('redis-py-c2')
+        r2[1].client_setname('redis-py-c3')
+        test_clients = [client for client in r.client_list()
+                        if client.get('name')
+                        in ['redis-py-c1', 'redis-py-c2', 'redis-py-c3']]
+        assert len(test_clients) == 3
+
+        resp = r.client_kill(test_clients[1].get('addr'))
+        assert isinstance(resp, bool) and resp is True
+
+        test_clients = [client for client in r.client_list()
+                        if client.get('name')
+                        in ['redis-py-c1', 'redis-py-c2', 'redis-py-c3']]
+        assert len(test_clients) == 2
+
+    @skip_if_server_version_lt('2.8.12')
+    def test_client_kill_filter_invalid_params(self, r):
+        # empty
+        with pytest.raises(exceptions.DataError):
+            r.client_kill_filter()
+
+        # invalid skipme
+        with pytest.raises(exceptions.DataError):
+            r.client_kill_filter(skipme="yeah")
+
+        # invalid type
+        with pytest.raises(exceptions.DataError):
+            r.client_kill_filter(_type="caster")
+
+    @skip_if_server_version_lt('2.8.12')
+    def test_client_kill_filter(self, r, r2):
+        r.client_setname('redis-py-c1')
+        r2[0].client_setname('redis-py-c2')
+        r2[1].client_setname('redis-py-c3')
+        test_clients = [client for client in r.client_list()
+                        if client.get('name')
+                        in ['redis-py-c1', 'redis-py-c2', 'redis-py-c3']]
+        assert len(test_clients) == 3
+
+        resp = r.client_kill_filter(_id=test_clients[1].get('id'))
+        assert isinstance(resp, int) and resp == 1
+
+        resp = r.client_kill_filter(addr=test_clients[2].get('addr'))
+        assert isinstance(resp, int) and resp == 1
+
+        test_clients = [client for client in r.client_list()
+                        if client.get('name')
+                        in ['redis-py-c1', 'redis-py-c2', 'redis-py-c3']]
+        assert len(test_clients) == 1
+
+    @skip_if_server_version_lt('2.6.9')
     def test_client_list_after_client_setname(self, r):
         r.client_setname('redis_py_test')
         clients = r.client_list()