#!/bin/bash

# Assuming by default that testclusters are on localhost
# Specified DB will also be used for storing pg_stat_statement snapshots from all clusters for easier analyzing
export PGHOST=/var/run/postgresql
export PGPORT=5432
export PGUSER=postgres
export PGDATABASE=postgres

PGBENCH=/usr/lib/postgresql/11/bin/pgbench
PSQL=/usr/lib/postgresql/11/bin/psql

# Add/remove hosts as needed here
connection_str[0]=''		# 1st instance doesn't need adjusting
connection_name[0]="10"
connection_str[1]='-p 5433'
connection_name[1]="11rc1"

TEST_MODE="analytics"
SCALE="100"	# 100 = 10mio pgbench_accounts rows
TEST_DURATION=600	# 10min
LOOPS_EACH_INSTANCE=36
DO_PGBENCH_INIT=1

TEST_QUERIES=$(cat <<-HERE
SELECT sum(abalance) FROM pgbench_accounts CROSS JOIN generate_series(1, 5) where aid % 2 = 0; \n
SELECT count(*) FROM pgbench_accounts JOIN pgbench_accounts_copy using (aid) where aid % 2 = 0; \n
SELECT count(*) FROM (SELECT aid, bid, count(*) FROM pgbench_accounts where aid % 2 = 0 GROUP BY CUBE (aid, bid)) a;\n
SELECT COUNT(DISTINCT aid) FROM pgbench_accounts where aid % 2 = 0;
HERE
)


# Test conn and do pg_stat_statement and pgbench setup
for conn_str in "${connection_str[@]}" ; do

	$PSQL $conn_str -qXc "select 1" &>/dev/null
	if [ "$?" -ne 0 ] ; then
		echo "could not connect to $conn_str, check connection params. exiting..."
		exit 1
	fi

	if [ "$DO_PGBENCH_INIT" -gt 0 ] ; then

	echo "init scale $SCALE ..."
	$PGBENCH "$conn_str" -i -q --unlogged-tables --foreign-keys -s $SCALE
	$PSQL $conn_str -qXc "drop table if exists pgbench_accounts_copy"
	$PSQL $conn_str -qXc "create unlogged table if not exists pgbench_accounts_copy as select * from pgbench_accounts"
	$PSQL $conn_str -qXc "create unique index if not exists pgbench_accounts_copy_aid_idx ON pgbench_accounts_copy (aid)"
	$PSQL $conn_str -qXc "vacuum analyze pgbench_accounts_copy"

	fi

	$PSQL $conn_str -qXc "checkpoint"
	$PSQL $conn_str -qXc "create extension if not exists pg_stat_statements"
	$PSQL $conn_str -qXc "select pg_stat_statements_reset()"
done


# Initialize the table for storing pg_stat_statement results
$PSQL -qc "create table if not exists my_pg_stat_statements as select ''::text as testset, ''::text as mode, 0 as scale, 0 as clients, now() as created_on, * from pg_stat_statements where false"
#$PSQL -qc "truncate table my_pg_stat_statements"


# Test
for loop in $(seq 1 ${LOOPS_EACH_INSTANCE}) ; do
	i=0
	for conn_str in "${connection_str[@]}" ; do
		echo "doing loop $loop on ${connection_name[i]} - scale $scale, duration $TEST_DURATION s ..."
		echo -e $TEST_QUERIES |  $PGBENCH $conn_str -T $TEST_DURATION -f-
		i=$((i + 1))
	done
done

# Store results for this run
i=0
for conn_str in "${connection_str[@]}" ; do
	echo "storing pg_stat_statements results for ${connection_name[i]}..."
	$PSQL $conn_str -qXc "copy (select '${connection_name[$i]}', '$TEST_MODE', ${SCALE:-NULL}, ${clients:-1}, now(), * from pg_stat_statements where query ~ '^(INSERT|UPDATE|SELECT).*pgbench_') to stdout" \
	| $PSQL -qXc "copy my_pg_stat_statements from stdin"
	i=$((i + 1))
done

echo "done"