pgbench_parse

#!/usr/bin/env python3

import os
import sys
import re
from csv_helper import csv
import collections
import datetime
import signal


class PGBenchParser(object):

    # progress: 560.0 s, 55376.5 tps, lat 0.866 ms stddev 0.268
    re_progress = re.compile(r"^progress: (\d+.\d+) s, (\d+.\d+) tps, lat (\d+.\d+) ms stddev (\d+.\d+|NaN)")

    # latency average = 2.591 ms
    # latency stddev = 1.300 ms
    re_latency = re.compile(r"^latency (average|stddev) = (\d+.\d+) ms")

    # tps = 18519.035816 (including connections establishing)
    re_tps = re.compile(r"^tps = (\d+.\d+) \((including|excluding) connections establishing\)")

    def __init__(self):
        self.outdir = None
        self.outputs = {}


    def parse_pgbench(self, fh, outdir):
        re_progress = self.re_progress
        re_latency = self.re_latency
        re_tps = self.re_tps

        self.outdir = outdir

        sum_stats = {}

        for line in fh:
            if line.startswith('progress'):
                match = re_progress.match(line)
                data = list(match.groups())
                if (float(data[0]) > 1000000000):
                    ts = datetime.datetime.fromtimestamp(float(match[1]))
                    ts = ts.astimezone()
                    data[0] = ts.isoformat(timespec = 'milliseconds')

                self.output_progress(data)

            elif line.startswith('latency'):
                match = re_latency.match(line)
                if match.group(1) == 'average':
                    sum_stats['lat_avg'] = match.group(2)
                if match.group(1) == 'stddev':
                    sum_stats['lat_stddev'] = match.group(2)

            elif line.startswith('tps'):
                match = re_tps.match(line)
                if match.group(2) == 'including':
                    sum_stats['tps_incl'] = match.group(1)
                if match.group(2) == 'excluding':
                    sum_stats['tps_excl'] = match.group(1)

        self.output_summary(sum_stats)

        self.close_outputs()
        self.outdir = None


    progress_fieldnames = ('ts', 'tps', 'lat', 'stddev')
    progress_filename = 'pgbench_progress.dat'

    def output_progress(self, data):
        name = self.progress_filename
        fieldnames = self.progress_fieldnames

        output = self.outputs.get(name, None)
        if output is None:
            output = self.CSVOutput(self.outdir, name, fieldnames)
            self.outputs[name] = output
            output.writer.writeheader()

        output.writer.writerow(dict(zip(fieldnames, data)))
        output.fh.flush()


    summary_fieldnames = ('tps_incl', 'tps_excl', 'lat_avg', 'lat_stddev')
    summary_filename = 'pgbench_summary.dat'

    def output_summary(self, data):
        name = self.summary_filename
        fieldnames = self.summary_fieldnames

        output = self.outputs.get(name, None)
        if output is None:
            output = self.CSVOutput(self.outdir, name, fieldnames)
            self.outputs[name] = output
            output.writer.writeheader()

        output.writer.writerow(data)
        output.fh.flush()


    def close_outputs(self):
        for key in list(self.outputs.keys()):
            output = self.outputs.pop(key)
            output.close()
        assert(len(self.outputs) == 0)


    class CSVOutput(object):
        __slots__ = ['fh', 'writer']

        def __init__(self, outdir, name, fieldnames):
            fname = os.path.join(outdir,  name)
            print(fname)
            self.fh = open(fname, 'w')
            self.writer = csv.DictWriter(self.fh, dialect='unix_tab', fieldnames=fieldnames)

        def close(self):
            del self.writer
            self.fh.close()
            del self.fh



def terminate_cleanly(signo, frame):
    sys.stderr.write("\n")
    sys.exit(0)

# FIXME: Add proper argument parsing
if __name__ == "__main__":
    signal.signal(signal.SIGINT, terminate_cleanly)
    signal.signal(signal.SIGTERM, terminate_cleanly)

    parser = PGBenchParser()
    for f in sys.argv[1:]:
        if f == '-':
            parser.parse_pgbench(sys.stdin, sys.argv[2])
            break;
        else:
            outdir = os.path.dirname(os.path.abspath(f))
            with open(f, "r") as fh:
                parser.parse_pgbench(fh, outdir)

andres@awork3:~/src/postgresql$ pgbench -n -M prepared -c 16 -j 16 -T 10 -P1 2>&1 |~/src/pgbenchrun/pgbench_parse - /tmp/out/
/tmp/out/pgbench_progress.dat
/tmp/out/pgbench_summary.dat
andres@awork3:~/src/postgresql$ cat /tmp/out/pgbench_summary.dat 
tps_incl	tps_excl	lat_avg	lat_stddev
1888.287624	1890.151450	8.351	11.367
andres@awork3:~/src/postgresql$ cat /tmp/out/pgbench_progress.dat
ts	tps	lat	stddev
1.0	1887.9	8.334	7.503
2.0	1965.0	8.136	8.122
3.0	1926.0	8.328	13.658
4.0	1982.0	8.038	6.448
5.0	1948.9	8.169	11.960
6.0	1840.1	8.774	15.656
7.0	1905.0	8.385	8.372
8.0	1907.0	8.393	8.246
9.0	1877.0	8.513	13.041
10.0	1928.1	8.108	10.451