Merge pull request sripathikrishnan#166 from Sekenre/no_ttl_snapshot

adding an options for discarding the TTL or amending it when generating protocol feed from the rdb.

.travis.yml

@@ -1,4 +1,5 @@
 language: python
+dist: trusty
 python:
   - "2.6"
   - "2.7"

README.md

@@ -153,7 +153,7 @@ To limit the size of the files, you can filter on keys using the `--key` option
 
 You can convert RDB file into a stream of [redis protocol](http://redis.io/topics/protocol) using the `protocol` command.
 
-    > rdb --c protocol /var/redis/6379/dump.rdb
+    > rdb -c protocol /var/redis/6379/dump.rdb
 
     *4
     $4
@@ -172,6 +172,12 @@ Read [Redis Mass Insert](http://redis.io/topics/mass-insert) for more informatio
 
 When printing protocol output, the `--escape` option can be used with `printable` or `utf8` to avoid non printable/control characters.
 
+By default, expire times are emitted verbatim if they are present in the rdb file, causing all keys that expire in the past to be removed.
+If this behaviour is unwanted the `-x/--no-expire` option will ignore all key expiry commands.
+
+Otherwise you may want to set an expiry time in the future with `-a/--amend-expire` option which adds an integer number of seconds to the expiry time of each key which is already set to expire.
+This will not change keys that do not already have an expiry set.
+
 # Using the Parser ##
 
     from rdbtools import RdbParser, RdbCallback

rdbtools/callbacks.py

@@ -356,15 +356,20 @@ def _unix_timestamp(dt):
 
 
 class ProtocolCallback(RdbCallback):
-    def __init__(self, out, string_escape=None):
+    def __init__(self, out, string_escape=None, emit_expire=True, amend_expire=0):
         super(ProtocolCallback, self).__init__(string_escape)
+        self._emit_expire = emit_expire
+        self._amend_expire = (amend_expire > 0)
+        self._expire_delta = calendar.datetime.timedelta(seconds=amend_expire)
         self._out = out
         self.reset()
 
     def reset(self):
         self._expires = {}
 
     def set_expiry(self, key, dt):
+        if self._amend_expire:
+            dt = dt + self._expire_delta
         self._expires[key] = dt
 
     def get_expiry_seconds(self, key):
@@ -376,7 +381,7 @@ def expires(self, key):
         return key in self._expires
 
     def pre_expiry(self, key, expiry):
-        if expiry is not None:
+        if expiry is not None and self._emit_expire:
             self.set_expiry(key, expiry)
 
     def post_expiry(self, key):

rdbtools/cli/rdb.py

@@ -2,7 +2,7 @@
 from __future__ import print_function
 import os
 import sys
-from optparse import OptionParser
+from argparse import ArgumentParser
 from rdbtools import RdbParser, JSONCallback, DiffCallback, MemoryCallback, ProtocolCallback, PrintAllKeys, KeysOnlyCallback, KeyValsOnlyCallback
 from rdbtools.encodehelpers import ESCAPE_CHOICES
 from rdbtools.parser import HAS_PYTHON_LZF as PYTHON_LZF_INSTALLED
@@ -14,36 +14,38 @@ def eprint(*args, **kwargs):
 
 VALID_TYPES = ("hash", "set", "string", "list", "sortedset")
 def main():
-    usage = """usage: %prog [options] /path/to/dump.rdb
+    usage = """usage: %(prog)s [options] /path/to/dump.rdb
 
-Example : %prog --command json -k "user.*" /var/redis/6379/dump.rdb"""
+Example : %(prog)s --command json -k "user.*" /var/redis/6379/dump.rdb"""
 
-    parser = OptionParser(usage=usage)
-    parser.add_option("-c", "--command", dest="command",
-                  help="Command to execute. Valid commands are json, diff, justkeys, justkeyvals, memory and protocol", metavar="FILE")
-    parser.add_option("-f", "--file", dest="output",
+    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")
+    parser.add_argument("-f", "--file", dest="output",
                   help="Output file", metavar="FILE")
-    parser.add_option("-n", "--db", dest="dbs", action="append",
+    parser.add_argument("-n", "--db", dest="dbs", action="append",
                   help="Database Number. Multiple databases can be provided. If not specified, all databases will be included.")
-    parser.add_option("-k", "--key", dest="keys", default=None,
+    parser.add_argument("-k", "--key", dest="keys", default=None,
                   help="Keys to export. This can be a regular expression")
-    parser.add_option("-o", "--not-key", dest="not_keys", default=None,
+    parser.add_argument("-o", "--not-key", dest="not_keys", default=None,
                   help="Keys Not to export. This can be a regular expression")
-    parser.add_option("-t", "--type", dest="types", action="append",
+    parser.add_argument("-t", "--type", dest="types", action="append",
                   help="""Data types to include. Possible values are string, hash, set, sortedset, list. Multiple typees can be provided. 
                     If not specified, all data types will be returned""")
-    parser.add_option("-b", "--bytes", dest="bytes", default=None,
+    parser.add_argument("-b", "--bytes", dest="bytes", default=None,
                   help="Limit memory output to keys greater to or equal to this value (in bytes)")
-    parser.add_option("-l", "--largest", dest="largest", default=None,
+    parser.add_argument("-l", "--largest", dest="largest", default=None,
                   help="Limit memory output to only the top N keys (by size)")
-    parser.add_option("-e", "--escape", dest="escape", choices=ESCAPE_CHOICES,
-                      help="Escape strings to encoding: %s (default), %s, %s, or %s." % tuple(ESCAPE_CHOICES))
+    parser.add_argument("-e", "--escape", dest="escape", choices=ESCAPE_CHOICES,
+                  help="Escape strings to encoding: %s (default), %s, %s, or %s." % tuple(ESCAPE_CHOICES))
+    expire_group = parser.add_mutually_exclusive_group(required=False)
+    expire_group.add_argument("-x", "--no-expire", dest="no_expire", default=False, action='store_true',
+                  help="With protocol command, remove expiry from all keys")
+    expire_group.add_argument("-a", "--amend-expire", dest="amend_expire", default=0, type=int, metavar='N',
+                  help="With protocol command, add N seconds to key expiry time")
+    parser.add_argument("dump_file", nargs=1, help="RDB Dump file to process")
 
-    (options, args) = parser.parse_args()
-
-    if len(args) == 0:
-        parser.error("Redis RDB file not specified")
-    dump_file = args[0]
+    options = parser.parse_args()
 
     filters = {}
     if options.dbs:
@@ -84,7 +86,10 @@ def main():
                 'justkeyvals': lambda f: KeyValsOnlyCallback(f, string_escape=options.escape),
                 'memory': lambda f: MemoryCallback(PrintAllKeys(f, options.bytes, options.largest),
                                                    64, string_escape=options.escape),
-                'protocol': lambda f: ProtocolCallback(f, string_escape=options.escape)
+                'protocol': lambda f: ProtocolCallback(f, string_escape=options.escape,
+                                                       emit_expire=not options.no_expire,
+                                                       amend_expire=options.amend_expire
+                                                      )
             }[options.command](out_file_obj)
         except:
             raise Exception('Invalid Command %s' % options.command)
@@ -98,7 +103,7 @@ def main():
             eprint("")
 
         parser = RdbParser(callback, filters=filters)
-        parser.parse(dump_file)
+        parser.parse(options.dump_file[0])
     finally:
         if options.output and out_file_obj is not None:
             out_file_obj.close()

tests/__init__.py

@@ -2,6 +2,7 @@
 from tests.parser_tests import RedisParserTestCase
 from tests.memprofiler_tests import MemoryCallbackTestCase
 from tests.callbacks_tests import ProtocolTestCase, JsonTestCase, DiffTestCase, KeysTestCase, KeyValsTestCase
+from tests.protocol_tests import ProtocolExpireTestCase
 
 
 def all_tests():
@@ -12,7 +13,8 @@ def all_tests():
                       JsonTestCase,
                       DiffTestCase,
                       KeysTestCase,
-                      KeyValsTestCase]
+                      KeyValsTestCase,
+                      ProtocolExpireTestCase]
     for case in test_case_list:
         suite.addTest(unittest.makeSuite(case))
     return suite

tests/protocol_tests.py

@@ -0,0 +1,57 @@
+import unittest
+import os
+import math
+from rdbtools import ProtocolCallback, RdbParser
+from io import BytesIO
+
+class ProtocolExpireTestCase(unittest.TestCase):
+    def setUp(self):
+        self.dumpfile = os.path.join(
+                os.path.dirname(__file__),
+                'dumps',
+                'keys_with_expiry.rdb')
+
+    def tearDown(self):
+        pass
+
+
+    def test_keys_with_expiry(self):
+        expected = (
+                b'*2\r\n$6\r\nSELECT\r\n$1\r\n0\r\n'
+                b'*3\r\n$3\r\nSET\r\n$20\r\nexpires_ms_precision\r\n'
+                b'$27\r\n2022-12-25 10:11:12.573 UTC\r\n'
+                b'*3\r\n$8\r\nEXPIREAT\r\n$20\r\nexpires_ms_precision\r\n'
+                b'$10\r\n1671963072\r\n'
+                )
+        buf = BytesIO()
+        parser = RdbParser(ProtocolCallback(buf))
+        parser.parse(self.dumpfile)
+        self.assertEquals(buf.getvalue(), expected)
+
+
+    def test_amend_expiry(self):
+        expected = (
+                b'*2\r\n$6\r\nSELECT\r\n$1\r\n0\r\n'
+                b'*3\r\n$3\r\nSET\r\n$20\r\nexpires_ms_precision\r\n'
+                b'$27\r\n2022-12-25 10:11:12.573 UTC\r\n'
+                b'*3\r\n$8\r\nEXPIREAT\r\n$20\r\nexpires_ms_precision\r\n'
+                b'$10\r\n1671965072\r\n'
+                )
+        buf = BytesIO()
+        parser = RdbParser(ProtocolCallback(buf, amend_expire=2000))
+        parser.parse(self.dumpfile)
+        self.assertEquals(buf.getvalue(), expected)
+
+
+    def test_skip_expiry(self):
+        expected = (
+                b'*2\r\n$6\r\nSELECT\r\n$1\r\n0\r\n'
+                b'*3\r\n$3\r\nSET\r\n$20\r\nexpires_ms_precision\r\n'
+                b'$27\r\n2022-12-25 10:11:12.573 UTC\r\n'
+                )
+        buf = BytesIO()
+        parser = RdbParser(ProtocolCallback(buf, emit_expire=False))
+        parser.parse(self.dumpfile)
+        self.assertEquals(buf.getvalue(), expected)
+
+