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

#!/usr/bin/env bash

set -euo pipefail

# PostgreSQL 17 Streaming Replication with PgBouncer Setup Script
# Production-ready installation and configuration

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

# Configuration variables
readonly SCRIPT_NAME=$(basename "$0")
readonly TOTAL_STEPS=12

# Function definitions
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_on_error() {
    log_error "Installation failed. Cleaning up..."
    systemctl stop postgresql* 2>/dev/null || true
    systemctl stop pgbouncer 2>/dev/null || true
}

trap cleanup_on_error ERR

show_usage() {
    cat << EOF
Usage: $SCRIPT_NAME <role> [options]

Roles:
  primary <primary_ip>              - Configure as primary server
  replica <primary_ip> <slot_name>  - Configure as replica server
  pgbouncer <primary_ip> <replica_ip1,replica_ip2,...>  - Configure PgBouncer

Examples:
  $SCRIPT_NAME primary 203.0.113.10
  $SCRIPT_NAME replica 203.0.113.10 replica_slot_1
  $SCRIPT_NAME pgbouncer 203.0.113.10 203.0.113.11,203.0.113.12
EOF
}

# Validate arguments
if [[ $# -lt 2 ]]; then
    show_usage
    exit 1
fi

readonly ROLE="$1"
readonly PRIMARY_IP="$2"
readonly SLOT_NAME="${3:-replica_slot_1}"
readonly REPLICA_IPS="${3:-}"

# Detect OS and package manager
if [[ ! -f /etc/os-release ]]; then
    log_error "/etc/os-release not found. Cannot detect distribution."
    exit 1
fi

. /etc/os-release

case "$ID" in
    ubuntu|debian)
        PKG_MGR="apt"
        PKG_UPDATE="apt update"
        PKG_INSTALL="apt install -y"
        PG_SERVICE="postgresql"
        PG_VERSION="17"
        PG_DATA_DIR="/var/lib/postgresql/17/main"
        PG_CONFIG_DIR="/etc/postgresql/17/main"
        ;;
    almalinux|rocky|centos|rhel|ol)
        PKG_MGR="dnf"
        PKG_UPDATE="dnf update -y"
        PKG_INSTALL="dnf install -y"
        PG_SERVICE="postgresql-17"
        PG_VERSION="17"
        PG_DATA_DIR="/var/lib/pgsql/17/data"
        PG_CONFIG_DIR="/var/lib/pgsql/17/data"
        ;;
    fedora)
        PKG_MGR="dnf"
        PKG_UPDATE="dnf update -y"
        PKG_INSTALL="dnf install -y"
        PG_SERVICE="postgresql"
        PG_VERSION="17"
        PG_DATA_DIR="/var/lib/pgsql/data"
        PG_CONFIG_DIR="/var/lib/pgsql/data"
        ;;
    amzn)
        PKG_MGR="yum"
        PKG_UPDATE="yum update -y"
        PKG_INSTALL="yum install -y"
        PG_SERVICE="postgresql"
        PG_VERSION="17"
        PG_DATA_DIR="/var/lib/pgsql/data"
        PG_CONFIG_DIR="/var/lib/pgsql/data"
        ;;
    *)
        log_error "Unsupported distribution: $ID"
        exit 1
        ;;
esac

# Check prerequisites
if [[ $EUID -ne 0 ]]; then
    log_error "This script must be run as root"
    exit 1
fi

# Step 1: Update system packages
echo "[1/$TOTAL_STEPS] Updating system packages..."
$PKG_UPDATE

# Step 2: Install PostgreSQL 17
echo "[2/$TOTAL_STEPS] Installing PostgreSQL 17..."
if [[ "$ID" == "ubuntu" || "$ID" == "debian" ]]; then
    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-client-17 postgresql-contrib-17
else
    $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 postgresql17
fi

# Step 3: Initialize PostgreSQL (if needed)
echo "[3/$TOTAL_STEPS] Initializing PostgreSQL..."
if [[ "$ID" == "ubuntu" || "$ID" == "debian" ]]; then
    systemctl enable postgresql
    if [[ "$ROLE" == "primary" ]]; then
        systemctl start postgresql
    fi
else
    if [[ "$ROLE" == "primary" ]]; then
        /usr/pgsql-17/bin/postgresql-17-setup initdb
        systemctl enable postgresql-17
        systemctl start postgresql-17
    else
        systemctl enable postgresql-17
    fi
fi

# Configure based on role
case "$ROLE" in
    "primary")
        echo "[4/$TOTAL_STEPS] Configuring primary server..."
        
        # Backup original config
        cp "$PG_CONFIG_DIR/postgresql.conf" "$PG_CONFIG_DIR/postgresql.conf.backup"
        
        # Configure postgresql.conf
        cat >> "$PG_CONFIG_DIR/postgresql.conf" << 'EOF'

