#!/bin/bash
#
# Вариант запуска для разработки: 'cd "pg_source"/se-test/pgacext && sudo ./runtests -d . -u -n'
#

export SCRIPT_VERSION="1.1"

COLOR_B_RED="\033[1;31m"	# format output text bold red
COLOR_CYAN="\033[0;36m"		# format output text cyan
COLOR_OFF="\033[0m"			# format output text reset

__user_exists() { id "$1" &>/dev/null; }  # Существует ли указанный пользователь

__timestamp() { date +"%s%N"; }  # Текущее время в наносекундах

function clean_everything() {
    echo ''

    # Сбрасываем перехват сигналов
    trap - EXIT SIGINT
    # По умолчанию выход с кодом 1, если другой не был передан
    local exit_code="${1:-1}"

    [[ "$ARG_CMD" != "clean" ]] && printf "\n"
    printf "Очистка"
    [[ "$ARG_CMD" != "clean" ]] && printf " перед выходом из скрипта:\n"
    [[ "$ARG_CMD" = "clean" ]] && printf ":\n"

    echo "- Возвращаем настройки pam"
    mv /etc/pam.d/su.bak /etc/pam.d/su &> /dev/null
    mv /etc/pam.d/sudo.bak /etc/pam.d/sudo &> /dev/null

    usermod -r -G shadow postgres &> /dev/null

    if [[ ! $ARG_NO_STOP ]]; then
        echo "- Остановка кластеров"
        pg_ctlcluster $PG_VERSION $PG_SETEST_CLUSTER stop &> /dev/null    &
        pg_ctlcluster $PG_VERSION $PG_SEFOREIGN_CLUSTER stop &> /dev/null &
        pg_ctlcluster $PG_VERSION $PG_FILES_CLUSTER stop &> /dev/null     &
    fi

    if [[ ! $ARG_NO_RM ]]; then
        __remove_clusters
    fi

    if [[ ! $ARG_NO_RM_USERS ]]; then
        echo "- Удаление тестовых пользователей"
        ./support/remove-users &> /dev/null &
    fi

    rm -rf $PG_TEMP_DIR
    unlink "/tmp/pgacext_support" &> /dev/null || true

    wait

    echo "Очистка завершена."
    exit $exit_code
}



# Форматированное времени, представленное в наносекундах
function __display_time() {
  local T=$1
  local D=$((T/1000000000/60/60/24))
  local H=$((T/1000000000/60/60%24))
  local M=$((T/1000000000/60%60))
  local S=$((T/1000000000%60))
  local MS=$((T/1000000%1000))

  [[ $D > 0 ]] && printf '%dd ' $D
  [[ $H > 0 ]] && printf '%dh ' $H
  [[ $M > 0 ]] && printf '%dm ' $M
  [[ $S > 0 ]] && printf '%ds ' $S

  printf '%dms' $MS
}



