# /// script
# requires-python = ">=3.10"
# dependencies = ["redis>=5"]
# ///
"""Copy keys, with their expiry, from one Redis to another.

Usage: uv run copy-keys.py SOURCE_URL TARGET_URL [--match PATTERN] [--replace]

URLs look like redis://user:password@host:6379/0 (rediss:// for TLS).
Keys are read with DUMP and written with RESTORE, so every data type
copies as is. Keys that already exist in the target are skipped, unless
you add --replace. --match copies only keys matching a pattern, as in
SCAN (such as 'user:42:*').
"""

import argparse

import redis

parser = argparse.ArgumentParser(description="Copy keys from one Redis to another.")
parser.add_argument("source")
parser.add_argument("target")
parser.add_argument("--match", default=None, help="copy only keys matching this pattern")
parser.add_argument("--replace", action="store_true", help="overwrite keys that exist in the target")
args = parser.parse_args()

src = redis.Redis.from_url(args.source)
dst = redis.Redis.from_url(args.target)
copied = skipped = 0


def copy(keys):
    global copied, skipped
    read = src.pipeline(transaction=False)
    for key in keys:
        read.dump(key)
        read.pttl(key)
    values = read.execute()
    write = dst.pipeline(transaction=False)
    written = 0
    for i in range(len(keys)):
        data, ttl = values[2 * i], values[2 * i + 1]
        if data is None:  # expired or deleted meanwhile
            continue
        write.restore(keys[i], max(ttl, 0), data, replace=args.replace)
        written += 1
    if not written:
        return
    for result in write.execute(raise_on_error=False):
        if isinstance(result, redis.ResponseError) and "BUSYKEY" in str(result):
            skipped += 1
        elif isinstance(result, Exception):
            raise result
        else:
            copied += 1


batch = []
for key in src.scan_iter(match=args.match, count=1000):
    batch.append(key)
    if len(batch) == 500:
        copy(batch)
        batch = []
if batch:
    copy(batch)
print(f"Copied {copied} keys; skipped {skipped} that already existed.")