Set up PostgreSQL 17 streaming replication with PgBouncer connection pooling and load balancing

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

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

# Usage function
usage() {
    echo "Usage: $0 [primary|replica] [primary_ip] [replica_password] [pgbouncer_password]"
    echo "  primary - Set up primary PostgreSQL server"
    echo "  replica - Set up replica PostgreSQL server"
    echo "  primary_ip - IP address of primary server (required for replica)"
    echo "  replica_password - Password for replication user"
    echo "  pgbouncer_password - Password for pgbouncer user"
    exit 1
}

# Error cleanup
cleanup() {
    echo -e "${RED}[ERROR] Installation failed. Check logs for details.${NC}"
    exit 1
}
trap cleanup ERR

# Check prerequisites
check_prereqs() {
    if [[ $EUID -ne 0 ]]; then
        echo -e "${RED}This script must be run as root${NC}"
        exit 1
    fi
}

# Detect distribution
detect_distro() {
    if [[ ! -f /etc/os-release ]]; then
        echo -e "${RED}Cannot detect distribution${NC}"
        exit 1
    fi
    
    . /etc/os-release
    case "$ID" in
        ubuntu|debian)
            PKG_MGR="apt"
            PKG_UPDATE="apt update && apt upgrade -y"
            PKG_INSTALL="apt install -y"
            PG_SERVICE="postgresql"
            PG_DATA_DIR="/var/lib/postgresql/17/main"
            PG_CONFIG_DIR="/etc/postgresql/17/main"
            PG_BIN_DIR="/usr/lib/postgresql/17/bin"
            ;;
        almalinux|rocky|centos|rhel|ol)
            PKG_MGR="dnf"
            PKG_UPDATE="dnf update -y"
            PKG_INSTALL="dnf install -y"
            PG_SERVICE="postgresql-17"
            PG_DATA_DIR="/var/lib/pgsql/17/data"
            PG_CONFIG_DIR="/var/lib/pgsql/17/data"
            PG_BIN_DIR="/usr/pgsql-17/bin"
            ;;
        fedora)
            PKG_MGR="dnf"
            PKG_UPDATE="dnf update -y"
            PKG_INSTALL="dnf install -y"
            PG_SERVICE="postgresql"
            PG_DATA_DIR="/var/lib/pgsql/data"
            PG_CONFIG_DIR="/var/lib/pgsql/data"
            PG_BIN_DIR="/usr/bin"
            ;;
        amzn)
            PKG_MGR="yum"
            PKG_UPDATE="yum update -y"
            PKG_INSTALL="yum install -y"
            PG_SERVICE="postgresql"
            PG_DATA_DIR="/var/lib/pgsql/data"
            PG_CONFIG_DIR="/var/lib/pgsql/data"
            PG_BIN_DIR="/usr/bin"
            ;;
        *)
            echo -e "${RED}Unsupported distribution: $ID${NC}"
            exit 1
            ;;
    esac
}

# Install PostgreSQL 17
install_postgresql() {
    echo -e "${GREEN}[2/10] Installing PostgreSQL 17...${NC}"
    
    case "$ID" in
        ubuntu|debian)
            wget --quiet -O - https://www.postgresql.org/media/keys/ACCC4CF8.asc | apt-key add -
            echo "deb http://apt.postgresql.org/pub/repos/apt/ $(lsb_release -cs)-pgdg main" > /etc/apt/sources.list.d/pgdg.list
            apt update
            $PKG_INSTALL postgresql-17 postgresql-contrib-17 pgbouncer
            ;;
        almalinux|rocky|centos|rhel|ol)
            $PKG_INSTALL https://download.postgresql.org/pub/repos/yum/reporpms/EL-9-x86_64/pgdg-redhat-repo-latest.noarch.rpm
            $PKG_INSTALL postgresql17-server postgresql17-contrib pgbouncer
            ;;
    esac
}

# Initialize PostgreSQL
init_postgresql() {
    echo -e "${GREEN}[3/10] Initializing PostgreSQL...${NC}"
    
    case "$ID" in
        ubuntu|debian)
            systemctl start postgresql
            systemctl enable postgresql
            ;;
        almalinux|rocky|centos|rhel|ol)
            $PG_BIN_DIR/postgresql-17-setup initdb
            systemctl enable --now postgresql-17
            ;;
    esac
}

# Configure primary server
configure_primary() {
    echo -e "${GREEN}[4/10] Configuring primary server...${NC}"
    
    # Create archive directory
    mkdir -p $PG_DATA_DIR/archive
    chown postgres:postgres $PG_DATA_DIR/archive
    chmod 700 $PG_DATA_DIR/archive
    
    # Configure postgresql.conf
    cat >> $PG_CONFIG_DIR/postgresql.conf << EOF

# Replication configuration
listen_addresses = '*'
max_connections = 200
shared_buffers = 256MB
effective_cache_size = 1GB
work_mem = 4MB
maintenance_work_mem = 64MB
wal_level = replica
max_wal_senders = 10
max_replication_slots = 10
wal_keep_size = 1GB
archive_mode = on
archive_command = 'cp %p $PG_DATA_DIR/archive/%f'
hot_standby = on
wal_receiver_timeout = 60s
wal_sender_timeout = 60s
EOF
    
    # Configure pg_hba.conf
    cat >> $PG_CONFIG_DIR/pg_hba.conf << EOF
host    replication     replica         0.0.0.0/0               scram-sha-256
host    all             pgbouncer       0.0.0.0/0               scram-sha-256
EOF
}

