0f562cd043
Reviewers: mislav.bradac Reviewed By: mislav.bradac Subscribers: pullbot Differential Revision: https://phabricator.memgraph.io/D602
105 lines
2.8 KiB
Python
105 lines
2.8 KiB
Python
#!/usr/bin/env python3
|
|
|
|
"""
|
|
A python script that launches the memgraph client,
|
|
executes a query and prints out a JSON dict of measurements
|
|
to stdout.
|
|
|
|
Takes a number of cmd-line arguments of the following structure:
|
|
Positional, mandatory:
|
|
- db_uri
|
|
- query
|
|
Named, optional:
|
|
- encrypt
|
|
|
|
Required the database URI to be passed as the single
|
|
cmd line argument.
|
|
|
|
The dict that is printed out contains:
|
|
- return_code of the client execution process
|
|
- error_msg (empty if not applicable)
|
|
- metedata dict
|
|
|
|
Note that 'metadata' are only valid if the return_code is 0
|
|
"""
|
|
|
|
import sys
|
|
import os
|
|
import time
|
|
import json
|
|
import io
|
|
from contextlib import redirect_stderr
|
|
from multiprocessing import Pool
|
|
from functools import partial
|
|
|
|
# tests/stress dir, that's the place of common.py.
|
|
sys.path.append(os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(
|
|
os.path.realpath(__file__)))), "stress"))
|
|
from common import connection_argument_parser, execute_till_success, \
|
|
argument_driver
|
|
|
|
|
|
# string constants
|
|
RETURN_CODE = "return_code"
|
|
ERROR_MSG = "error_msg"
|
|
WALL_TIME = "wall_time"
|
|
|
|
|
|
def _prepare_for_json(obj):
|
|
if isinstance(obj, dict):
|
|
return {k: _prepare_for_json(v) for k, v in obj.items()}
|
|
if isinstance(obj, list):
|
|
return [_prepare_for_json(elem) for elem in obj]
|
|
if isinstance(obj, (str, int, float, type(None))):
|
|
return obj
|
|
return None
|
|
|
|
|
|
def _print_dict(fname, d):
|
|
with open(fname, "a") as f:
|
|
f.write(json.dumps(_prepare_for_json(d), indent=2))
|
|
|
|
|
|
def _run_query(args, query, self):
|
|
if not hasattr(self, "driver"):
|
|
# TODO: this driver and session is never closed.
|
|
self.driver = argument_driver(args)
|
|
self.session = self.driver.session()
|
|
return execute_till_success(self.session, query)[2]
|
|
_run_query.__defaults__ = (_run_query,)
|
|
|
|
|
|
def main():
|
|
argp = connection_argument_parser()
|
|
argp.add_argument("--num-workers", type=int, default=1)
|
|
argp.add_argument("--output", type=str)
|
|
|
|
# Parse args and ensure that stdout is not polluted by argument parsing.
|
|
try:
|
|
f = io.StringIO()
|
|
with redirect_stderr(f):
|
|
args = argp.parse_args()
|
|
except:
|
|
_print_dict(args.output, {RETURN_CODE: 1, ERROR_MSG: "Invalid cmd-line arguments"})
|
|
sys.exit(1)
|
|
|
|
queries = filter(lambda x: x.strip() != '', sys.stdin.read().split("\n"))
|
|
|
|
# Execute the queries.
|
|
metadatas = []
|
|
with Pool(args.num_workers) as pool:
|
|
start = time.time()
|
|
metadatas = list(pool.map(partial(_run_query, args), queries))
|
|
end = time.time()
|
|
delta_time = end - start
|
|
|
|
_print_dict(args.output, {
|
|
RETURN_CODE: 0,
|
|
WALL_TIME: (None if not queries else delta_time),
|
|
"metadatas": metadatas
|
|
})
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|