function set_variables() {
    export PG_SETEST_CLUSTER="setest"
    export PG_SEFOREIGN_CLUSTER="seforeign"
    export PG_FILES_CLUSTER="sefiles"
    export PG_SETEST_PORT=6000
    export PG_SEFOREIGN_PORT=6001
    export PG_FILES_PORT=6002

    export PG_CLUSTERS="$PG_SETEST_CLUSTER $PG_SEFOREIGN_CLUSTER $PG_FILES_CLUSTER"

    if [[ "$ARG_PG_VERSION" != "" ]]; then
        export PG_VERSION="$ARG_PG_VERSION"
    else
        # выставляем отдельно ещё раз на случай. если эта функция будет вызвана отдельно от main()
        export PG_VERSION=$(/usr/share/postgresql-common/supported-versions | tail -n 1)
    fi

    if [[ "$ARG_TIME_QUANT" != "" && $(($ARG_TIME_QUANT + 0)) ]]; then
        export TIME_QUANT=$ARG_TIME_QUANT
    else
        export TIME_QUANT=1
    fi

    PG_WORKDIR="$ARG_WORKING_DIR"

    PG_CONFIG_DIR="/etc/postgresql/$PG_VERSION"
    PG_DATA_DIR="/var/lib/postgresql/$PG_VERSION"

    PG_TABLE_SPACE="$PG_DATA_DIR/testspace"
    PG_TABLE_SPACE_2="/sefiles_tablespace"

    PG_TEMP_DIR="/tmp_test"

    PG_TESTS_LIST=''

    PG_CLUSTERS_LIST=($PG_SETEST_CLUSTER $PG_SEFOREIGN_CLUSTER $PG_FILES_CLUSTER)
    PG_PORTS_LIST=($PG_SETEST_PORT $PG_SEFOREIGN_PORT $PG_FILES_PORT)

    if [[ "$ARG_RES_DIR" != "" ]]; then
        PG_TESTS_RES_DIR="$(realpath "$ARG_RES_DIR")"
    elif [[ "${SUDO_UID}" == "" ]]; then
    	PG_TESTS_RES_DIR="/root/pg-$PG_VERSION-se-test"
    else
    	PG_TESTS_RES_DIR="$(getent passwd | grep x:${SUDO_UID}: | cut -d : -f 6)/pg-$PG_VERSION-se-test"
    fi

    # TODO: может быть сделать определение данных переменных в psql скриптах
    # Сохраняем переменные во временный файл
#    PG_VARS_FILE="/tmp/pg_se-test_vars"
#    rm -f $PG_VARS_FILE
    if [[ "$1" = "print" ]]; then
        for var in 'PG_VERSION' 'PG_WORKDIR' 'PG_SETEST_CLUSTER' 'PG_SEFOREIGN_CLUSTER' 'PG_FILES_CLUSTER' \
                   'PG_SETEST_PORT' 'PG_SEFOREIGN_PORT' 'PG_FILES_PORT' 'PG_TABLE_SPACE' 'PG_TABLE_SPACE_2' \
                   'PG_CLUSTERS' 'PG_TESTS_RES_DIR'
        do
            printf 'export %s="%s"\n' "$var" "${!var}" # >> $PG_VARS_FILE
        done
    fi
#    chmod 777 $PG_VARS_FILE
}



function __remove_clusters() {
    echo "- Удаление тестовых табличных пространств в ФС"
    rm -rf $PG_TABLE_SPACE   &
    rm -rf $PG_TABLE_SPACE_2 &

    echo "- Удаление кластеров"
    pg_dropcluster $PG_VERSION $PG_SETEST_CLUSTER --stop &> /dev/null    &
    pg_dropcluster $PG_VERSION $PG_SEFOREIGN_CLUSTER --stop &> /dev/null &
    pg_dropcluster $PG_VERSION $PG_FILES_CLUSTER --stop &> /dev/null     &

    wait

    rm -rf $PG_CONFIG_DIR/$PG_SETEST_CLUSTER    &
    rm -rf $PG_CONFIG_DIR/$PG_SEFOREIGN_CLUSTER &
    rm -rf $PG_CONFIG_DIR/$PG_FILES_CLUSTER     &

    rm -rf $PG_DATA_DIR/$PG_SETEST_CLUSTER      &
    rm -rf $PG_DATA_DIR/$PG_SEFOREIGN_CLUSTER   &
    rm -rf /$PG_FILES_CLUSTER                   &
}



function __create_clusters() {
    echo "+ Создание кластеров"

    local count=0
    local option=""
    local dummy_locale="$PG_WORKDIR/support/locale"

    chmod +x "$dummy_locale"

    [[ $ARG_CLUSTERS == "" || $ARG_CLUSTERS != *"$PG_FILES_CLUSTER"* ]] && mkdir /$PG_FILES_CLUSTER &&
    pdpl-file 3::3:ccnr,irelax /$PG_FILES_CLUSTER &&
    chown postgres:postgres /$PG_FILES_CLUSTER

    for cluster in $PG_CLUSTERS
    do
        # Если указан запуск только конкретных кластеров, проверяем, нужно ли создавать $cluster
        [[ $ARG_CLUSTERS != "" && $ARG_CLUSTERS != *"$cluster"* ]] && continue
        [[ "$cluster" = "$PG_FILES_CLUSTER" ]] && option="-D /$PG_FILES_CLUSTER"

        # Процессы PostgreSQL пока что запускаются с серверной меткой {0,0}
        # Для initdb в /$PG_FILES_CLUSTER {3,3} нужны привилегии:
        # PARSEC_CAP_IGNMACLVL | PARSEC_CAP_IGNMACCAT | PARSEC_CAP_IGNMACINT (execaps -c 0x2030)
        unshare --mount bash -c "mount --bind $dummy_locale /usr/bin/locale &&
        execaps -c 0x2030 -- pg_createcluster $PG_VERSION $cluster $option --port ${PG_PORTS_LIST[count]} --start-conf=manual > /dev/null" &

        (( count++ ))
    done

    wait
}



