Add COUNT + WITHLABELS options for MRANGE (#37)

* count + withlabels options

* None -> True

* per review

Author: Ariel Shtul

redistimeseries/client.py

@@ -109,17 +109,22 @@ def __init__(self, *args, **kwargs):
     def appendRetention(params, retention):
         if retention is not None:
             params.extend(['RETENTION', retention])
-            
+
     @staticmethod
     def appendLabels(params, labels):
         if labels:
             params.append('LABELS')
             for k, v in labels.items():
                 params.extend([k,v])
+            
+    @staticmethod
+    def appendCount(params, count):
+        if count is not None:
+            params.extend(['COUNT', count])
 
     @staticmethod
     def appendAggregation(params, aggregation_type, 
-                          bucket_size_msec):     
+                          bucket_size_msec):  
         params.append('AGGREGATION')
         params.extend([aggregation_type, bucket_size_msec])
 
@@ -218,32 +223,39 @@ def deleterule(self, source_key, dest_key):
         """Deletes a compaction rule"""
         return self.execute_command(self.DELETERULE_CMD, source_key, dest_key)
 
-    def range(self, key, from_time, to_time, 
+    def range(self, key, from_time, to_time, count=None,
                 aggregation_type=None, bucket_size_msec=0):
         """
         Query a range from ``key``, from ``from_time`` to ``to_time``.
+        ``count`` limits the number of results.
         Can Aggregate for ``bucket_size_msec`` where an ``aggregation_type``
         can be ['avg', 'sum', 'min', 'max', 'range', 'count', 'first',
         'last', 'std.p', 'std.s', 'var.p', 'var.s']
         """
         params = [key, from_time, to_time]
-        if aggregation_type != None:
+        self.appendCount(params, count)
+        if aggregation_type is not None:
             self.appendAggregation(params, aggregation_type, bucket_size_msec)
 
         return self.execute_command(self.RANGE_CMD, *params)
 
-    def mrange(self, from_time, to_time, filters,
-                     aggregation_type=None, bucket_size_msec=0):
+    def mrange(self, from_time, to_time, filters, count=None,
+                     aggregation_type=None, bucket_size_msec=0, with_labels=False):
         """
         Query a range based on filters,retention_msecs from ``from_time`` to ``to_time``.
+        ``count`` limits the number of results.
         ``filters`` are a list strings such as ['Test=This'].
         Can Aggregate for ``bucket_size_msec`` where an ``aggregation_type``
         can be ['avg', 'sum', 'min', 'max', 'range', 'count', 'first',
         'last', 'std.p', 'std.s', 'var.p', 'var.s']
+        ``WITHLABELS`` appends labels to results.
         """
         params = [from_time, to_time]
-        if aggregation_type != None:
+        self.appendCount(params, count)
+        if aggregation_type is not None:
             self.appendAggregation(params, aggregation_type, bucket_size_msec)
+        if with_labels:
+            params.extend(['WITHLABELS'])
         params.extend(['FILTER'])
         params += filters
         return self.execute_command(self.MRANGE_CMD, *params)

test_commands.py

@@ -37,6 +37,7 @@ def testAlter(self):
         self.assertEqual(10, rts.info(1).retention_msecs)
         pipe = rts.pipeline()
         self.assertTrue(pipe.create(2))
+
     def testAdd(self):
         '''Test TS.ADD calls'''
 
@@ -94,11 +95,13 @@ def testRange(self):
 
         for i in range(100):
             rts.add(1, i, i % 7)
-        self.assertTrue(100, len(rts.range(1, 0, 200)))
+        self.assertEqual(100, len(rts.range(1, 0, 200)))
         for i in range(100):
             rts.add(1, i+200, i % 7)
-        self.assertTrue(200, len(rts.range(1, 0, 500)))
-        self.assertTrue(20, len(rts.range(1, 0, 500, aggregation_type='avg', bucket_size_msec=10)))
+        self.assertEqual(200, len(rts.range(1, 0, 500)))
+        #last sample isn't returned
+        self.assertEqual(19, len(rts.range(1, 0, 500, aggregation_type='avg', bucket_size_msec=10)))
+        self.assertEqual(10, len(rts.range(1, 0, 500, count=10)))
 
     def testMultiRange(self):
         '''Test TS.MRANGE calls which returns range by filter'''
@@ -108,11 +111,25 @@ def testMultiRange(self):
         for i in range(100):
             rts.add(1, i, i % 7)
             rts.add(2, i, i % 11)
-        self.assertTrue(100, len(rts.mrange(0, 200, filters=['Test=This'])))
+                
+        res = rts.mrange(0, 200, filters=['Test=This'])
+        self.assertEqual(2, len(res))
+        self.assertEqual(100, len(res[0]['1'][1]))
+        
+        res = rts.mrange(0, 200, filters=['Test=This'], count=10)
+        self.assertEqual(10, len(res[0]['1'][1]))
+
         for i in range(100):
             rts.add(1, i+200, i % 7)
-        self.assertTrue(20, len(rts.mrange(0, 500, filters=['Test=This'],
-                        aggregation_type='avg', bucket_size_msec=10)))
+        res = rts.mrange(0, 500, filters=['Test=This'],
+                        aggregation_type='avg', bucket_size_msec=10)
+        self.assertEqual(2, len(res))
+        self.assertEqual(19, len(res[0]['1'][1]))
+        
+        #test withlabels
+        self.assertEqual({}, res[0]['1'][0])
+        res = rts.mrange(0, 200, filters=['Test=This'], with_labels=True)
+        self.assertEqual({'Test': 'This'}, res[0]['1'][0])
 
     def testGet(self):
         '''Test TS.GET calls'''
@@ -138,7 +155,7 @@ def testInfo(self):
         '''Test TS.INFO calls'''
         rts.create(1, retention_msecs=5, labels={'currentLabel' : 'currentData'})
         info = rts.info(1)
-        self.assertTrue(info.retention_msecs == 5)
+        self.assertEqual(5, info.retention_msecs)
         self.assertEqual(info.labels['currentLabel'], 'currentData')
 
     def testQueryIndex(self):