# Replication settings
listen_addresses = '*'
wal_level = replica
max_wal_senders = 10
max_replication_slots = 10
wal_keep_size = 1024
archive_mode = on
archive_command = '/bin/true'
checkpoint_completion_target = 0.9
wal_compression = on
EOF
        
        echo "[5/$TOTAL_STEPS] Configuring authentication..."
        cat >> "$PG_CONFIG_DIR/pg_hba.conf" << EOF

# Replication connections
host replication replicator 0.0.0.0/0 md5
host all pgbouncer 0.0.0.0/0 md5
EOF
        
        echo "[6/$TOTAL_STEPS] Creating replication users..."
        systemctl restart $PG_SERVICE
        sleep 5
        
        sudo -u postgres psql -c "CREATE USER replicator REPLICATION LOGIN ENCRYPTED PASSWORD 'SecureReplPassword123!';" || true
        sudo -u postgres psql -c "CREATE USER pgbouncer LOGIN ENCRYPTED PASSWORD 'SecurePgBouncerPass123!';" || true
        
        log_info "Primary server configured successfully"
        ;;
        
    "replica")
        echo "[4/$TOTAL_STEPS] Configuring replica server..."
        
        systemctl stop $PG_SERVICE 2>/dev/null || true
        rm -rf "$PG_DATA_DIR"/*
        
        echo "[5/$TOTAL_STEPS] Creating base backup..."
        sudo -u postgres PGPASSWORD='SecureReplPassword123!' pg_basebackup -h "$PRIMARY_IP" -D "$PG_DATA_DIR" -U replicator -P -v -R -W -C -S "$SLOT_NAME"
        
        echo "[6/$TOTAL_STEPS] Configuring replica settings..."
        cat >> "$PG_CONFIG_DIR/postgresql.conf" << 'EOF'

# Hot standby settings
listen_addresses = '*'
hot_standby = on
max_standby_archive_delay = 30s
max_standby_streaming_delay = 30s
wal_receiver_status_interval = 2s
hot_standby_feedback = on
EOF
        
        chown -R postgres:postgres "$PG_DATA_DIR"
        chmod 750 "$PG_DATA_DIR"
        
        systemctl start $PG_SERVICE
        log_info "Replica server configured successfully"
        ;;
        
    "pgbouncer")
        echo "[4/$TOTAL_STEPS] Installing PgBouncer..."
        $PKG_INSTALL pgbouncer
        
        echo "[5/$TOTAL_STEPS] Configuring PgBouncer..."
        
        cat > /etc/pgbouncer/pgbouncer.ini << EOF
[databases]
* = host=$PRIMARY_IP port=5432

[pgbouncer]
listen_addr = *
listen_port = 6432
auth_type = md5
auth_file = /etc/pgbouncer/userlist.txt
pool_mode = session
max_client_conn = 1000
default_pool_size = 25
reserve_pool_size = 5
log_connections = 1
log_disconnections = 1
EOF
        
        echo "[6/$TOTAL_STEPS] Creating user authentication..."
        echo '"pgbouncer" "md5$(echo -n SecurePgBouncerPass123!pgbouncer | md5sum | cut -d\  -f1)"' > /etc/pgbouncer/userlist.txt
        
        chown pgbouncer:pgbouncer /etc/pgbouncer/userlist.txt
        chmod 640 /etc/pgbouncer/userlist.txt
        
        systemctl enable pgbouncer
        systemctl start pgbouncer
        log_info "PgBouncer configured successfully"
        ;;
        
    *)
        log_error "Invalid role: $ROLE"
        show_usage
        exit 1
        ;;
esac

# Configure firewall
echo "[7/$TOTAL_STEPS] Configuring firewall..."
if command -v firewall-cmd &> /dev/null; then
    firewall-cmd --permanent --add-port=5432/tcp
    if [[ "$ROLE" == "pgbouncer" ]]; then
        firewall-cmd --permanent --add-port=6432/tcp
    fi
    firewall-cmd --reload
elif command -v ufw &> /dev/null; then
    ufw allow 5432/tcp
    if [[ "$ROLE" == "pgbouncer" ]]; then
        ufw allow 6432/tcp
    fi
fi

# Final verification
echo "[8/$TOTAL_STEPS] Verifying installation..."
if [[ "$ROLE" == "pgbouncer" ]]; then
    systemctl is-active --quiet pgbouncer && log_info "PgBouncer is running" || log_error "PgBouncer is not running"
else
    systemctl is-active --quiet $PG_SERVICE && log_info "PostgreSQL is running" || log_error "PostgreSQL is not running"
fi

echo -e "\n${GREEN}Installation completed successfully!${NC}"
echo "Role: $ROLE"
[[ "$ROLE" != "primary" ]] && echo "Primary IP: $PRIMARY_IP"
[[ "$ROLE" == "replica" ]] && echo "Slot name: $SLOT_NAME"

log_warn "Remember to:"
log_warn "- Change default passwords in production"
log_warn "- Restrict IP access in pg_hba.conf"
log_warn "- Configure SSL certificates for secure connections"
log_warn "- Set up monitoring and backup procedures"