Create a gist now

Instantly share code, notes, and snippets.

What would you like to do?
#!/usr/bin/env python3
import re
import os
import subprocess
import statistics
class NoQueryPlanOutput(Exception):
pass
def main(dbname, sql_file, times=10):
if not os.path.isfile(sql_file):
raise FileNotFoundError(sql_file)
with open(sql_file) as f:
sql = f.read()
sql = '\n'.join(
x for x in sql.splitlines()
if not x.strip().startswith('--')
)
if not sql.lower().strip().startswith('explain '):
# In case you forgot to put the 'EXPLAIN ANALYZE' before the query
sql = 'EXPLAIN ANALYZE {}'.format(sql)
execution_times = []
planning_times = []
for i in range(times):
p = subprocess.Popen(
['psql', dbname],
stdin=subprocess.PIPE,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
)
out, err = p.communicate(sql.encode('utf-8'))
if err:
raise Exception(err.decode('utf-8'))
output = out.decode('utf-8')
if 'QUERY PLAN' not in output:
raise NoQueryPlanOutput(output)
planning_time, = re.findall('Planning time: ([\d\.]+) ms', output)
execution_time, = re.findall('Execution time: ([\d\.]+) ms', output)
planning_times.append(float(planning_time))
execution_times.append(float(execution_time))
def fmt(ms):
return '{:.2f}ms'.format(ms)
print("Execution time")
print(
"\tBEST ",
fmt(min(execution_times)).ljust(15),
'<-- most important number'
)
print("\tmean ", fmt(statistics.mean(execution_times)))
print("\tmedian ", fmt(statistics.median(execution_times)))
print("Planning time")
print("\tbest ", fmt(min(planning_times)))
print("\tmean ", fmt(statistics.mean(planning_times)))
print("\tmedian ", fmt(statistics.median(planning_times)))
if __name__ == '__main__':
import argparse
parser = argparse.ArgumentParser(description='Best EXPLAIN ANALYZE')
parser.add_argument('dbname', type=str,
help='Database name (e.g. `psql $this`)')
parser.add_argument('sql_file', type=str,
help='The .sql file (e.g. `psql $dbname < $this`)')
parser.add_argument('--times', dest='times', type=int,
default=10,
help='Number of times to run (default 10)')
args = parser.parse_args()
args = vars(args)
main(args.pop('dbname'), args.pop('sql_file'), **args)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment