Merge pull request sripathikrishnan#167 from oranagra/module_aux

handle rdb files with module aux (key is None)

rdbtools/callbacks.py

@@ -124,6 +124,8 @@ def end_stream(self, key, items, last_entry_id, cgroups):
         self._out.write(b'}')
 
     def start_module(self, key, module_name, expiry, info):
+        if key is None:
+            key = "__aux__"
         self._start_key(key, 0)
         self._out.write(self.encode_key(key) + b':{')
         return False
@@ -172,7 +174,8 @@ def start_stream(self, key, listpacks_count, expiry, info):
         self._keyout(key)
 
     def start_module(self, key, module_name, expiry, info):
-        self._keyout(key)
+        if key is not None:
+            self._keyout(key)
         return False
 
 class KeyValsOnlyCallback(RdbCallback):
@@ -258,6 +261,8 @@ def end_stream(self, key, items, last_entry_id, cgroups):
         self._end_key(key)
 
     def start_module(self, key, module_name, expiry, info):
+        if key is None:
+            return False
         self._start_key(key, 0)
         self._out.write(self.encode_key(key) + b' ')
         return False
@@ -343,6 +348,8 @@ def end_stream(self, key, items, last_entry_id, cgroups):
         self.newline()
 
     def start_module(self, key, module_name, expiry, info):
+        if key is None:
+            key = "__aux__"
         self._out.write(self.dbstr() + self.encode_key(key) + b' -> module-name=' + codecs.encode(module_name, 'ascii'))
         self.newline()
         return False

rdbtools/memprofiler.py

@@ -317,7 +317,9 @@ def end_list(self, key, info):
     def start_module(self, key, module_id, expiry, info):
         self._key_expiry = expiry
         self._current_encoding = module_id
-        self._current_size = self.top_level_object_overhead(key, expiry)
+        self._current_size = 0
+        if key is not None:
+            self._current_size += self.top_level_object_overhead(key, expiry)
         self._current_size += 8 + 1  # add the module id length and EOF byte
 
         return False  # don't build the full key buffer

rdbtools/parser.py

@@ -461,9 +461,10 @@ def parse_fd(self, fd):
                         self.read_object(f, data_type)
                     else:
                         self.skip_object(f, data_type)
-                    self._key = None
                 else :
                     self.skip_key_and_object(f, data_type)
+                self._key = None
+
 
     def read_length_with_encoding(self, f):
         length = 0

tests/memprofiler_tests.py

@@ -25,10 +25,15 @@
 
 class Stats(object):
     def __init__(self):
+        self.sums = {}
         self.records = {}
 
     def next_record(self, record):
-        self.records[record.key] = record
+        if record.type not in self.sums:
+            self.sums[record.type] = 0
+        self.sums[record.type] += record.bytes
+        if record.key is not None:
+            self.records[record.key] = record
 
 
 def get_stats(file_name):
@@ -38,6 +43,13 @@ def get_stats(file_name):
     parser.parse(os.path.join(os.path.dirname(__file__), 'dumps', file_name))
     return stats.records
 
+def get_sums(file_name):
+    stats = Stats()
+    callback = MemoryCallback(stats, 64)
+    parser = RdbParser(callback)
+    parser.parse(os.path.join(os.path.dirname(__file__), 'dumps', file_name))
+    return stats.sums
+
 def get_csv(dump_file_name):
     buff = BytesIO()
     callback = MemoryCallback(PrintAllKeys(buff, None, None), 64)
@@ -90,6 +102,10 @@ def test_rdb_with_module(self):
                                        len_largest_element=101, expiry=None)
         self.assertEquals(stats['foo'], expected_record)
 
+    def test_rdb_with_module_aux(self):
+        sums = get_sums('redis_60_with_module_aux.rdb')
+        self.assertEquals(sums['module'], 32)
+
     def test_rdb_with_stream(self):
         stats = get_stats('redis_50_with_streams.rdb')