function __create_tablespaces() {
    echo "+ Создание тестовых табличных пространств в ФС"

    if [[ ! -e "$PG_TABLE_SPACE" ]]; then
        mkdir -p "$PG_TABLE_SPACE"
        chown postgres:postgres "$PG_TABLE_SPACE"
    fi

    if [[ ! -e "$PG_TABLE_SPACE_2" ]]; then
        mkdir -p "$PG_TABLE_SPACE_2"
        chown postgres:postgres "$PG_TABLE_SPACE_2"
        chmod 770 "$PG_TABLE_SPACE_2"
        pdpl-file 3:0:3:ccnr "$PG_TABLE_SPACE_2"
    fi
}



function __configure_clusters() {
    echo "+ Настройка конфигурации кластеров"

    # Настройка конфигурации основного сервера
    if [[ $ARG_CLUSTERS == "" || $ARG_CLUSTERS == *"$PG_SETEST_CLUSTER"* ]]; then
        pg_conftool $PG_VERSION $PG_SETEST_CLUSTER set shared_preload_libraries 'pg_hint_plan'  # , online_analyze, plantuner
        pg_conftool $PG_VERSION $PG_SETEST_CLUSTER set wal_level 'logical'
        pg_conftool $PG_VERSION $PG_SETEST_CLUSTER set max_logical_replication_workers 8
        pg_conftool $PG_VERSION $PG_SETEST_CLUSTER set max_worker_processes 16
        cp "./support/pg_hba.conf.tst" "$PG_CONFIG_DIR/$PG_SETEST_CLUSTER/pg_hba.conf"
    fi

    # Настройка конфигурации внешнего сервера
    if [[ $ARG_CLUSTERS == "" || $ARG_CLUSTERS == *"$PG_SEFOREIGN_CLUSTER"* ]]; then
        pg_conftool $PG_VERSION $PG_SEFOREIGN_CLUSTER set ac_ignore_socket_maclabel false
        pg_conftool $PG_VERSION $PG_SEFOREIGN_CLUSTER set wal_level 'logical'
        pg_conftool $PG_VERSION $PG_SEFOREIGN_CLUSTER set max_logical_replication_workers 8
        pg_conftool $PG_VERSION $PG_SEFOREIGN_CLUSTER set max_worker_processes 16
        cp "./support/pg_hba_foreign.conf.tst" "$PG_CONFIG_DIR/$PG_SEFOREIGN_CLUSTER/pg_hba.conf"
    fi

    # Настройка конфигурации дополнительного сервера с метками на файлах
    if [[ $ARG_CLUSTERS == "" || $ARG_CLUSTERS == *"$PG_FILES_CLUSTER"* ]]; then
        pg_conftool $PG_VERSION $PG_FILES_CLUSTER set ac_enable_maclabels_on_files true
        cp "./support/pg_hba.conf.tst" "$PG_CONFIG_DIR/$PG_FILES_CLUSTER/pg_hba.conf"
    fi

    local data_dir="$PG_DATA_DIR"
    for cluster in $PG_CLUSTERS
    do
        # Если указан запуск только конкретных кластеров, проверяем, нужно ли настраивать $cluster
        [[ $ARG_CLUSTERS != "" && $ARG_CLUSTERS != *"$cluster"* ]] && continue

        [[ "$cluster" = "$PG_FILES_CLUSTER" ]] && data_dir=""
        cp "./support/auto.conf.$cluster" "$data_dir/$cluster/postgresql.auto.conf"
        chown -R postgres:postgres "$PG_CONFIG_DIR/$cluster"
        chown -R postgres:postgres "$data_dir/$cluster"
    done

    if [[ $ARG_CMD == "pgbench" || $ARG_CMD == "pb" ]]; then
        cat "./support/auto.conf.$PG_FILES_CLUSTER.pgbench" >> "/$PG_FILES_CLUSTER/postgresql.auto.conf"
    fi
}


