New metadata command

rdbtools/__init__.py

@@ -1,10 +1,10 @@
 from rdbtools.parser import RdbCallback, RdbParser, DebugCallback
-from rdbtools.callbacks import JSONCallback, DiffCallback, ProtocolCallback, KeyValsOnlyCallback, KeysOnlyCallback
+from rdbtools.callbacks import JSONCallback, DiffCallback, ProtocolCallback, KeyValsOnlyCallback, KeysOnlyCallback, MetadataCallback
 from rdbtools.memprofiler import MemoryCallback, PrintAllKeys, StatsAggregator, PrintJustKeys
 
 __version__ = '0.1.15'
 VERSION = tuple(map(int, __version__.split('.')))
 
 __all__ = [
-    'RdbParser', 'RdbCallback', 'JSONCallback', 'DiffCallback', 'MemoryCallback', 'ProtocolCallback', 'KeyValsOnlyCallback', 'KeysOnlyCallback', 'PrintJustKeys']
+    'RdbParser', 'RdbCallback', 'JSONCallback', 'DiffCallback', 'MemoryCallback', 'ProtocolCallback', 'KeyValsOnlyCallback', 'KeysOnlyCallback', 'MetadataCallback', 'PrintJustKeys']
 

rdbtools/callbacks.py

@@ -474,3 +474,23 @@ def select(self, db_number):
 
     def expireat(self, key, timestamp):
         self.emit(b'EXPIREAT', key, timestamp)
+
+
+class MetadataCallback(RdbCallback):
+    def __init__(self, out, string_escape=None):
+        super(MetadataCallback, self).__init__(string_escape)
+        self._out = out
+        self._metadata = {}
+
+    def aux_field(self, key, value):
+        """Collect metadata fields like redis-ver, redis-bits, ctime, used-mem"""
+        if isinstance(key, bytes):
+            key = key.decode('utf-8', 'ignore')
+        if isinstance(value, bytes):
+            value = value.decode('utf-8', 'ignore')
+        self._metadata[key] = value
+
+    def end_rdb(self):
+        """Output collected metadata as JSON"""
+        json_output = json.dumps(self._metadata, indent=2)
+        self._out.write(json_output.encode('utf-8'))

rdbtools/cli/rdb.py

@@ -4,6 +4,7 @@
 import sys
 from argparse import ArgumentParser
 from rdbtools import RdbParser, JSONCallback, DiffCallback, MemoryCallback, ProtocolCallback, PrintAllKeys, KeysOnlyCallback, KeyValsOnlyCallback
+from rdbtools.callbacks import MetadataCallback
 from rdbtools.encodehelpers import ESCAPE_CHOICES
 from rdbtools.parser import HAS_PYTHON_LZF as PYTHON_LZF_INSTALLED
 
@@ -20,7 +21,7 @@ def main():
 
     parser = ArgumentParser(prog='rdb', usage=usage)
     parser.add_argument("-c", "--command", dest="command", required=True,
-                  help="Command to execute. Valid commands are json, diff, justkeys, justkeyvals, memory and protocol", metavar="CMD")
+                  help="Command to execute. Valid commands are json, diff, justkeys, justkeyvals, memory, protocol and metadata", metavar="CMD")
     parser.add_argument("-f", "--file", dest="output",
                   help="Output file", metavar="FILE")
     parser.add_argument("-n", "--db", dest="dbs", action="append",
@@ -89,7 +90,8 @@ def main():
                 'protocol': lambda f: ProtocolCallback(f, string_escape=options.escape,
                                                        emit_expire=not options.no_expire,
                                                        amend_expire=options.amend_expire
-                                                      )
+                                                      ),
+                'metadata': lambda f: MetadataCallback(f, string_escape=options.escape)
             }[options.command](out_file_obj)
         except:
             raise Exception('Invalid Command %s' % options.command)
@@ -103,7 +105,11 @@ def main():
             eprint("")
 
         parser = RdbParser(callback, filters=filters)
-        parser.parse(options.dump_file[0])
+        if options.command == 'metadata':
+            # Use optimized metadata-only parsing for better performance
+            parser.parse_metadata_only(options.dump_file[0])
+        else:
+            parser.parse(options.dump_file[0])
     finally:
         if options.output and out_file_obj is not None:
             out_file_obj.close()

rdbtools/parser.py

@@ -1055,6 +1055,47 @@ def lzf_decompress(self, compressed, expected_length):
                 raise Exception('lzf_decompress', 'Expected lengths do not match %d != %d for key %s' % (len(out_stream), expected_length, self._key))
             return bytes(out_stream)
 
+    def parse_metadata_only(self, filename):
+        """
+        Parse only metadata (AUX fields) from a redis rdb dump file.
+        This is much faster than full parsing for metadata extraction.
+        """
+        self.parse_metadata_only_fd(open(filename, "rb"))
+
+    def parse_metadata_only_fd(self, fd):
+        with fd as f:
+            self.verify_magic_string(f.read(5))
+            self.verify_version(f.read(4))
+            self._callback.start_rdb()
+
+            # Parse only until we encounter database content
+            while True:
+                data_type = read_unsigned_char(f)
+
+                if data_type == REDIS_RDB_OPCODE_AUX:
+                    aux_key = self.read_string(f)
+                    aux_val = self.read_string(f)
+                    ret = self._callback.aux_field(aux_key, aux_val)
+                    if ret:
+                        break  # Callback signals to stop
+                    continue
+
+                if data_type == REDIS_RDB_OPCODE_MODULE_AUX:
+                    self.read_module(f)
+                    continue
+
+                if data_type == REDIS_RDB_OPCODE_EOF:
+                    # Handle empty RDB files
+                    self._callback.end_rdb()
+                    if self._rdb_version >= 5:
+                        f.read(8)
+                    break
+
+                # Any other opcode means we've reached database content
+                # Stop parsing here for metadata-only extraction
+                self._callback.end_rdb()
+                break
+
 def skip(f, free):
     if free :
         f.read(free)