#!/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) CockroachDB data loader.

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

Options:
  -l PORT        PostgreSQL PORT
  --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
  -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
}

RATE=100
ROWS_PER_COMMIT=0
SEED=""
WAREHOUSES=1
while [ "${#}" -gt 0 ] ; do
	case "${1}" in
	(-l)
		shift
		PORT="${1}"
		;;
	(-s)
		shift
		SEED="${1}"
		;;
	(--rows-per-commit)
		shift
		ROWS_PER_COMMIT="${1}"
		;;
	(--rows-per-commit=?*)
		ROWS_PER_COMMIT="${1#*--rows-per-commit=}"
		;;
	(-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

	# simulate ceiling function
	T=$(( ($WAREHOUSES + $RATE - 1) / $RATE ))
	P=`seq 1 ${T}`
	for i in ${P}; do
		run_datagen $T $i $TABLE $WAREHOUSES $SEED
	done
}

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
}

echo "loading $RATE warehouses at a time"

# 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
