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

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

Options:
  -l PORT        YugabyteDB 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}"
		;;
	(--rows-per-commit)
		shift
		ROWS_PER_COMMIT="${1}"
		;;
	(--rows-per-commit=?*)
		ROWS_PER_COMMIT="${1#*--rows-per-commit=}"
		;;
	(-s)
		shift
		SEED="${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 [ ! "${PORT}" = "" ]; then
	export PGPORT="${PORT}"
fi

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

# Generate a seed randomly if one is not provided.
if [ "${SEED}" = "" ]; 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

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

echo "loading up to $RATE warehouses at a time"

# Load one table at a time, and one partition of each table at a
# time.
for TABLE in item warehouse stock district customer history orders new_order;
do
	load_table "$TABLE" "$WAREHOUSES" "$SEED"
done
