#!/usr/bin/env bash
#
# Partition Table Performance Check
#
HERE="$(cd -P "$(dirname "${BASH_SOURCE:-$0}")" && pwd -P)"

# trap INT QUIT TERM
trap 'exit' 2 3 15

export PGDATA="$HERE/pgdata"
export PGDATABASE=postgres
PSQL="psql -q"
#PSQL="psql -a"

# Options
declare -A args

current_time() { date +'%Y-%m-%d %H:%M:%S'; }

# Help
usage() {
  cat <<EOF
usage: $(basename "$0") [setup|bench] <args>
  setup [--part|--all] [<range>] [<tablenum>]
  bench <[part|part_child|nopart]> [--procnum <num>] [--dup]
                                   [<range>] [<use_range>]
                                   [<table_min>] [<table_max>]

EOF
}

# Do vacuum analyze
vacuum_analyze() {
  echo "vacuum analyze" | $PSQL
}

# Dump record numbers per table
dump_table_stats() {
  $PSQL <<-SQL
	select relname, n_live_tup
	from pg_stat_user_tables
	where schemaname = 'public'
	order by relname;
SQL
}

# Setup non-partition tables.
init_non_part_tables() {
  local range="$1"
  local tablenum="$2"
  local enable_index="$3"

  for i in $(seq 1 $tablenum); do
    table="nopart_tbl_${i}"
    val_from=$(((i - 1) * range + 1))
    val_to=$((val_from + range - 1))
    $PSQL <<-SQL
	create table ${table}(id int, msg text);
	insert into ${table} select i, 'test-' || i
	  from generate_series(${val_from}, ${val_to}) as i;
	-- select min(id), max(id) from ${table};
SQL

    if ((enable_index == 1 )); then
      $PSQL -c "create index ${table}_idx on ${table} (id);"
    fi
  done
}

# Setup partition tables.
init_part_tables() {
  local range="$1"
  local tablenum="$2"
  local enable_index="$3"

  parent_table="part_tbl"
  $PSQL <<-SQL
	create table ${parent_table}(id int, msg text)
	partition by range (id);
SQL

  for i in $(seq 1 $tablenum); do
    part_table="${parent_table}_${i}"
    range_from=$(( (i - 1) * range + 1 ))
    range_to=$(( range_from + range ))
    val_from=$range_from
    val_to=$(( range_to - 1 ))

    $PSQL <<-SQL
	create table ${part_table} partition of ${parent_table}
	  for values from ($range_from) to ($range_to);
	insert into ${parent_table} select i, 'test-' || i
	  from generate_series($val_from, $val_to) as i;
--	select min(id), max(id) from ${part_table};
SQL

    if ((enable_index == 1 )); then
      $PSQL -c "create index ${part_table}_idx on ${part_table} (id);"
    fi
  done
}

# Do benchmark test.
do_bench() {
  echo "---------------------------"

  echo "# conditions:"
  local k v
  for k in ${!args[@]}; do
    printf "    [%s]=%s\n" "$k" "${args[$k]}"
  done
  echo

  # Single process
  if [ ${args[procnum]} -eq 1 ]; then
    echo "# start with single proc: $(current_time)"
    time python "$HERE"/bench.py "${args[table]}" ${args[range]} \
         ${args[table_min]} ${args[table_max]} ${args[use_range]}
  else
    # Multi process
    if [ ${args[dup]} -eq 0 ]; then
      local tablenum=$(( ${args[table_max]} - ${args[table_min]} + 1 ))
      local divmod=$(( tablenum % ${args[procnum]} ))
      local task_per_proc=$(( tablenum / ${args[procnum]} ))
      (( divmod != 0 )) && (( task_per_proc++ ))

      echo "#   procnum: ${args[procnum]}"
      echo "#   task_per_proc: $task_per_proc"

      local table_first=
      local table_last=$(( table_no + task_per_proc - 1 ))
      echo "# start with multi proc: $(current_time)"
      for table_first in $(seq ${args[table_min]} $task_per_proc ${args[table_max]}); do
        table_last=$(( table_first + task_per_proc - 1 ))
        echo "table_no: $table_first - $table_last"
        time python "$HERE"/bench.py "${args[table]}" ${args[range]} \
             $table_first $table_last ${args[use_range]} &

        (( i++ ))
      done
    else
      for proc_no in $(seq 1 ${args[procnum]}); do
        echo "table_no: ${args[table_min]} - ${args[table_max]}"
        time python "$HERE"/bench.py "${args[table]}" ${args[range]} \
             ${args[table_min]} ${args[table_max]} ${args[use_range]} &

        (( i++ ))
      done


    fi
    # Wait until all processes have done.
    wait
  fi

  echo
  echo "# end: $(current_time)"
  echo "---------------------------"
  echo
}




com="$1" ; shift
case "$com" in

  # Usage
  (help)
    usage && exit
    ;;

  # Setup database
  (setup)

    args=(
      [part]=0
      [all]=0
      [index]=0
    )

    while [ $# -gt 0 ]; do
      case "$1" in
        -p|--part)
          args[part]=1
          ;;
        -a|--all)
          args[all]=1
          ;;
        -i|--index)
          args[index]=1
          ;;
        *)
          break
          ;;
      esac
      shift
    done

    range=${1:-1000}
    tablenum=${2:-100}

    cat <<EOF
===========================================
Test Environment $(current_time)
===========================================
Partition Range  : $range
Max Table Nubmers: $tablenum
-------------------------------------------
EOF

      # Clean workplace.
      {
        pg_ctl stop -m i
        /bin/rm -rf "$PGDATA"
        initdb

        cat <<-CONF >> "$PGDATA"/postgresql.conf
	shared_buffers = 1024MB
	log_destination = 'stderr'
	logging_collector = on
	log_directory = 'log'
	log_filename = 'postgresql-%Y-%m-%d.log'
	#log_statement = all
	enable_indexscan = off
	checkpoint_completion_target = 0.9
CONF

        pg_ctl start

      } > /dev/null 2>&1

      if [ ${args[part]} -eq 0 ] || [ ${args[all]} -eq 1 ]; then
        # Setup non-partition tables.
        init_non_part_tables $range $tablenum ${args[index]}
      fi

      if [ ${args[part]} -eq 1 ] || [ ${args[all]} -eq 1 ]; then
        # Setup partition tables.
        init_part_tables $range $tablenum ${args[index]}
      fi

      vacuum_analyze

      # Make sure that records are put into tables.
#      dump_table_stats

      ;;


  # Do trivial bench
  (bench)

    [ -z "$1" ] && usage && exit

    args=(
      [procnum]=1
      [dup]=0
    )

    while [ $# -gt 0 ]; do
      case "$1" in
        --dup)
          args[dup]=1
          ;;

        -P|--procnum)
          shift ; args[procnum]="$1"
          ;;
        *)
          break
          ;;
      esac
      shift
    done

    if [ "$1" = "part" ]; then
      args[table]="part_tbl"
    elif [ "$1" = "part_child" ]; then
      args[table]="part_tbl_%d"
    elif [ "$1" = "nopart" ]; then
      args[table]="nopart_tbl_%d"
    fi

    args[range]=${2:-1000}
    args[table_min]=${3:-1}
    args[table_max]=${4:-10}
    args[use_range]=${5:-${args[range]}}

    echo "==> start bench: ${args[table]}"
    do_bench

    ;;

  (*)
    echo "invalid command: $1" 2>&1
    exit 1
    ;;
esac
