diff options
Diffstat (limited to 'scripts')
-rwxr-xr-x | scripts/mkdb.py | 73 |
1 files changed, 73 insertions, 0 deletions
diff --git a/scripts/mkdb.py b/scripts/mkdb.py new file mode 100755 index 0000000..d3786f1 --- /dev/null +++ b/scripts/mkdb.py @@ -0,0 +1,73 @@ +#!/usr/bin/env python3 + +import sys +import sqlite3 +import string +from argparse import ArgumentParser +from pathlib import Path +from random import choice, randrange + +characters = string.ascii_letters + string.digits + + +class Range: + + def __init__(self, arg: str): + start, end = arg.split("-") + + self.start = int(start) + self.end = int(end) + + +def truncate(path: Path) -> None: + with open(path, "w") as file: + file.truncate(0) + + +def main() -> int: + parser = ArgumentParser() + + parser.add_argument("--database", type=Path, required=True) + parser.add_argument("--word-list", type=Path, required=True) + parser.add_argument("--num-records", type=int, required=True) + parser.add_argument("--key-length", type=Range, required=True) + parser.add_argument("--value-length", type=Range, required=True) + + args = parser.parse_args() + + words = args.word_list.read_text().split("\n") + + truncate(args.database) + + def generate_key(length: int) -> str: + return "".join(choice(characters) for _ in range(length)) + + def generate_value(length: int) -> str: + return "-".join(choice(words) for _ in range(length)) + + with sqlite3.connect(args.database) as connection: + connection.execute("PRAGMA jorunal_mode = WAL") + connection.execute("create table kv(key text, value text)") + connection.execute("create index keys on kv(key)") + + for i in range(args.num_records): + key = generate_key(randrange(args.key_length.start, args.key_length.end)) + + val = generate_value( + randrange(args.value_length.start, args.value_length.end) + ) + + print(f"inserting record number {i}", file=sys.stderr) + + connection.execute( + "insert into kv values(?, ?)", + ( + key, + val, + ), + ) + + return 0 + + +sys.exit(main()) |