Implement MariaDB connection pooling with ProxySQL for high availability

#!/usr/bin/env bash
set -euo pipefail

# ProxySQL + MariaDB Connection Pooling Installation Script
# Production-quality script for high availability database setup

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m'

# Configuration variables
PROXYSQL_VERSION="2.5.5"
ADMIN_PASSWORD="${PROXYSQL_ADMIN_PASSWORD:-$(openssl rand -base64 12)}"
PRIMARY_DB="${PRIMARY_DB_IP:-}"
SECONDARY_DB="${SECONDARY_DB_IP:-}"

# Functions
log_info() { echo -e "${GREEN}[INFO]${NC} $1"; }
log_warn() { echo -e "${YELLOW}[WARN]${NC} $1"; }
log_error() { echo -e "${RED}[ERROR]${NC} $1"; }

cleanup() {
    log_error "Installation failed. Cleaning up..."
    systemctl stop proxysql 2>/dev/null || true
    if [[ "$PKG_MGR" == "apt" ]]; then
        dpkg --remove proxysql 2>/dev/null || true
    else
        rpm -e proxysql 2>/dev/null || true
    fi
    exit 1
}

show_usage() {
    cat << EOF
Usage: $0 [OPTIONS]
Install ProxySQL with MariaDB connection pooling for high availability

Options:
    -p, --primary-db IP      Primary MariaDB server IP
    -s, --secondary-db IP    Secondary MariaDB server IP
    -h, --help              Show this help message

Environment Variables:
    PROXYSQL_ADMIN_PASSWORD  Custom admin password (auto-generated if not set)

Example:
    $0 -p 10.0.1.10 -s 10.0.1.11
EOF
}

check_prerequisites() {
    if [[ $EUID -ne 0 ]]; then
        log_error "This script must be run as root"
        exit 1
    fi

    if ! command -v systemctl >/dev/null 2>&1; then
        log_error "systemd is required for this installation"
        exit 1
    fi

    if ! command -v openssl >/dev/null 2>&1; then
        log_error "openssl is required for password generation"
        exit 1
    fi
}

detect_distro() {
    if [ -f /etc/os-release ]; then
        . /etc/os-release
        case "$ID" in
            ubuntu|debian) 
                PKG_MGR="apt"
                PKG_INSTALL="apt install -y"
                PKG_UPDATE="apt update && apt upgrade -y"
                ;;
            almalinux|rocky|centos|rhel|ol|fedora) 
                PKG_MGR="dnf"
                PKG_INSTALL="dnf install -y"
                PKG_UPDATE="dnf update -y"
                ;;
            amzn) 
                PKG_MGR="yum"
                PKG_INSTALL="yum install -y"
                PKG_UPDATE="yum update -y"
                ;;
            *) 
                log_error "Unsupported distribution: $ID"
                exit 1
                ;;
        esac
    else
        log_error "Cannot detect operating system"
        exit 1
    fi
    log_info "Detected OS: $PRETTY_NAME"
}

install_proxysql() {
    log_info "[2/8] Installing ProxySQL..."
    
    cd /tmp
    if [[ "$PKG_MGR" == "apt" ]]; then
        local package="proxysql_${PROXYSQL_VERSION}-ubuntu22_amd64.deb"
        wget -q "https://github.com/sysown/proxysql/releases/download/v${PROXYSQL_VERSION}/${package}"
        dpkg -i "$package" || apt install -f -y
        rm -f "$package"
    else
        local package="proxysql-${PROXYSQL_VERSION}-1-centos8.x86_64.rpm"
        wget -q "https://github.com/sysown/proxysql/releases/download/v${PROXYSQL_VERSION}/${package}"
        $PKG_INSTALL "$package"
        rm -f "$package"
    fi
}

configure_firewall() {
    log_info "[4/8] Configuring firewall..."
    
    if command -v firewall-cmd >/dev/null 2>&1; then
        firewall-cmd --permanent --add-port=6032/tcp --add-port=6033/tcp
        firewall-cmd --reload
    elif command -v ufw >/dev/null 2>&1; then
        ufw allow 6032/tcp
        ufw allow 6033/tcp
    fi
}

