run-bench

#!/usr/bin/env python3

import shutil
import os
import shlex
import subprocess
import sys
import argparse
import datetime

class Benchmark:
    def __init__(self):
        self.env = dict(os.environ)
        self.now = datetime.datetime.utcnow()
        self.db = None

    def read_conffile(self, fd):
        for line in fd:
            line = line.strip()
            if not line or line.startswith("#"): continue
            if '=' in line:
                key, val = line.split("=", 1)
                val = shlex.split(val)[0]
                if val.startswith("$"):
                    val = self.env[val[1:]]
                self.env[key] = val

    def select_db(self, name):
        self.db = name
        if name in ("pg", "postgresql"):
            self.db = "postgresql"
            self.env["DBA_DB"] = self.env["DBA_DB_POSTGRESQL"]
        elif name == "mysql":
            self.env["DBA_DB"] = self.env["DBA_DB_MYSQL"]
        elif name == "sqlite":
            self.env["DBA_DB"] = self.env["DBA_DB_SQLITE"]
        elif name == "mem":
            self.db = "mem"
            self.env["DBA_DB"] = "mem:"
        elif name == "sqlitev7":
            self.env["DBA_DB"] = self.env["DBA_DB_SQLITE"]
            self.env["DBA_DB_FORMAT"] = "V7"
        elif name in ("pgv7", "postgresqlv7"):
            self.db = "postgresqlv7"
            self.env["DBA_DB"] = self.env["DBA_DB_POSTGRESQL"]
            self.env["DBA_DB_FORMAT"] = "V7"
        elif name == "mysqlv7":
            self.db = "mysqlv7"
            self.env["DBA_DB"] = self.env["DBA_DB_MYSQL"]
            self.env["DBA_DB_FORMAT"] = "V7"
        else:
            raise RuntimeError("unknown database type {}".format(name))

    def add_extra_env(self, assignments):
        for arg in assignments:
            if '=' in arg:
                key, val = arg.split("=", 1)
                self.env[key] = val

    def build(self):
        subprocess.check_call(["make", "-C", "dballe"])
        subprocess.check_call(["make", "-C", "dballe", "bench-run"])

    def run(self):
        self.env["ARGS"] = self.db
        csvout = subprocess.check_output(["../extra/runtest", "bench-run"], env=self.env, cwd="dballe", universal_newlines=True)
        shasum = subprocess.check_output(["git", "rev-parse", "HEAD"], universal_newlines=True).strip()
        ts = self.now.strftime("%Y%m%d%H%M%S")
        fname = os.path.join("bench", "_".join((ts, self.db, shasum)) + ".csv")
        with open(fname, "wt") as fd:
            # Skip the Running... printed at the beginning by runtest
            for line in csvout.splitlines()[1:]:
                print(line, file=fd)

def main():
    parser = argparse.ArgumentParser(description="Run DB-All.e benchmarks.")
    parser.add_argument("env", nargs="*", help="Extra env var assignments")
    parser.add_argument("-d", "--db", default=None, help="Database to use (pg/postgresql, mysql, sqlite, mem, sqlitev7, pgv7/postgresqlv7, mysqlv7)")
    args = parser.parse_args()

    bench = Benchmark()
    for conffile in ("./run-check.conf", ".git/run-check.conf"):
        if not os.path.exists(conffile): continue
        with open(conffile, "rt") as fd:
            bench.read_conffile(fd)

    bench.add_extra_env(args.env)
    bench.build()

    if args.db is None:
        for db in ("pg", "mysql", "sqlite", "mem", "sqlitev7", "postgresqlv7", "mysqlv7"):
            print("Running benchmarks for {}...".format(db))
            bench.select_db(db)
            bench.run()
    else:
        bench.select_db(args.db)
        bench.run()


if __name__ == "__main__":
    main()