From 9d498df60e2c545a5a7af57c125704a9c9511c7f Mon Sep 17 00:00:00 2001 From: Kio Smallwood Date: Sun, 3 May 2020 12:41:17 +0100 Subject: [PATCH] Squashed commit of the following: Add missing tests Fix review query, add expire to existing timestamp Update README with expiry usage Add an option to modify key expiry to N seconds in the future Upgraded to ArgumentParser in order to handle positional file argument and mutually exclusive option group since it makes no sense to modify the expire time and remove it from output. Fix travis build to one with Python 2.6 available Restoring a protocol snapshot without TTL --- .travis.yml | 1 + README.md | 8 +++++- rdbtools/callbacks.py | 9 +++++-- rdbtools/cli/rdb.py | 49 +++++++++++++++++++---------------- tests/__init__.py | 4 ++- tests/protocol_tests.py | 57 +++++++++++++++++++++++++++++++++++++++++ 6 files changed, 102 insertions(+), 26 deletions(-) create mode 100644 tests/protocol_tests.py diff --git a/.travis.yml b/.travis.yml index dc0e4d1..86f79b2 100644 --- a/.travis.yml +++ b/.travis.yml @@ -1,4 +1,5 @@ language: python +dist: trusty python: - "2.6" - "2.7" diff --git a/README.md b/README.md index 5aa1fc8..a390baa 100644 --- a/README.md +++ b/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 diff --git a/rdbtools/callbacks.py b/rdbtools/callbacks.py index 67f2b83..775eae7 100644 --- a/rdbtools/callbacks.py +++ b/rdbtools/callbacks.py @@ -356,8 +356,11 @@ 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() @@ -365,6 +368,8 @@ 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): diff --git a/rdbtools/cli/rdb.py b/rdbtools/cli/rdb.py index f9d5817..1cbcf4b 100755 --- a/rdbtools/cli/rdb.py +++ b/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() diff --git a/tests/__init__.py b/tests/__init__.py index fab64de..95ce871 100644 --- a/tests/__init__.py +++ b/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 diff --git a/tests/protocol_tests.py b/tests/protocol_tests.py new file mode 100644 index 0000000..cb91d18 --- /dev/null +++ b/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) + +