diff --git a/pyredis/__main__.py b/pyredis/__main__.py index 34fee5f..ac4adbc 100644 --- a/pyredis/__main__.py +++ b/pyredis/__main__.py @@ -33,6 +33,11 @@ async def amain(args): datastore = DataStore() + if args.restore and not AppendOnlyPersister.restore_from_file( + "ccdb.aof", datastore + ): + return -1 + loop = asyncio.get_running_loop() loop.create_task(acheck_expiry_task(datastore)) @@ -76,6 +81,7 @@ def main(args): ) parser.add_argument("--asyncio", action=argparse.BooleanOptionalAction) parser.add_argument("--trio", action=argparse.BooleanOptionalAction) + parser.add_argument("--restore", action=argparse.BooleanOptionalAction) parser.add_argument( "-v", "--verbose", diff --git a/pyredis/persistence.py b/pyredis/persistence.py index 2d1639e..731add8 100644 --- a/pyredis/persistence.py +++ b/pyredis/persistence.py @@ -1,3 +1,7 @@ +from pyredis.commands import handle_command +from pyredis.protocol import extract_frame_from_buffer + + class AppendOnlyPersister: def __init__(self, filename): self._filename = filename @@ -8,3 +12,23 @@ def log_command(self, command): for item in command: self._file.write(item.resp_encode()) + + @staticmethod + def restore_from_file(filename=None, database=None): + buffer = bytearray() + with open(filename, "rb") as f: + while True: + data = f.read(1024) + if data: + buffer.extend(data) + else: + break + while True: + frame, frame_size = extract_frame_from_buffer(buffer) + + if frame: + buffer = buffer[frame_size:] + handle_command(frame, database, None) + else: + break + return True