setup_proxysql_config() {
    log_info "[5/8] Configuring ProxySQL..."
    
    # Create ProxySQL configuration directory if needed
    mkdir -p /etc/proxysql
    chown proxysql:proxysql /etc/proxysql
    chmod 755 /etc/proxysql
    
    # Start ProxySQL to initialize database
    systemctl enable proxysql
    systemctl start proxysql
    
    # Wait for ProxySQL to be ready
    sleep 5
    
    # Create admin configuration script
    cat > /tmp/proxysql_config.sql << EOF
-- Update admin credentials
UPDATE global_variables SET variable_value='admin:${ADMIN_PASSWORD}' WHERE variable_name='admin-admin_credentials';
LOAD ADMIN VARIABLES TO RUNTIME;
SAVE ADMIN VARIABLES TO DISK;

-- Configure connection pooling settings
UPDATE global_variables SET variable_value='250' WHERE variable_name='mysql-max_connections';
UPDATE global_variables SET variable_value='10' WHERE variable_name='mysql-default_max_connections';
UPDATE global_variables SET variable_value='1000' WHERE variable_name='mysql-default_max_latency_ms';
UPDATE global_variables SET variable_value='true' WHERE variable_name='mysql-monitor_enabled';

LOAD MYSQL VARIABLES TO RUNTIME;
SAVE MYSQL VARIABLES TO DISK;
EOF

    # Add backend servers if provided
    if [[ -n "$PRIMARY_DB" ]]; then
        cat >> /tmp/proxysql_config.sql << EOF

-- Add MariaDB backend servers
INSERT INTO mysql_servers(hostgroup_id, hostname, port, weight, comment) VALUES
(0, '${PRIMARY_DB}', 3306, 1000, 'Primary MariaDB Server');
EOF
        
        if [[ -n "$SECONDARY_DB" ]]; then
            cat >> /tmp/proxysql_config.sql << EOF
INSERT INTO mysql_servers(hostgroup_id, hostname, port, weight, comment) VALUES
(0, '${SECONDARY_DB}', 3306, 900, 'Secondary MariaDB Server');
EOF
        fi
        
        cat >> /tmp/proxysql_config.sql << EOF

LOAD MYSQL SERVERS TO RUNTIME;
SAVE MYSQL SERVERS TO DISK;

-- Configure query routing rules
INSERT INTO mysql_query_rules(rule_id, active, match_pattern, destination_hostgroup, apply, comment) VALUES
(1, 1, '^SELECT.*', 1, 1, 'Route SELECT to read hostgroup'),
(2, 1, '^INSERT|^UPDATE|^DELETE|^CREATE|^ALTER|^DROP', 0, 1, 'Route writes to primary hostgroup');

LOAD MYSQL QUERY RULES TO RUNTIME;
SAVE MYSQL QUERY RULES TO DISK;
EOF
    fi
    
    # Apply configuration
    mysql -u admin -padmin -h 127.0.0.1 -P6032 < /tmp/proxysql_config.sql
    rm -f /tmp/proxysql_config.sql
}

create_credentials_file() {
    log_info "[6/8] Creating credentials file..."
    
    cat > /etc/proxysql/credentials << EOF
# ProxySQL Admin Credentials
# Admin Interface: mysql -u admin -p${ADMIN_PASSWORD} -h 127.0.0.1 -P6032
# Application Interface: mysql -u <username> -p<password> -h 127.0.0.1 -P6033

PROXYSQL_ADMIN_PASSWORD="${ADMIN_PASSWORD}"
PROXYSQL_ADMIN_PORT="6032"
PROXYSQL_MYSQL_PORT="6033"
EOF

    chown root:proxysql /etc/proxysql/credentials
    chmod 640 /etc/proxysql/credentials
}

verify_installation() {
    log_info "[7/8] Verifying installation..."
    
    # Check ProxySQL service status
    if ! systemctl is-active --quiet proxysql; then
        log_error "ProxySQL service is not running"
        return 1
    fi
    
    # Check if ProxySQL is listening on required ports
    if ! ss -tulpn | grep -q ":6032.*proxysql"; then
        log_error "ProxySQL admin interface not listening on port 6032"
        return 1
    fi
    
    if ! ss -tulpn | grep -q ":6033.*proxysql"; then
        log_error "ProxySQL MySQL interface not listening on port 6033"
        return 1
    fi
    
    # Test admin connection
    if ! mysql -u admin -p"${ADMIN_PASSWORD}" -h 127.0.0.1 -P6032 -e "SELECT 1;" >/dev/null 2>&1; then
        log_error "Cannot connect to ProxySQL admin interface"
        return 1
    fi
    
    log_info "ProxySQL installation verified successfully"
}

show_summary() {
    log_info "[8/8] Installation completed successfully!"
    
    cat << EOF

${GREEN}ProxySQL Configuration Summary:${NC}
- Admin Interface: 127.0.0.1:6032
- MySQL Interface: 127.0.0.1:6033
- Admin Password: ${ADMIN_PASSWORD}
- Credentials File: /etc/proxysql/credentials

${YELLOW}Next Steps:${NC}
1. Configure your MariaDB backend servers if not done already
2. Create application database users in MariaDB
3. Add users to ProxySQL configuration:
   mysql -u admin -p${ADMIN_PASSWORD} -h 127.0.0.1 -P6032
4. Update your applications to connect via ProxySQL (port 6033)

${YELLOW}Useful Commands:${NC}
- View ProxySQL logs: journalctl -u proxysql -f
- ProxySQL admin: mysql -u admin -p${ADMIN_PASSWORD} -h 127.0.0.1 -P6032
- Check status: systemctl status proxysql
EOF
}

# Parse command line arguments
while [[ $# -gt 0 ]]; do
    case $1 in
        -p|--primary-db)
            PRIMARY_DB="$2"
            shift 2
            ;;
        -s|--secondary-db)
            SECONDARY_DB="$2"
            shift 2
            ;;
        -h|--help)
            show_usage
            exit 0
            ;;
        *)
            log_error "Unknown option: $1"
            show_usage
            exit 1
            ;;
    esac
done

# Set up error handling
trap cleanup ERR

# Main installation process
log_info "Starting ProxySQL + MariaDB connection pooling installation..."

check_prerequisites
detect_distro

log_info "[1/8] Updating system packages..."
eval "$PKG_UPDATE"

install_proxysql

log_info "[3/8] Installing MariaDB client..."
$PKG_INSTALL mariadb-client 2>/dev/null || $PKG_INSTALL mariadb || true

configure_firewall
setup_proxysql_config
create_credentials_file
verify_installation
show_summary

log_info "ProxySQL installation completed successfully!"