perf-data-loader

#!/bin/bash

# Initialize variables
CONFIG=""
TABLE_NAME=""
ROWS=""
GENERATE_ONLY=false
LOAD_ONLY=false
GEN_LOAD_ORDER=false
JAR_PATH=""

# ANSI color codes
GREEN="\033[32m"
RED="\033[31m"
RESET="\033[0m"

# Function to display help
function display_help {
    echo "Usage: $0 --config <config_file> --table-name <table_name> --rows <rows> [--gen-config-only] [--load-only] [--gen-load-order]"
    echo "Short forms: -c <config_file> -t <table_name> -r <rows>"
    echo "Options:"
    echo "  -c, --config                Configuration file"
    echo "  -t, --table-name            Table name"
    echo "  -r, --rows                  Number of rows"
    echo "  --gen-config-only           Only generate the loader/config file"
    echo "  --load-only                 Only load data into the database"
    echo "  --gen-load-order            Generate table load order in the provided database"
    echo "  -h, --help                  Display this help message"
    exit 0
}

# Function to find the jar file
function find_jar {
    # Check current directory first
    if [[ -f "benchbase.jar" ]]; then
        JAR_PATH="benchbase.jar"
        return
    fi

    # Check target/benchbase-yugabyte/
    if [[ -f "target/benchbase-yugabyte/benchbase.jar" ]]; then
        JAR_PATH="target/benchbase-yugabyte/benchbase.jar"
        return
    fi

    # Check target/benchbase-postgres/
    if [[ -f "target/benchbase-postgres/benchbase.jar" ]]; then
        JAR_PATH="target/benchbase-postgres/benchbase.jar"
        return
    fi

    echo "Error: benchbase.jar not found in any expected location."
    exit 1
}

# Function to show a fake progress bar while a command is running
show_progress_bar() {
    local CMD_PID=$1
    local SECONDS=0

    while kill -0 "$CMD_PID" 2>/dev/null; do
        printf "\r["
        for ((i=0; i<50; i++)); do
            if (( i < (SECONDS % 50) )); then
                printf "="
            else
                printf " "
            fi
        done
        printf "] %d%%" $((SECONDS % 50 * 2))
        sleep 0.1
    done
}

# Function to check command success or failure
# Function to check command success or failure
check_command_status() {
    local CMD_PID=$1
    local SUCCESS_MSG=$2
    local FAILURE_MSG=$3
    local LOG_FILE=$4

    wait "$CMD_PID"
    CMD_EXIT_CODE=$?
    # Clear the progress bar line
    printf "\r\033[K"
    if [ $CMD_EXIT_CODE -eq 0 ]; then
        # Ensure the progress bar finishes at 100% on success
        printf "\r[==================================================] 100%%\n"
        echo -e "${GREEN}$SUCCESS_MSG${RESET}"
    else
        # Avoid showing 100% if the command fails
        echo -e "${RED}$FAILURE_MSG${RESET}"
        echo -e "${RED}\nConsole log:${RESET}"
        cat "$LOG_FILE"
    fi
}

# Parse arguments
while [[ $# -gt 0 ]]; do
    case $1 in
        -c|--config)
            CONFIG="$2"
            shift # past argument
            shift # past value
            ;;
        -t|--table-name)
            TABLE_NAME="$2"
            shift # past argument
            shift # past value
            ;;
        -r|--rows)
            ROWS="$2"
            shift # past argument
            shift # past value
            ;;
        --gen-config-only)
            GENERATE_ONLY=true
            shift # past argument
            ;;
        --load-only)
            LOAD_ONLY=true
            shift # past argument
            ;;
        --gen-load-order)
            GEN_LOAD_ORDER=true
            shift # past argument
            ;;
        -h|--help)
            display_help
            ;;
        *)
            echo "Invalid option: $1"
            display_help
            ;;
    esac
done

# Ensure that both --gen-config-only and --load-only are not provided simultaneously
if { [ "$GENERATE_ONLY" = true ] && [ "$LOAD_ONLY" = true ]; } || \
   { [ "$GENERATE_ONLY" = true ] && [ "$GEN_LOAD_ORDER" = true ]; } || \
   { [ "$LOAD_ONLY" = true ] && [ "$GEN_LOAD_ORDER" = true ]; }; then
    echo "Error: Cannot use these conflicting parameters simultaneously. --gen-config-only, --load-only, --gen-load-order ."
    exit 1
fi

# Find the JAR file
find_jar

if [ "$LOAD_ONLY" = true ]; then
    if [ -z "$CONFIG" ]; then
        echo "Error: --config parameter is required with --load-only."
        display_help
    fi

    # Create a temporary file to capture output and errors
    LOG_FILE=$(mktemp)
    echo "Loading the data into the database now..."
    # Start the Java command in the background
    java -jar "$JAR_PATH" -b featurebench -c "$CONFIG" --load=True > "$LOG_FILE" 2>&1 &
    CMD_PID=$!
    # Show the progress bar while the command is running
    show_progress_bar "$CMD_PID"
    # Check if the command was successful and print the appropriate message
    check_command_status "$CMD_PID" \
        "Data load to the table \`${TABLE_NAME}\` is successful." \
        "Failed to load data into the table: ${TABLE_NAME}." \
        "$LOG_FILE"

    # Clean up the log file
    rm "$LOG_FILE"
    exit 0
