ec32ae8bad
Reviewers: buda Reviewed By: buda Subscribers: pullbot Differential Revision: https://phabricator.memgraph.io/D539
103 lines
2.4 KiB
Python
103 lines
2.4 KiB
Python
#!/usr/bin/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 time
|
|
import json
|
|
from argparse import ArgumentParser
|
|
from contextlib import redirect_stderr
|
|
import io
|
|
|
|
from neo4j.v1 import GraphDatabase, basic_auth
|
|
|
|
|
|
# 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(d):
|
|
print(json.dumps(_prepare_for_json(d), indent=2))
|
|
|
|
|
|
def main():
|
|
argp = ArgumentParser("Bolt client execution process")
|
|
# positional args
|
|
argp.add_argument("db_uri")
|
|
# named, optional
|
|
argp.add_argument("--encrypt", action="store_true")
|
|
|
|
# parse ags, ensure that stdout is not polluted by argument parsing
|
|
try:
|
|
f = io.StringIO()
|
|
with redirect_stderr(f):
|
|
args = argp.parse_args()
|
|
except:
|
|
_print_dict({RETURN_CODE: 1, ERROR_MSG: "Invalid cmd-line arguments"})
|
|
sys.exit(1)
|
|
|
|
queries = sys.stdin.read().split("\n")
|
|
|
|
driver = GraphDatabase.driver(
|
|
args.db_uri,
|
|
auth=basic_auth("", ""),
|
|
encrypted=args.encrypt)
|
|
|
|
session = driver.session()
|
|
|
|
# execute the queries
|
|
metadatas = []
|
|
start = time.time()
|
|
for query in queries:
|
|
result = session.run(query)
|
|
metadatas.append(result.summary().metadata)
|
|
end = time.time()
|
|
delta_time = end - start
|
|
|
|
_print_dict({
|
|
RETURN_CODE: 0,
|
|
WALL_TIME: (None if not queries else
|
|
delta_time / float(len(queries))),
|
|
"metadatas": metadatas
|
|
})
|
|
|
|
session.close()
|
|
driver.close()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|