Merge pull request #71 from ioga/scan

Support for SCAN commands.

README.md

@@ -638,6 +638,7 @@ Thanks to (in no particular order):
 
   - Free connection selection algorithm for pools.
   - Non-unicode charset fixes.
+  - SCAN commands
 
 - Matt Pizzimenti (mjpizz)
 
@@ -649,4 +650,4 @@ Thanks to (in no particular order):
 
 - Evgeny Tataurov (etataurov)
 
-  - Ability to use hiredis protocol parser
+  - Ability to use hiredis protocol parser

tests/test_scan.py

@@ -0,0 +1,76 @@
+#!/usr/bin/env python
+# coding: utf-8
+#
+# Copyright 2014 Ilia Glazkov
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from txredisapi import Connection
+
+from twisted.internet.defer import inlineCallbacks
+from twisted.trial import unittest
+
+from .mixins import RedisVersionCheckMixin, REDIS_HOST, REDIS_PORT
+
+
+class TestScan(unittest.TestCase, RedisVersionCheckMixin):
+    KEYS = ['_scan_test_' + str(v).zfill(4) for v in range(100)]
+    SKEY = ['_scan_test_set']
+    SUFFIX = '12'
+    PATTERN = '_scan_test_*' + SUFFIX
+    FILTERED_KEYS = [k for k in KEYS if k.endswith(SUFFIX)]
+
+    @inlineCallbacks
+    def setUp(self):
+        self.db = yield Connection(REDIS_HOST, REDIS_PORT, reconnect=False)
+        self.redis_2_8_0 = yield self.checkVersion(2, 8, 0)
+        yield self.db.delete(*self.KEYS)
+        yield self.db.delete(self.SKEY)
+
+    @inlineCallbacks
+    def tearDown(self):
+        yield self.db.delete(*self.KEYS)
+        yield self.db.delete(self.SKEY)
+        yield self.db.disconnect()
+
+    @inlineCallbacks
+    def test_scan(self):
+        self._skipCheck()
+        yield self.db.mset(dict((k, 'value') for k in self.KEYS))
+
+        cursor, result = yield self.db.scan(pattern=self.PATTERN)
+
+        while cursor != 0:
+            cursor, keys = yield self.db.scan(cursor, pattern=self.PATTERN)
+            result.extend(keys)
+
+        self.assertEqual(set(result), set(self.FILTERED_KEYS))
+
+    @inlineCallbacks
+    def test_sscan(self):
+        self._skipCheck()
+        yield self.db.sadd(self.SKEY, self.KEYS)
+
+        cursor, result = yield self.db.sscan(self.SKEY, pattern=self.PATTERN)
+
+        while cursor != 0:
+            cursor, keys = yield self.db.sscan(self.SKEY, cursor,
+                                               pattern=self.PATTERN)
+            result.extend(keys)
+
+        self.assertEqual(set(result), set(self.FILTERED_KEYS))
+
+    def _skipCheck(self):
+        if not self.redis_2_8_0:
+            skipMsg = "Redis version < 2.8.0 (found version: %s)"
+            raise unittest.SkipTest(skipMsg % self.redis_version)

txredisapi.py

@@ -554,6 +554,26 @@ def keys(self, pattern="*"):
         """
         return self.execute_command("KEYS", pattern)
 
+    @staticmethod
+    def _build_scan_args(cursor, pattern, count):
+        """
+        Construct arguments list for SCAN, SSCAN, HSCAN, ZSCAN commands
+        """
+        args = [cursor]
+        if pattern is not None:
+            args.extend(("MATCH", pattern))
+        if count is not None:
+            args.extend(("COUNT", count))
+
+        return args
+
+    def scan(self, cursor=0, pattern=None, count=None):
+        """
+        Incrementally iterate the keys in database
+        """
+        args = self._build_scan_args(cursor, pattern, count)
+        return self.execute_command("SCAN", *args)
+
     def randomkey(self):
         """
         Return a random key from the key space
@@ -1013,6 +1033,10 @@ def srandmember(self, key):
         """
         return self.execute_command("SRANDMEMBER", key)
 
+    def sscan(self, key, cursor=0, pattern=None, count=None):
+        args = self._build_scan_args(cursor, pattern, count)
+        return self.execute_command("SSCAN", key, *args)
+
     # Commands operating on sorted zsets (sorted sets)
     def zadd(self, key, score, member, *args):
         """
@@ -1218,6 +1242,10 @@ def _zaggregate(self, command, dstkey, keys, aggregate):
             pieces.extend(("AGGREGATE", aggregate))
         return self.execute_command(*pieces)
 
+    def zscan(self, key, cursor=0, pattern=None, count=None):
+        args = self._build_scan_args(cursor, pattern, count)
+        return self.execute_command("ZSCAN", key, *args)
+
     # Commands operating on hashes
     def hset(self, key, field, value):
         """
@@ -1306,6 +1334,10 @@ def hgetall(self, key):
         f = lambda d: dict(zip(d[::2], d[1::2]))
         return self.execute_command("HGETALL", key, post_proc=f)
 
+    def hscan(self, key, cursor=0, pattern=None, count=None):
+        args = self._build_scan_args(cursor, pattern, count)
+        return self.execute_command("HSCAN", key, *args)
+
     # Sorting
     def sort(self, key, start=None, end=None, by=None, get=None,
              desc=None, alpha=False, store=None):