Commit 6d459962 authored by Petr Špaček's avatar Petr Špaček

qprep: add proper command-line interface and help texts

parent eee27a71
import lmdb
env_open = {
'map_size': 1024**4,
'max_readers': 64,
......@@ -24,3 +26,18 @@ def key_stream(lenv, db):
def qid2key(qid):
"""Encode query ID to database key"""
return str(qid).encode('ascii')
def db_exists(envdir, dbname):
"""
Determine if named DB exists in environment specified by path.
"""
config = env_open.copy()
config['path'] = envdir
config['readonly'] = True
try:
with lmdb.Environment(**config) as env:
db = env.open_db(key=dbname, **db_open, create=False)
return True
except (lmdb.NotFoundError, lmdb.Error):
return False
#!/usr/bin/env python3
import argparse
import logging
import multiprocessing.pool as pool
import sys
......@@ -9,6 +13,9 @@ import blacklist
import dbhelper
QUERIES_DB_NAME = b'queries'
def read_lines(instream):
"""
Yield (line number, stripped line text). Skip empty lines.
......@@ -84,8 +91,26 @@ def wrk_process_line(args):
def main():
logging.basicConfig()
parser = argparse.ArgumentParser(
description='Convert text list of queries from standard input '
'and store wire format into LMDB "queries" DB. '
'Expected query format is: "<qname> <RR type>", '
'one query per line.')
parser.add_argument('envpath', type=str,
help='path where to create LMDB environment')
args = parser.parse_args()
if dbhelper.db_exists(args.envpath, QUERIES_DB_NAME):
logging.critical(
'LMDB environment "%s" already contains DB %s! '
'Overwritting it would invalidate data in the environment, '
'terminating.',
args.envpath, QUERIES_DB_NAME)
sys.exit(1)
qstream = read_lines(sys.stdin)
with pool.Pool(initializer=wrk_lmdb_init, initargs=(sys.argv[1],)) as workers:
with pool.Pool(initializer=wrk_lmdb_init, initargs=(args.envpath,)) as workers:
for _ in workers.imap_unordered(wrk_process_line, qstream, chunksize=1000):
pass
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment