mkvdb0-2_benchmark.py

#!/usr/bin/env python3
import argparse
import asyncio
import time

from mkvdb import Mkv


# ------------------------------------------------------------------------------
# SYNC BENCHMARKS
# ------------------------------------------------------------------------------

def bench_sync(uri: str, db_name: str, collection_name: str, count: int):
    print(f"\n=== mkvDB Sync Benchmark ({count:,} keys) ===")
    db = Mkv(uri, db_name=db_name, collection_name=collection_name)

    # Start from a clean collection
    try:
        db.purge()
    except Exception:
        # If purge isn't implemented in sync yet, you can drop this
        pass

    # SET
    t0 = time.perf_counter()
    for i in range(count):
        db.set(f"key_{i}", i)
    t1 = time.perf_counter()

    # GET
    for i in range(count):
        db.get(f"key_{i}")
    t2 = time.perf_counter()

    set_rate = count / (t1 - t0)
    get_rate = count / (t2 - t1)

    print(f"SET (sync): {set_rate:,.0f} ops/sec  ({t1 - t0:.2f}s)")
    print(f"GET (sync): {get_rate:,.0f} ops/sec  ({t2 - t1:.2f}s)")

    # Optional: clean up
    try:
        db.purge()
    except Exception:
        pass

    db.close()


# ------------------------------------------------------------------------------
# ASYNC BENCHMARKS
# ------------------------------------------------------------------------------

async def bench_async(uri: str, db_name: str, collection_name: str, count: int):
    print(f"\n=== mkvDB Async Benchmark ({count:,} keys) ===")
    db = Mkv(uri, db_name=db_name, collection_name=collection_name)

    # Start from a clean collection
    try:
        await db.purge()
    except Exception:
        pass

    # SET
    t0 = time.perf_counter()
    for i in range(count):
        await db.set(f"key_{i}", i)
    t1 = time.perf_counter()

    # GET
    for i in range(count):
        await db.get(f"key_{i}")
    t2 = time.perf_counter()

    set_rate = count / (t1 - t0)
    get_rate = count / (t2 - t1)

    print(f"SET (async): {set_rate:,.0f} ops/sec  ({t1 - t0:.2f}s)")
    print(f"GET (async): {get_rate:,.0f} ops/sec  ({t2 - t1:.2f}s)")

    # Optional: clean up
    try:
        await db.purge()
    except Exception:
        pass

    await db.close()


# ------------------------------------------------------------------------------
# CLI ENTRY
# ------------------------------------------------------------------------------

def main():
    parser = argparse.ArgumentParser(description="mkvDB performance benchmark.")
    parser.add_argument(
        "-n", "--count",
        type=int,
        default=100_000,
        help="Number of key-value pairs (default: 100k)",
    )
    parser.add_argument(
        "--async",
        action="store_true",
        dest="async_mode",
        help="Run async version instead of sync",
    )
    parser.add_argument(
        "--uri",
        type=str,
        default="mongodb://localhost:27017",
        help='MongoDB URI (default: "mongodb://localhost:27017")',
    )
    parser.add_argument(
        "--db-name",
        type=str,
        default="mkvdb_benchmark",
        help='Database name (default: "mkvdb_benchmark")',
    )
    parser.add_argument(
        "--collection-name",
        type=str,
        default="kv",
        help='Collection name (default: "kv")',
    )

    args = parser.parse_args()

    if args.async_mode:
        asyncio.run(
            bench_async(args.uri, args.db_name, args.collection_name, args.count)
        )
    else:
        bench_sync(args.uri, args.db_name, args.collection_name, args.count)


if __name__ == "__main__":
    main()