#!/bin/bash
#set -x

#
# API performance test
#

#set -x

#
# parse args
#

FETCH_ROWS=1000
METHOD=none

function usage {
  echo
  echo "Usage: $(basename $0) [-r rows_to_fetch] [-m api_method]"
  echo "  API Methods:  SYNC, ASYNC, MSQ"
  echo
  exit 1
}

echo
while getopts 'abm:r:h' opt; do
  case "$opt" in
    a)
      echo "Processing option 'a'"
      ;;

    b)
      echo "Processing option 'b'"
      ;;

    m)
      METHOD="$OPTARG"
      ;;
   
    r)
      FETCH_ROWS="$OPTARG"
      ;;
   
    ?|h)
      usage
      ;;
  esac
done
shift "$(($OPTIND -1))"

if [[ $METHOD != "SYNC" && $METHOD != "ASYNC" && $METHOD != "MSQ" ]]; then
  echo "Invalid API Method: $METHOD"
  usage
fi
echo "API Method: $METHOD"
echo "Fetching $FETCH_ROWS rows"

#
# Assign SQL Query
#
#SQLTEXT="SELECT trip_id, pickup_longitude FROM trips_xaa where pickup_longitude != 0 limit $FETCH_ROWS"
SQLTEXT="SELECT count(*) Cnt FROM wikipedia"
echo "Query is: "$SQLTEXT

#
# START TIMER
#
sdt=$(date +%s)

#
# Async API Method
#

if [ $METHOD == 'ASYNC' ]; then

 echo
 echo "Using Async SQL API ..."
 echo

  # Launch select statement
  
  echo "Launching query ..."
  
  export TASK_ID=`curl --silent --location --request POST 'http://localhost:8888/druid/v2/sql/async/' \
       --header 'Content-Type: application/json' \
       --data-raw "{\"query\" : \"$SQLTEXT\", \"resultFormat\":\"csv\"}" | jq -r '.asyncResultId'`
  
  echo "Task ID :$TASK_ID"
  
  # Check status
  
  STATUS=RUNNING
  while [ $STATUS != 'COMPLETE' ]
  do
    echo -n "Checking status ... "
    export STATUS=`curl --silent --location --request GET "http://localhost:8888/druid/v2/sql/async/$TASK_ID/status" | jq -r '.state'`
    echo $STATUS
    sleep 1
  done
  
  # pull results
  
  echo "Pulling results to wc -l ..."
  
  curl --silent --location --request GET "http://localhost:8888/druid/v2/sql/async/$TASK_ID/results" | wc -l

#
# MSQ SQL API Method
#

elif [ $METHOD == 'MSQ' ]; then

 echo
 echo "Using MSQ SQL API ..."
 echo

  # Launch select statement
  
  echo "Launching query ..."
  
  export TASK_ID=`curl --silent --location --request POST 'http://localhost:8888/druid/v2/sql/statements' \
       --header 'Content-Type: application/json' \
       --data-raw "{\"query\" : \"$SQLTEXT\", \"context\": {\"executionMode\":\"ASYNC\"}, \"maxNumTasks\":2}, \"resultFormat\":\"csv\"}" | jq -r '.queryId'` 
  
  echo "Task ID: $TASK_ID"
  
  # Check status
  
  STATUS=RUNNING
  while [ $STATUS != 'SUCCESS' ]
  do
    echo -n "Checking status ... "
    export STATUS=`curl --silent --location --request GET "http://localhost:8888/druid/v2/sql/statements/$TASK_ID" | jq -r '.state'`
    echo $STATUS
    sleep 1
  done
  
  # pull results
  
  echo "Pulling results to wc -l ..."
  
  curl --silent --location --request GET "http://localhost:8888/druid/indexer/v1/task/$TASK_ID/reports" | jq | wc -l

#
# Sync SQL API Method
#

elif [ $METHOD == 'SYNC' ]; then

 echo
 echo "Using Sync SQL API ..."
 echo

  # Launch select statement
  
  echo "Launching query ..."
  
  curl --silent --location --request POST 'http://localhost:8888/druid/v2/sql/' \
     --header 'Content-Type: application/json' \
     --data-raw "{ \"query\" : \"$SQLTEXT\", \"resultFormat\":\"csv\"}" | wc -l

#
# End options
#

fi 

#
# END TIMER
#

edt=$(date +%s)

echo "$edt $sdt" | awk ' { printf("Time: %.2f seconds.\n", ($1-$2)); } '

#
# END
#