partitioning/postgresql/docker/run_test_env.sh

#!/bin/bash

# Default values
PG_VERSION=""
ZABBIX_VERSION=""

# Color codes
GREEN='\033[0;32m'
RED='\033[0;31m'
NC='\033[0m' # No Color

usage() {
    echo "Usage: $0 --pg <16|17|18> --zabbix <7.0|7.4> [--rds] [--rds-drop]"
    echo "Example: $0 --pg 16 --zabbix 7.0 [--rds-drop]"
    exit 1
}

# Parse arguments
USE_RDS=false
DROP_RDS=false
while [[ "$#" -gt 0 ]]; do
    case $1 in
        --pg) PG_VERSION="$2"; shift ;;
        --zabbix) ZABBIX_VERSION="$2"; shift ;;
        --rds) USE_RDS=true ;;
        --rds-drop) USE_RDS=true; DROP_RDS=true ;;
        *) echo "Unknown parameter: $1"; usage ;;
    esac
    shift
done

if [[ -z "$PG_VERSION" || -z "$ZABBIX_VERSION" ]]; then
    echo -e "${RED}Error: detailed arguments required.${NC}"
    usage
fi

# Map Zabbix version to sql-scripts folder
if [[ "$ZABBIX_VERSION" == "7.0" ]]; then
    SQL_DIR="../sql-scripts-70"
elif [[ "$ZABBIX_VERSION" == "7.4" ]]; then
    SQL_DIR="../sql-scripts-74"
else
    echo -e "${RED}Error: Unsupported Zabbix version. Use 7.0 or 7.4.${NC}"
    exit 1
fi

echo -e "${GREEN}Preparing environment for PostgreSQL $PG_VERSION and Zabbix $ZABBIX_VERSION...${NC}"

# Cleanup previous run
echo "Cleaning up containers and volumes..."
docker compose down -v > /dev/null 2>&1
rm -rf init_scripts
mkdir -p init_scripts

# Symlink SQL scripts
echo "Setting up initialization scripts from $SQL_DIR..."

# 0. Extra Users
if [[ -f "../init_extra_users.sql" ]]; then
    cp "../init_extra_users.sql" ./init_scripts/00_init_extra_users.sql
    echo "Copied extra user init script."
fi

# 1. Schema
if [[ -f "$SQL_DIR/schema.sql" ]]; then
    # Use 01_00 to ensure it comes before 01_10
    cp "$SQL_DIR/schema.sql" ./init_scripts/01_00_schema.sql
    
    # 1.1 Partitioning Infrastructure
    if [[ -f "../procedures/00_schema_create.sql" ]]; then
       cp "../procedures/00_schema_create.sql" ./init_scripts/01_10_schema_create.sql
    fi
     if [[ -f "../procedures/01_maintenance.sql" ]]; then
       cp "../procedures/01_maintenance.sql" ./init_scripts/01_30_maintenance.sql
     fi
     if [[ -f "../procedures/02_enable_partitioning.sql" ]]; then
       cp "../procedures/02_enable_partitioning.sql" ./init_scripts/01_40_enable.sql
     fi
     if [[ -f "../procedures/03_monitoring_view.sql" ]]; then
       cp "../procedures/03_monitoring_view.sql" ./init_scripts/01_50_monitoring.sql
     fi
else
    echo -e "${RED}Error: schema.sql not found in $SQL_DIR${NC}"
    exit 1
fi

# 2. Images
if [[ -f "$SQL_DIR/images.sql" ]]; then
    cp "$SQL_DIR/images.sql" ./init_scripts/02_images.sql
else
    echo -e "${RED}Error: images.sql not found in $SQL_DIR${NC}"
    exit 1
fi

# 3. Data
if [[ -f "$SQL_DIR/data.sql" ]]; then
    cp "$SQL_DIR/data.sql" ./init_scripts/03_data.sql
else
    echo -e "${RED}Error: data.sql not found in $SQL_DIR${NC}"
    exit 1
fi

# 4. Mock History Data
if [[ -f "../z_gen_history_data.sql" ]]; then
    cp "../z_gen_history_data.sql" ./init_scripts/04_gen_data.sql
    echo "Copied mock data generator."
else
     echo -e "${RED}Warning: z_gen_history_data.sql not found!${NC}"
fi

# Check logic for 7.4 vs 7.0 (file names might slightly differ or be organized differently if using packages, 
# but assuming source layout provided)

# Export variable for Docker Compose
export PG_VERSION=$PG_VERSION

if [ "$USE_RDS" = "true" ]; then
    echo -e "${GREEN}Deploying directly to RDS environment...${NC}"
    if [ ! -f "../db_credentials" ]; then
        echo -e "${RED}Error: ../db_credentials file not found. Please create it first.${NC}"
        exit 1
    fi
    
    # Initialize RDS (create/drop user and db)
    if [ "$DROP_RDS" = "true" ]; then
        echo "Initializing Zabbix RDS user and database (with DROP requested)..."
        bash ../setup_rds.sh --drop
    else
        echo "Initializing Zabbix RDS user and database..."
        bash ../setup_rds.sh
    fi
    
    source ../db_credentials
    export PGPASSWORD="$ZBX_DB_PASSWORD"
    
    echo "Applying scripts from init_scripts/ to RDS..."
    for sql_file in $(ls ./init_scripts/*.sql | sort); do
        echo "Executing $sql_file..."
        psql "host=$DB_HOST port=$DB_PORT dbname=$ZBX_DB_NAME user=$ZBX_DB_USER sslmode=$DB_SSL_MODE sslrootcert=../$DB_SSL_ROOT_CERT" -f "$sql_file" -v ON_ERROR_STOP=1
    done
    
    echo -e "${GREEN}RDS Environment ready.${NC}"
    echo "Connect: psql \"host=$DB_HOST port=$DB_PORT dbname=$ZBX_DB_NAME user=$ZBX_DB_USER sslmode=$DB_SSL_MODE sslrootcert=../$DB_SSL_ROOT_CERT\""
else
    # Run Docker Compose
    echo -e "${GREEN}Starting PostgreSQL container...${NC}"
    docker compose up -d

    echo -e "${GREEN}Waiting for database to be ready...${NC}"
    # Simple wait loop
    for i in {1..30}; do
        if docker exec zabbix-db-test pg_isready -U zabbix > /dev/null 2>&1; then
            echo -e "${GREEN}Database is ready!${NC}"
            break
        fi
        echo -n "."
        sleep 1
    done

    # Check if data generation finished
    echo "To follow initialization logs, run: docker logs -f zabbix-db-test"
    echo -e "${GREEN}Environment ready.${NC}"
    echo "Connect: psql -h localhost -p 5432 -U zabbix -d zabbix"
fi