function __additional_environment_configurations() {
    echo "+ Генерация файлов из шаблонов"
    mkdir -p "expected/filtered"
    python3 template.py all &

    echo "+ Настройка необходимых прав пользователю postgres"
    setfacl -m u:postgres:rx /etc/parsec/macdb
    setfacl -m u:postgres:rx /etc/parsec/capdb
    setfacl -d -m u:postgres:r /etc/parsec/macdb
    setfacl -d -m u:postgres:r /etc/parsec/capdb
    setfacl -R -m u:postgres:r /etc/parsec/macdb/*
    setfacl -R -m u:postgres:r /etc/parsec/capdb/*
    usermod -a -G shadow postgres

    echo "+ Меняем настройки pam"
    sed -i.bak 's/account required pam_su.so/#account required pam_su.so/g' /etc/pam.d/su
    sed -i.bak 's/account required pam_sudo.so/#account required pam_sudo.so/g' /etc/pam.d/sudo

    echo "- Удаление результатов предыдущих тестов"
    rm -rf "$PG_TESTS_RES_DIR"/*

    echo "+ Создание и настройка директорий"

    [[ $ARG_LOGS_TO_RES_DIR -eq 1 ]] && mkdir -p "$PG_TESTS_RES_DIR/logs"
    mkdir -p "$PG_TESTS_RES_DIR/results"
    chmod -R 777 "$PG_TESTS_RES_DIR"

    mkdir -p "$PG_TEMP_DIR"
    pdpl-file 3:63:3:ccnr,irelax "$PG_TEMP_DIR"
    chmod 777 "$PG_TEMP_DIR"

    ln -sf "$PG_WORKDIR/support" "/tmp/pgacext_support"

    ### Дополнительные настройки в режиме разработки
     if [[ "$PG_WORKDIR" != "/usr/share/postgresql/$PG_VERSION/test/pgacext" ]]; then
         # logrotate /etc/logrotate.d/syslog-ng-mod-astra

         # Позволяем другим пользователям (в частности postgres)
         # выполнять команды в нашей домашней директории без выдачи ошибки.
         chmod o+X $(dirname $PG_TESTS_RES_DIR)
     fi
    ###

    wait
}


function __parse_file_as_a_list() {
    local _file=$1

    # Удаляем возможные пробелы и табуляцию с конца строк
    sed -ie 's/[ \t]*$//' $_file
    # Убираем все пустые строки, а также закомментированные строки
    local tests="$(cat $_file | grep -v '#' | grep -v '^$' | sort)"

    for t in $tests
    do
        PG_TESTS_LIST="$PG_TESTS_LIST sql/$t.sql"
    done
}


function __set_which_tests_to_run() {
    if [[ "$ARG_FILE" != "" ]]; then
        [[ ! "$ARG_FILE" =~ ^.*\.sql$ ]] && ARG_FILE+='.sql'
        [[ ! "$ARG_FILE" =~ ^sql/.*$ ]] && ARG_FILE='sql/'$ARG_FILE
        PG_TESTS_LIST="$ARG_FILE"

    elif [[ "$ARG_LIST" != "" ]]; then
		__parse_file_as_a_list $ARG_LIST

    else
        local search_str='not-found'
        for type in $ARG_TYPE
        do
            if [[ "$type" == 'all' ]]; then
                search_str=''
            else
                search_str="$type"
            fi

            PG_TESTS_LIST="$PG_TESTS_LIST $(find sql -name "${search_str}*.sql" -type f | sort | tr '\n' ' ')"

            # Если режим Воронеж или Орёл, необходимо отфильтровать список
            if [[ $(astra-modeswitch get) -le 1 ]]; then
                local ALL_TESTS=$PG_TESTS_LIST
                PG_TESTS_LIST=''
                __parse_file_as_a_list "./support/tests_voronezh_orel.lst"

                local filtered=$(printf '%s\n' $ALL_TESTS | grep -xF -f <(printf '%s\n' $PG_TESTS_LIST))
                PG_TESTS_LIST=$(echo "$filtered" | paste -sd ' ')

                echo "* Список тестов был изменён из-за режима \"Воронеж\" или \"Орёл\""
            fi

            if [[ "$type" == 'all' ]]; then
                break
            fi
        done
    fi

    PG_TESTS_LIST="$(echo $PG_TESTS_LIST | xargs)"
}


function __start_clusters() {
    local count=0
    local wait_pids=""
    echo "+ Запуск кластеров"

    for cluster in $PG_CLUSTERS
    do
        # Если указан запуск только конкретных кластеров, проверяем, нужно ли запускать $cluster
        [[ $ARG_CLUSTERS != "" && $ARG_CLUSTERS != *"$cluster"* ]] && continue

        # Запуск кластера обычным способом или под реверсивным отладчиком
        if [[ $ARG_RR != *"$cluster"* ]]; then
            pg_ctlcluster $PG_VERSION $cluster start &
            wait_pids+="$! "
        else
            # Логи запуска будут храниться по пути "/root/rr_$cluster.log"
            local data_dir="$PG_DATA_DIR"
            local config="--config_file=$PG_CONFIG_DIR/$cluster/postgresql.conf"
            [[ "$cluster" == "$PG_FILES_CLUSTER" ]] && data_dir=""

            # TODO: think about syncing 'ac_audit_destination' with 'logging_collector' instead of only 'ac_audit_mode'
            rr -M execaps -c 0x400 -- perl support/postgres_with_caps.pl $PG_VERSION \
                -p ${PG_PORTS_LIST[count]} -D "$data_dir/$cluster" $config \
                --ac_audit_mode=none --logging_collector=off --log_statement=all --log_min_messages=DEBUG5 \
                &> "/root/rr_$cluster.log" &
        fi

        (( count++ ))
    done

    [[ $wait_pids != "" ]] && wait $wait_pids
    [[ $ARG_RR != "" ]] && sleep $(($TIME_QUANT*3))

    if [[ $ARG_LOGS_TO_RES_DIR -eq 1 ]]; then
        for cluster in $PG_CLUSTERS; do
            local log_file="$PG_DATA_DIR/$cluster/pg_log/"*.log
            ln -s $log_file "$PG_TESTS_RES_DIR/logs/$cluster.log"
        done
    fi
}


function prepare_test_environment() {
    echo "Подготовка к проведению тестирования PostgreSQL:"

    # Устанавливаем перехват сигналов, чтобы всё точно очищалось при любом завершении скрипта
    trap clean_everything EXIT SIGINT

    ulimit -c unlimited

    if ! __user_exists "u_0_00"; then
        echo "+ Создание тестовых пользователей"

        if ! ./support/create-users &> /dev/null; then
            echo -e "${COLOR_B_RED}Ошибка создания пользователей!"
            echo -e "Выполните \"sudo bash $PG_WORKDIR/support/create-users\" отдельно, чтобы увидеть её.${COLOR_OFF}"
            exit 1
        fi
    fi

    __remove_clusters
    __create_clusters

    __create_tablespaces &
    __configure_clusters &

    __additional_environment_configurations &

    wait

    __set_which_tests_to_run

    __start_clusters

    pg_lsclusters
    if echo "$(pg_lsclusters)" | egrep -e "$PG_SETEST_CLUSTER.* down " \
                                       -e "$PG_SEFOREIGN_CLUSTER.* down " \
                                       -e "$PG_FILES_CLUSTER.* down " \
                                 > /dev/null; then
        echo -e "${COLOR_B_RED}Не все тестовые кластеры успешно запустились!${COLOR_OFF}"
        exit 1
    fi

    echo "Подготовка к проведению тестирования завершена."
}



function run_tests() {
    export PGPORT="$PG_SETEST_PORT"  # кластер для подключения по умолчанию

    local psql="/usr/lib/postgresql/$PG_VERSION/bin/psql"

    local tests_failed=0
    local tests_started=0
    local tests_all=0

    let tests_all=$(printf "$PG_TESTS_LIST" | grep -o ' ' | wc -l)+1

    printf "\n+--------------------------------------------------------------------------------------------------------\n"
    echo "| Тестирование функциональных возможностей PostgreSQL (кол-во тестов: $tests_all)"
    echo "| Результаты выполнения будут сохранены в \"$PG_TESTS_RES_DIR\""
    echo "+--------------------------------------------------------------------------------------------------------"

    tests_all=0
    local total_begin=$(__timestamp)

    for test_file in $PG_TESTS_LIST; do
        local test_name=$(basename "$test_file" .sql)
        local test_res_path="$PG_TESTS_RES_DIR/results/$test_name.out"
        local test_expect_path="expected/filtered/$test_name.out"

        printf '%-4s' "$((tests_all+1))) "  # номер теста
        printf '%-27s' "$test_name"         # имя теста

        # Заголовок тестов начинается с: "-- RBT-TEST"
        local test_header=$(grep -m 1 'RBT-TEST' "$test_file" | cut -c 15-)

        if [[ "$test_header" == "" ]]; then
            echo "Пропущен (не найден заголовок)"
        else
            while [ "${#test_header}" -lt 50 ]; do
                test_header="$test_header."
            done;
            echo -n "$test_header "          # описание теста

            # Принудительно очищаем кластеры
            $psql -q -U postgres -f "support/clear-clusters.sql" template1 &> /dev/null

            # # # Выполняем тест
            local script_begin=$(__timestamp) && \
            local script_begin_time=$(date +"%T.%3N")

            # Выполняем от execaps для исключения случайных привилегий роли
            execaps -c 0x14 -- $psql -a -q -U postgres --variable=time_quant=$TIME_QUANT -f "$test_file" template1 &> $test_res_path \
            || \
            execaps -c 0x14 -- $psql -a -q -U postgres --variable=time_quant=$TIME_QUANT -f "$test_file.ext" template1 &>> $test_res_path

            local script_end=$(__timestamp)
            let script_time=$script_end-$script_begin
            # # #

            # Обрабатываем вывод и ожидаемые файлы
            # Если индивидуальный шаблон для теста есть - сначала применяем его, потом - общие шаблоны
            if [[ -f "support/sed_filters/$test_name" ]]; then
                sed -i -f "support/sed_filters/$test_name" "$test_res_path"
            fi
            if [[ -f "support/sed_filters/$test_name.e" ]]; then
                sed -E -i -f "support/sed_filters/$test_name.e" "$test_res_path"
            fi
            sed -i -f "support/sed_filters/all_results" -f "support/sed_filters/all_results_and_expected" "$test_res_path"
            sed -f "support/sed_filters/all_results_and_expected" expected/$test_name.out > "$test_expect_path"

            # Ищем различия
            if [ "$test_name" == "misc-audit" ]; then n_rows=6; else n_rows=3; fi
            diff_res=$(diff -C${n_rows} -N -Z "$test_expect_path" "$test_res_path")

            if [[ $? -eq 0 ]]; then
                echo -e "${COLOR_CYAN}успех${COLOR_OFF}    ($(__display_time $script_time))"
            else
                local display_time=$(__display_time $script_time)
                echo -e "${COLOR_B_RED}!ОШИБКА!${COLOR_OFF} ($(printf '%-12s' "$display_time") | $script_begin_time)"
                ((tests_failed++))

                # Сохраняем различия
                echo "$diff_res" >> $PG_TESTS_RES_DIR/se-test.diffs
            fi

            ((tests_started++))
        fi

        ((tests_all++))
    done

    local total_end=$(__timestamp)
    let total_time=$total_end-$total_begin

    echo "+--------------------------------------------------------------------------------------------------------"
    echo "| Запущено = $tests_started, ошибочных = $tests_failed, общее время выполнения = $(__display_time $total_time)"
    echo "+--------------------------------------------------------------------------------------------------------"

    return $tests_failed
}



function run_pgbench() {
    # Кластер для подключения по умолчанию
    # (на самом деле PG_FILES_PORT, но т.к. кластер создаётся только один, а
    #  порты у них не заданы жестко, то получается вот так)
    export PGPORT="$PG_SETEST_PORT"

    local psql="/usr/lib/postgresql/$PG_VERSION/bin/psql"
    local DB_NAME="mtest"

    echo -e "${COLOR_CYAN}\n+ Инициализация данных для теста${COLOR_OFF}"

    $psql -U postgres -f './support/init.sql' template1 >& /dev/null
    $psql -U postgres -f './support/init-mac.sql' $DB_NAME >& /dev/null
    pgbench -U postgres --initialize --scale=$ARG_SCALE $DB_NAME
    $psql -U postgres -c "GRANT ALL ON pgbench_accounts TO public;" $DB_NAME >& /dev/null
    $psql -U postgres -c "GRANT ALL ON pgbench_branches TO public;" $DB_NAME >& /dev/null
    $psql -U postgres -c "GRANT ALL ON pgbench_history TO public;" $DB_NAME >& /dev/null
    $psql -U postgres -c "GRANT ALL ON pgbench_tellers TO public;" $DB_NAME >& /dev/null

    echo -e "${COLOR_CYAN}+ Запуск нагрузочного тестирования${COLOR_OFF}"
    echo -e "${COLOR_CYAN}Аргументы: -U $ARG_USERNAME --scale=$ARG_SCALE --client=$ARG_CLIENTS --jobs=$ARG_JOBS $ARG_UNKNOWN${COLOR_OFF}"

    local PERF=$([[ $ARG_PERF -eq 1 ]] && echo 'perf record -F 99 -a -g --' || echo '')
    time $PERF pgbench -U $ARG_USERNAME --client=$ARG_CLIENTS --jobs=$ARG_JOBS $ARG_UNKNOWN $DB_NAME

    if [[ $ARG_PERF -eq 1 ]]; then
        echo -e "${COLOR_CYAN}+ Создание flamegraph-а${COLOR_OFF}"

        perf script > /tmp/astra_pgbench.perf
        /usr/local/src/FlameGraph/stackcollapse-perf.pl /tmp/astra_pgbench.perf > /tmp/astra_pgbench.folded.perf
        /usr/local/src/FlameGraph/flamegraph.pl /tmp/astra_pgbench.folded.perf > "$PG_TESTS_RES_DIR"/flamegraph.svg
    fi
}



function parse_cmd() {
    # Основные команды скрипта
    case "$ARG_CMD" in
    "main")
        prepare_test_environment
        run_tests
        tests_result=$?

        printf "\nКонец тестирования: $(date +"%T")"
        clean_everything $tests_result
    ;;
    "prepare_environment"|"pe")
        prepare_test_environment
        trap - EXIT SIGINT
        exit 0
    ;;
    "clean_environment"|"ce")
        ARG_CMD='clean'
        clean_everything 0
    ;;
    "create_users"|"cu")
        bash ./support/create-users
    ;;
    "remove_users"|"ru")
        bash ./support/remove-users
    ;;
    "pgbench"|"pb")
        if [[ $ARG_INSTALL_DEPS -eq 1 ]]; then
            echo -e "${COLOR_CYAN}* Установка зависимостей для нагрузочного тестирования${COLOR_OFF}"
            set -x
            apt-get install -y linux-tools-$(uname -r) git
            git clone --depth 1 https://github.com/brendangregg/FlameGraph.git /usr/local/src/FlameGraph
            chmod -R 777 /usr/local/src/FlameGraph
            exit 0
        fi

        prepare_test_environment
        run_pgbench
        clean_everything $?
    ;;
    *)
        echo "Нераспознанная команда '$ARG_CMD'!"
    esac
}



function main() {
    export PG_VERSION=$(/usr/share/postgresql-common/supported-versions | tail -n 1)

    # # Обработка аргументов скрипта
    # совместимость с одним из старых аргументов
    if [[ "${@:1:1}" = "-all" ]]; then
        set -- "${@:2}"
    fi
    local args=$(python3 support/arg_parse.py "$@")
    if [[ "$args" != "ARG_"* ]]; then
        echo "$args"  # вызов справки или ошибка
        exit 0
    fi
#    printf "$args\n\n"

    # Выставление полученных ранее переменных в данном процессе bash
    # (двойные кавычки важны, чтобы символы новых строк не заменялись пробелами)
    eval "$args"

    # Выставление переменных или их печать
    case "$ARG_CMD" in
    "variables"|"vars")
        set_variables print
        exit 0
    ;;
    *)
        set_variables
    esac

    # Проверка на наличие прав суперпользователя
    if [[ "$UID" -ne "0" ]]; then
        echo -e "${COLOR_B_RED}Требуются права суперпользователя для запуска скрипта.${COLOR_OFF}"
        exit 1
    fi

    not_exist() { echo "Директории '$PG_WORKDIR' не существует!" && exit 1; }
    pushd "$PG_WORKDIR" &> /dev/null || not_exist

    parse_cmd
}


main "$@"