# Create users
create_users() {
    echo -e "${GREEN}[5/10] Creating replication users...${NC}"
    
    systemctl restart $PG_SERVICE
    sleep 5
    
    sudo -u postgres psql -c "CREATE ROLE replica REPLICATION LOGIN PASSWORD '$REPLICA_PASS';" || true
    sudo -u postgres psql -c "CREATE ROLE pgbouncer LOGIN PASSWORD '$PGBOUNCER_PASS';" || true
}

# Configure replica server
configure_replica() {
    echo -e "${GREEN}[4/10] Configuring replica server...${NC}"
    
    systemctl stop $PG_SERVICE
    rm -rf $PG_DATA_DIR/*
    
    # Create base backup
    sudo -u postgres PGPASSWORD=$REPLICA_PASS $PG_BIN_DIR/pg_basebackup -h $PRIMARY_IP -D $PG_DATA_DIR -U replica -v -P -W
    
    chown -R postgres:postgres $PG_DATA_DIR
    chmod 700 $PG_DATA_DIR
    
    # Create standby.signal
    touch $PG_DATA_DIR/standby.signal
    chown postgres:postgres $PG_DATA_DIR/standby.signal
    
    # Configure primary_conninfo
    cat >> $PG_CONFIG_DIR/postgresql.conf << EOF
primary_conninfo = 'host=$PRIMARY_IP port=5432 user=replica password=$REPLICA_PASS'
EOF
}

# Configure PgBouncer
configure_pgbouncer() {
    echo -e "${GREEN}[6/10] Configuring PgBouncer...${NC}"
    
    cat > /etc/pgbouncer/pgbouncer.ini << EOF
[databases]
* = host=127.0.0.1 port=5432

[pgbouncer]
listen_port = 6432
listen_addr = *
auth_type = scram-sha-256
auth_file = /etc/pgbouncer/userlist.txt
logfile = /var/log/pgbouncer/pgbouncer.log
pidfile = /var/run/pgbouncer/pgbouncer.pid
admin_users = pgbouncer
pool_mode = session
server_reset_query = DISCARD ALL
max_client_conn = 1000
default_pool_size = 50
EOF
    
    echo "\"pgbouncer\" \"$PGBOUNCER_PASS\"" > /etc/pgbouncer/userlist.txt
    chown pgbouncer:pgbouncer /etc/pgbouncer/userlist.txt
    chmod 640 /etc/pgbouncer/userlist.txt
    
    systemctl enable pgbouncer
    systemctl start pgbouncer
}

# Configure firewall
configure_firewall() {
    echo -e "${GREEN}[7/10] Configuring firewall...${NC}"
    
    if command -v firewall-cmd &> /dev/null; then
        firewall-cmd --permanent --add-port=5432/tcp
        firewall-cmd --permanent --add-port=6432/tcp
        firewall-cmd --reload
    elif command -v ufw &> /dev/null; then
        ufw allow 5432/tcp
        ufw allow 6432/tcp
    fi
}

# Verify installation
verify_installation() {
    echo -e "${GREEN}[8/10] Verifying installation...${NC}"
    
    if systemctl is-active --quiet $PG_SERVICE; then
        echo -e "${GREEN}PostgreSQL is running${NC}"
    else
        echo -e "${RED}PostgreSQL is not running${NC}"
        exit 1
    fi
    
    if systemctl is-active --quiet pgbouncer; then
        echo -e "${GREEN}PgBouncer is running${NC}"
    else
        echo -e "${YELLOW}PgBouncer is not running${NC}"
    fi
}

# Main function
main() {
    if [[ $# -lt 4 ]]; then
        usage
    fi
    
    MODE=$1
    PRIMARY_IP=${2:-""}
    REPLICA_PASS=$3
    PGBOUNCER_PASS=$4
    
    check_prereqs
    
    echo -e "${GREEN}[1/10] Updating system packages...${NC}"
    detect_distro
    $PKG_UPDATE
    $PKG_INSTALL wget ca-certificates lsb-release
    
    install_postgresql
    
    case "$MODE" in
        "primary")
            init_postgresql
            configure_primary
            create_users
            configure_pgbouncer
            ;;
        "replica")
            if [[ -z "$PRIMARY_IP" ]]; then
                echo -e "${RED}Primary IP is required for replica setup${NC}"
                exit 1
            fi
            init_postgresql
            configure_replica
            systemctl start $PG_SERVICE
            configure_pgbouncer
            ;;
        *)
            usage
            ;;
    esac
    
    configure_firewall
    verify_installation
    
    echo -e "${GREEN}[9/10] PostgreSQL 17 with PgBouncer setup complete!${NC}"
    echo -e "${GREEN}[10/10] PostgreSQL: port 5432, PgBouncer: port 6432${NC}"
}

main "$@"