Merge pull request #236 from CybercentreCanada/dispatcher-task-recovery

Dispatcher task reclaiming

Author: cccs-douglass

assemblyline/remote/datatypes/hash.py

@@ -25,13 +25,37 @@
 """
 
 
+class HashIterator:
+    def __init__(self, hash_object):
+        self.hash_object = hash_object
+        self.cursor = 0
+        self.buffer = []
+        self._load_next()
+
+    def __next__(self):
+        while True:
+            if self.buffer:
+                return self.buffer.pop(0)
+            if self.cursor == 0:
+                raise StopIteration()
+            self._load_next()
+
+    def _load_next(self):
+        self.cursor, data = retry_call(self.hash_object.c.hscan, self.hash_object.name, self.cursor)
+        for key, value in data.items():
+            self.buffer.append((key.decode('utf-8'), json.loads(value)))
+
+
 class Hash(object):
     def __init__(self, name, host=None, port=None):
         self.c = get_client(host, port, False)
         self.name = name
         self._pop = self.c.register_script(h_pop_script)
         self._limited_add = self.c.register_script(_limited_add)
 
+    def __iter__(self):
+        return HashIterator(self)
+
     def __enter__(self):
         return self
 

test/test_remote_datatypes.py

@@ -1,4 +1,4 @@
-
+import itertools
 import time
 
 from threading import Thread
@@ -36,6 +36,17 @@ def test_hash(redis_connection):
             assert h.increment("a") == 2
             assert h.increment("a", 10) == 12
             assert h.increment("a", -22) == -10
+            h.delete()
+
+            # Load a bunch of items and test iteration
+            data_before = [''.join(_x) for _x in itertools.product('abcde', repeat=5)]
+            data_before = {_x: _x + _x for _x in data_before}
+            h.multi_set(data_before)
+
+            data_after = {}
+            for key, value in h:
+                data_after[key] = value
+            assert data_before == data_after
 
 
 # noinspection PyShadowingNames