testing.py

import subprocess
import tempfile
import time
import signal
import os

dir = r'D:/Dev/postgres/install/bin/'

patched_postgres = r'D:/Dev/postgres/install/bin/'
vanilla_postgres = r'D:/Dev/postgres/install_vanilla/bin/'
pg_data = r'D:/Dev/postgres/pgdata'

scales = [10, 100]
correlations = [1.0, 0.95, 0.90, 0.85, 0.8, 0.75, 0.5, 0.25, 0.10, 0.01]
selectivities = [1.0, 0.66, 0.33, 0.1, 0.01] #[1.0, 0.75, 0.5, 0.25, 0.1, 0.01]
projections = [True]
page_visibilities = [True]
protocols = ['simple']
limit = 50
clients = 1
threads = 1
warm_time = 3
test_time = 5

database = "postgres"


def start_postgres(directory):
    postgres = subprocess.Popen(["postgres", "-D", pg_data], cwd=directory, shell=True, stdout=subprocess.PIPE,
                                stderr=subprocess.PIPE)

    time.sleep(1)
    success = False

    while not success:
        try:
            subprocess.check_output(
                ["psql", "-c", "SELECT 1;", database],
                cwd=dir,
                shell=True, stderr=subprocess.STDOUT)
            success = True
        except subprocess.CalledProcessError:
            success = False

    return postgres


def stop_postgres(postgres):
    try:
        postgres.send_signal(signal.CTRL_C_EVENT)
        postgres.wait(timeout=60)
        return True
    except KeyboardInterrupt:
        time.sleep(5)
        postgres.terminate()
        return False
    os.system("taskkill /f /im  postgres.exe")


def run_pgbench(protocol):
    # warn run
    subprocess.check_output(
        ["pgbench", "-s", str(scale), "-c", str(clients), "-j", str(threads), "-T",
         str(warm_time), "-M", protocol,
         "-n", "-f", temp.name, database], cwd=dir,
        shell=True,
        stderr=subprocess.STDOUT)
    # test run
    bench_output = subprocess.check_output(
        ["pgbench", "-s", str(scale), "-c", str(clients), "-j", str(threads), "-T",
         str(test_time), "-M", protocol,
         "-n", "-f", temp.name, database], cwd=dir,
        shell=True,
        stderr=subprocess.STDOUT)
    tps = int(float((str(bench_output).split("\\n")[9]).split(' ')[2]))
    return tps


try:
    os.system("taskkill /f /im  postgres.exe")
    header = "{}\t{}\t{}\t{}\t{}\t{}\t{}\t{}\t{}".format("scale", "projection", "index_only_fetch",
                                                     "page_visibility", "correlation", "selectivity", "protocol", "vanilla_tps",
                                                     "patched_tps")
    print(header)

    postgres = start_postgres(patched_postgres)

    for scale in scales:
        for correlation in correlations:
            subprocess.check_output(
                ["psql", "-c",
                 "DROP SCHEMA public CASCADE;CREATE SCHEMA public;GRANT ALL ON SCHEMA public TO public;CREATE EXTENSION pg_visibility;",
                 database],
                cwd=dir, shell=True,
                stderr=subprocess.STDOUT)

            subprocess.check_output(["pgbench", "-i", "-s", str(scale), "-q", "--unlogged-tables", database], cwd=dir, shell=True,
                                    stderr=subprocess.STDOUT)

            subprocess.check_output(
                ["psql", "-c", "ALTER TABLE pgbench_accounts drop constraint pgbench_accounts_pkey;", database],
                cwd=dir,
                shell=True, stderr=subprocess.STDOUT)

            subprocess.check_output(
                ["psql", "-c",
                 "WITH X AS (SELECT count(*) AS x FROM pgbench_accounts) UPDATE pgbench_accounts SET bid = TRUNC(RANDOM() * (SELECT x FROM X) + 1);",
                 database],
                cwd=dir, shell=True,
                stderr=subprocess.STDOUT)

            subprocess.check_output(
                ["psql", "-c",
                 "WITH X AS (SELECT count(*) AS x FROM pgbench_accounts) UPDATE pgbench_accounts SET aid = TRUNC(RANDOM() * (SELECT x FROM X) + 1) WHERE random() <= (1.0 - {});".format(
                     correlation), database],
                cwd=dir, shell=True,
                stderr=subprocess.STDOUT)

            subprocess.check_output(
                ["psql", "-c", "CREATE index test_index ON pgbench_accounts USING btree(aid, bid);", database], cwd=dir,
                shell=True, stderr=subprocess.STDOUT)
            subprocess.check_output(["psql", "-c", "VACUUM FULL;", database], cwd=dir, shell=True,
                                    stderr=subprocess.STDOUT)

            for page_visibility in page_visibilities:
                if (page_visibility):
                    subprocess.check_output(["psql", "-c", "VACUUM ANALYZE pgbench_accounts;", database], cwd=dir,
                                            shell=True, stderr=subprocess.STDOUT)
                else:
                    subprocess.check_output(
                        ["psql", "-c", "SELECT pg_truncate_visibility_map('pgbench_accounts'::regclass);", database],
                        cwd=dir,
                        shell=True, stderr=subprocess.STDOUT)

                test_base = "\set r random(1, 100000 * :scale)\r\nSELECT {} FROM pgbench_accounts WHERE aid > :r and bid % 100 <= {} order by aid limit {};"
                plan_base = "EXPLAIN SELECT {} FROM pgbench_accounts WHERE aid > 5000 and bid % 100 <= {} order by aid limit {};"

                for projection in projections:
                    if (projection):
                        proj = "*"
                    else:
                        proj = "aid * 2, bid * 3, abalance * 4, filler"
                    for selectivity in selectivities:
                        test = test_base.format(proj, str(100 * selectivity), limit)
                        plan = plan_base.format(proj, str(100 * selectivity), limit)

                        plan_output = str(
                            subprocess.check_output(["psql", "-c", plan, database], cwd=dir, shell=True,
                                                    stderr=subprocess.STDOUT))
                        if ("Index Only" in plan_output):
                            index_only_fetch = True
                        elif ("Index Scan" in plan_output):
                            index_only_fetch = False
                        else:
                            raise ValueError('bad query plan ' + plan_output)
                        with tempfile.NamedTemporaryFile(delete=False) as temp:
                            temp.write(test.encode())
                            temp.flush
                            temp.close()

                            for protocol in protocols:
                                patched_tps = run_pgbench(protocol)

                                stop_postgres(postgres)

                                postgres = start_postgres(vanilla_postgres)
                                vanilla_tps = run_pgbench(protocol)
                                stop_postgres(postgres)

                                postgres = start_postgres(patched_postgres)

                                result = "{}\t{}\t{}\t{}\t{}\t{}\t{}\t{}\t{}".format(scale, projection,
                                                                                 index_only_fetch, page_visibility,
                                                                                 correlation, selectivity, protocol, vanilla_tps,
                                                                                 patched_tps)
                                print(result)

except subprocess.CalledProcessError as exc:
    print("Status : FAIL  \n {} \n {}".format(exc.returncode, exc.output))