fi

if [ "$GENERATE_ONLY" = true ]; then
    if [ -z "$CONFIG" ] || [ -z "$TABLE_NAME" ] || [ -z "$ROWS" ]; then
        echo "Error: --config, --table-name, and --rows parameters are required with --gen-config-only."
        display_help
    fi
    # Create a temporary file to capture output and errors
    LOG_FILE=$(mktemp)
    echo "Generating loader file for the table..."
    # Start the Java command in the background
    java -jar "$JAR_PATH" -b perf-dataloader -c "$CONFIG" -p tableName="$TABLE_NAME" -p rows="$ROWS" --load=True > "$LOG_FILE" 2>&1 &
    CMD_PID=$!
    # Show the progress bar while the command is running
    show_progress_bar "$CMD_PID"
    # Check if the command was successful and print the appropriate message
    check_command_status "$CMD_PID" \
        "Loader file generated successfully: ${TABLE_NAME}_loader.yaml" \
        "Failed to generate loader file." \
        "$LOG_FILE"

    # Clean up the log file
    rm "$LOG_FILE"
    exit 0
fi

if [ "$GEN_LOAD_ORDER" = true ]; then
    if [ -z "$CONFIG" ]; then
        echo "Error: --config parameter is required with --gen-load-order."
        display_help
    fi
    # Create a temporary file
    TEMP_CONFIG=$(mktemp)
    # Create a temporary file to capture output and errors
    LOG_FILE=$(mktemp)
    # Copy the content of the original CONFIG file to the temporary file
    cp "$CONFIG" "$TEMP_CONFIG"

    # Append text to the temporary config file (you can customize the text as needed)
    echo >> "$TEMP_CONFIG"
    echo "gen-db-load-order: true" >> "$TEMP_CONFIG"
    echo "Generating load order based on the provided config file..."
    java -jar "$JAR_PATH" -b perf-dataloader -c "$TEMP_CONFIG" -p tableName="dummy" -p rows="1" --load=True > "$LOG_FILE" 2>&1 &
    CMD_PID=$!
    # Show the progress bar while the command is running
    show_progress_bar "$CMD_PID"
    # Check if the command was successful and print the appropriate message
    check_command_status "$CMD_PID" \
        "Load order generated successfully in file \`load_order.json\` and is as follows: " \
        "Failed to generate load order for the database." \
        "$LOG_FILE"
    # If successful, print the contents of the load order file
    if [ $CMD_EXIT_CODE -eq 0 ]; then
        cat "load_order.json"
        echo ""
    fi
    # Clean up: delete the temporary config file
    rm -f "$TEMP_CONFIG"
    # Clean up the log file
    rm "$LOG_FILE"
    exit 0
fi


# Check if required parameters are provided for the default operation
if [ -z "$CONFIG" ] || [ -z "$TABLE_NAME" ] || [ -z "$ROWS" ]; then
    echo "Error: Missing required parameters."
    display_help
fi

# If no specific option is provided, do both generate and load with default config file location
# Create a temporary file to capture output and errors
LOG_FILE=$(mktemp)
echo "Generating loader file for the table..."
# Start the Java command in the background
java -jar "$JAR_PATH" -b perf-dataloader -c "$CONFIG" -p tableName="$TABLE_NAME" -p rows="$ROWS" --load=True > "$LOG_FILE" 2>&1 &
CMD_PID=$!
# Show the progress bar while the command is running
show_progress_bar "$CMD_PID"
# Check if the command was successful and print the appropriate message
check_command_status "$CMD_PID" \
    "Loader file generated successfully: ${TABLE_NAME}_loader.yaml" \
    "Failed to generate loader file." \
    "$LOG_FILE"

# Clean up the log file
rm "$LOG_FILE"

DEFAULT_CONFIG="${TABLE_NAME}_loader.yaml"

# Create a temporary file to capture output and errors
LOG_FILE=$(mktemp)
echo "Loading the data into the database now..."
# Start the Java command in the background
java -jar "$JAR_PATH" -b featurebench -c "$DEFAULT_CONFIG" --load=True > "$LOG_FILE" 2>&1 &
CMD_PID=$!
# Show the progress bar while the command is running
show_progress_bar "$CMD_PID"
# Check if the command was successful and print the appropriate message
check_command_status "$CMD_PID" \
    "Data load to the table \`${TABLE_NAME}\` is successful." \
    "Failed to load data into the table: ${TABLE_NAME}." \
    "$LOG_FILE"

# Clean up the log file
rm "$LOG_FILE"