Verify: migrate from Aerospike to KiviDB

verify.py compares Aerospike with what the loader wrote: record counts per set, then a random sample of records field by field (with the loader's own conversion), then the TTLs of the sampled records. Save it next to as2kivi.py, which it imports.

#!/usr/bin/env python3
"""Compare an Aerospike namespace with what as2kivi.py wrote to KiviDB.

1. Record counts per set (Aerospike's set statistics against a SCAN of the
   matching key names in KiviDB).
2. A random sample of records, bin by bin, using the same conversion as
   as2kivi.py (keep the two files side by side).
3. TTLs of the sampled records, within a tolerance.
"""
import argparse
import json
import math
import random
import sys

import aerospike
import redis

from as2kivi import NEVER_EXPIRES, Stats, convert_hash, convert_json, key_part


def aerospike_set_counts(asc, ns):
    """Master objects per set, summed over the cluster."""
    counts = {}
    for _node, (err, text) in asc.info_all(f"sets/{ns}").items():
        if err:
            continue
        for entry in text.strip().rstrip(";").split(";"):
            fields = dict(kv.split("=", 1) for kv in entry.split(":") if "=" in kv)
            if "set" in fields:
                counts[fields["set"]] = counts.get(fields["set"], 0) + int(fields["objects"])
    return counts


def as_float(b):
    try:
        return float(b)
    except ValueError:
        return None


def as_json(b):
    try:
        v = json.loads(b)
        return v if isinstance(v, (list, dict)) else None
    except ValueError:
        return None


def same(want, got, rounded):
    """Equal, except that a float may differ in its last digit (counted)."""
    if isinstance(want, float) and isinstance(got, (int, float)) and want != got:
        if math.isclose(want, got, rel_tol=1e-15):
            rounded[0] += 1
            return True
        return False
    if isinstance(want, dict) and isinstance(got, dict):
        return want.keys() == got.keys() and all(same(want[k], got[k], rounded) for k in want)
    if isinstance(want, list) and isinstance(got, list):
        return len(want) == len(got) and all(same(w, g, rounded) for w, g in zip(want, got))
    return want == got


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--aerospike", required=True)
    ap.add_argument("--as-user"); ap.add_argument("--as-password")
    ap.add_argument("--namespace", required=True)
    ap.add_argument("--replication-factor", type=int, default=1,
                    help="the namespace's replication-factor: set counts include replicas")
    ap.add_argument("--set", action="append", dest="sets", help="only these sets, as given to as2kivi.py")
    ap.add_argument("--kividb", required=True)
    ap.add_argument("--ca")
    ap.add_argument("--cluster", action="store_true")
    ap.add_argument("--mode", choices=["hash", "json"], default="hash")
    ap.add_argument("--prefix", default="")
    ap.add_argument("--gen-field")
    ap.add_argument("--no-key", choices=["digest", "skip"], default="digest")
    ap.add_argument("--sample", type=int, default=1000)
    ap.add_argument("--ttl-slack", type=int, default=300,
                    help="seconds of TTL difference allowed (the copy took time)")
    a = ap.parse_args()

    host, port = a.aerospike.rsplit(":", 1)
    cfg = {"hosts": [(host, int(port))]}
    if a.as_user:
        cfg["user"], cfg["password"] = a.as_user, a.as_password
    asc = aerospike.client(cfg).connect()
    kw = {"ssl_ca_certs": a.ca} if a.ca else {}
    kv = (redis.RedisCluster if a.cluster else redis.Redis).from_url(a.kividb, **kw)
    ns, failures = a.namespace, 0

    # 1. counts
    for set_name, n in sorted(aerospike_set_counts(asc, ns).items()):
        if a.sets and set_name not in a.sets:
            continue
        n //= a.replication_factor
        got = sum(1 for _ in kv.scan_iter(match=f"{a.prefix}{ns}:{set_name}:*", count=1000))
        status = "ok" if got == n else "MISMATCH"
        failures += got != n
        print(f"count {set_name}: aerospike {n}, kividb {got} {status}")

    # 2 and 3. a sample from random partitions
    convert = convert_hash if a.mode == "hash" else convert_json
    stats, checked, rounded = Stats(), 0, [0]
    partitions = random.sample(range(4096), 4096)
    for pid in partitions:
        if checked >= a.sample:
            break
        recs = asc.query(ns, None).results({"partition_filter": {"begin": pid, "count": 1}})
        for key, meta, bins in recs:
            if checked >= a.sample:
                break
            set_name, user_key, digest = key[1] or "", key[2], key[3]
            if a.sets and set_name not in a.sets:
                continue
            if user_key is None and a.no_key == "skip":
                continue
            kpart = key_part(user_key) if user_key is not None else b"#" + bytes(digest).hex().encode()
            name = f"{a.prefix}{ns}:{set_name}:".encode() + kpart
            want, _ = convert(bins, stats)
            if a.gen_field:
                want[a.gen_field] = meta["gen"] if a.mode == "json" else str(meta["gen"])
            if a.mode == "hash":
                got = kv.hgetall(name)
                want = {k.encode(): (v if isinstance(v, bytes) else v.encode()) for k, v in want.items()}
            else:
                raw = kv.execute_command("JSON.GET", name, "$")
                got = json.loads(raw)[0] if raw else None
            checked += 1
            if a.mode == "hash" and got.keys() == want.keys():
                # a float written by another loader may be formatted differently (1.0E7)
                for f in got:
                    if got[f] == want[f]:
                        continue
                    if as_float(got[f]) is not None and as_float(got[f]) == as_float(want[f]):
                        got[f] = want[f]
                    elif as_json(got[f]) is not None and as_json(got[f]) == as_json(want[f]):
                        got[f] = want[f]   # the same list or map, keys in another order
            if not same(want, got, rounded):
                failures += 1
                print(f"DIFF {name!r}: aerospike {want!r} kividb {got!r}")
                continue
            ttl = kv.ttl(name)
            if meta["ttl"] == NEVER_EXPIRES:
                ok = ttl == -1
            else:
                ok = ttl > 0 and abs(ttl - meta["ttl"]) <= a.ttl_slack
            if not ok:
                failures += 1
                print(f"TTL {name!r}: aerospike {meta['ttl']} kividb {ttl}")

    if rounded[0]:
        print(f"{rounded[0]} JSON numbers differ from Aerospike in the last digit only")
    print(f"sampled {checked} records, {failures} problems")
    sys.exit(1 if failures else 0)


if __name__ == "__main__":
    main()

Run it with the same --namespace, --mode, --prefix, --set, --gen-field and --no-key you gave the loader, and the namespace's replication factor:

python verify.py --aerospike 10.0.0.5:3000 --namespace app --replication-factor 2 \
  --kividb redis://:"$KIVIDB_PASSWORD"@kividb.internal:6380 --mode json --sample 5000

A clean run:

count events: aerospike 200, kividb 200 ok
count orders: aerospike 5000, kividb 5000 ok
count users: aerospike 5000, kividb 5000 ok
sampled 5000 records, 0 problems

It exits with status 1 when anything differs, and prints each difference:

count orders: aerospike 5000, kividb 4999 MISMATCH
DIFF b'app:orders:5': aerospike {b'total': b'7.5', b'items': b'5', b'user': b'u5'} kividb {}
TTL b'app:users:u3': aerospike 3403 kividb -1

While the application is still writing, counts and values can differ by records written since the copy; run it after the catch-up and prune passes, or with writes stopped. The count pass scans every key in the namespace's prefix, so run it against a replica or off-peak on very large datasets.

On a self-hosted KiviDB, INFO keyspace gives the total keys (and how many have an expiry) per database, a quick first check against the sum of Aerospike's set counts.