#!/bin/sh

#
# This file is released under the terms of the Artistic License.
# Please see the file LICENSE, included in this package, for details.
#
# Copyright The DBT-2 Authors
#

if [ -z ${DBT2DBNAME} ]; then
	echo "DBT2DBNAME not defined."
	exit 1
fi

usage()
{
	cat << EOF
$(basename "${0}") is the Database Test 2 (DBT-2) PostgreSQL database builder.

Usage:
  $(basename "${0}") [OPTION]

Options:
  -l PORT        PostgreSQL PORT
  -p LIST        Sepecific LIST of parttitions to load, e.g. \"1 2\",
                 default: all
  --rows-per-commit ROWS
                 number of ROWS to load per transaction, default is to load all
                 data in a single transaction
  -s INTEGER     set random number generation seed
  -t PARTITIONS  total number of data PARTITIONS, default: detected number of
                 processors
  -w WAREHOUSES  number of WAREHOUSES to build, default: 1
  -?, --help     show this help, then exit

Database Test 2 (DBT-2) project page: https://github.com/osdldbt/dbt2
EOF
}

ROWS_PER_COMMIT=0
SEED=""
WAREHOUSES=1
while [ "${#}" -gt 0 ] ; do
	case "${1}" in
	(-l)
		shift
		PORT="${1}"
		;;
	(-p)
		shift
		PART_CURRENT="${1}"
		;;
	(--rows-per-commit)
		shift
		ROWS_PER_COMMIT="${1}"
		;;
	(--rows-per-commit=?*)
		ROWS_PER_COMMIT="${1#*--rows-per-commit=}"
		;;
	(-s)
		shift
		SEED="${1}"
		;;
	(-t)
		shift
		TOTAL_PARTS="${1}"
		;;
	(-w)
		shift
		WAREHOUSES="${1}"
		;;
	(-\? | --help)
		usage
		exit 0
		;;
	(--* | -*)
		echo "$(basename "${0}"): invalid option -- '${1}'"
		echo "try \"$(basename "${0}") --help\" for more information."
		exit 1
		;;
	(*)
		break
		;;
	esac
	shift
done

if [ ! "x${PORT}" = "x" ]; then
	PORTARG="-p ${PORT}"
fi

if [ "${ROWS_PER_COMMIT}" -gt 0 ]; then
	DATAGENARGS="--rows-per-commit=${ROWS_PER_COMMIT}"
fi

PSQL="psql -X ${PORTARG} -d ${DBT2DBNAME}"

# Generate a seed randomly if one is not provided.
if [ "x${SEED}" = "x" ]; then
	SEED=$(dbt2 rand -1 1 15)
fi

load_table() {
	TABLE=$1
	WAREHOUSES=$2
	SEED=$3

	if [ "${TOTAL_PARTS}" = "" ]; then
		# Detect the number of processors to determining how much to
		# parallelize the data loading.
		CPUS="$(nproc)"
		if [ "${WAREHOUSES}" -eq 1 ]; then
			# Don't parallelize the load if there is only 1 warehouse.
			CPUS=1
		elif [ "${WAREHOUSES}" -lt "${CPUS}" ]; then
			# Don't parallelize more than the number of warehouses if there are
			# less warehouses to build than the number of processors available.
			CPUS=$WAREHOUSES
		else
			echo "detected $CPUS processors for parallel loading"
		fi
		T=$CPUS
	else
		T=$TOTAL_PARTS
		CPUS=${TOTAL_PARTS}
	fi

	# When TOTAL_PARTS is 1, data isn't being partitiond so it doesn't need to
	# be parallelized.
	if [ "${TOTAL_PARTS}" = "1" ]; then
		P=1
	else
		if [ "${PART_CURRENT}" = "" ]; then
			P="$(seq 1 "${CPUS}")"
		else
			P="${PART_CURRENT}"
		fi
	fi

	$PSQL -c "ALTER TABLE $TABLE SET (autovacuum_enabled = false, toast.autovacuum_enabled = false);"
	for i in $P; do
		run_datagen $T $i $TABLE $WAREHOUSES $SEED &
	done
	wait
	$PSQL -c "ALTER TABLE $TABLE SET (autovacuum_enabled = true, toast.autovacuum_enabled = true);"
}

run_datagen()
{
	TOTAL=$1
	PART=$2
	TABLE=$3
	WAREHOUSES=$4
	SEED=$5

	eval dbt2 datagen --direct -w "${WAREHOUSES}" --pgsql --table "${TABLE}" \
			--seed "${SEED}" -P "${TOTAL}" -p "${PART}" "${DATAGENARGS}"
	if [ $? -ne 0 ]; then
		echo "ERROR: datagen couldn't load part $PART of $TOTAL with $TABLE"
		exit 1
	fi
}

# Load one table at a time, but each table load will be parallelized.
for TABLE in item warehouse stock district customer history orders new_order;
do
	load_table $TABLE $WAREHOUSES $SEED
done

ROWCOUNT="$(eval "${PSQL} -At" <<- EOF
	SELECT 1
	FROM pg_proc, pg_namespace
	WHERE pg_proc.pronamespace = pg_namespace.oid
	  AND nspname = 'bdr'
	  AND proname = 'wait_slot_confirm_lsn'
	LIMIT 1;
EOF
)"
if [ "${ROWCOUNT}" = "1" ]; then
	eval "${PSQL} -e" <<- EOF
		SELECT bdr.wait_slot_confirm_lsn(NULL, NULL);
	EOF
fi

exit 0
