Configure MongoDB sharding with zone-based data distribution for geographic workloads

#!/usr/bin/env bash

set -euo pipefail

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

# Global variables
MONGODB_VERSION="8.0"
CONFIG_DIR="/etc/mongod"
DATA_DIR="/var/lib/mongodb"
LOG_DIR="/var/log/mongodb"
MONGODB_USER="mongod"
REPLICA_SET_NAME="${1:-rs0}"
SHARD_NAME="${2:-shard01}"

# Usage function
usage() {
    echo "Usage: $0 [replica_set_name] [shard_name]"
    echo "  replica_set_name: Name for the replica set (default: rs0)"
    echo "  shard_name: Name for the shard (default: shard01)"
    echo ""
    echo "Example: $0 configrs01 shard-us-east"
    exit 1
}

# Cleanup function
cleanup() {
    local exit_code=$?
    if [ $exit_code -ne 0 ]; then
        echo -e "${RED}[ERROR] Script failed. Cleaning up...${NC}"
        systemctl stop mongod 2>/dev/null || true
        systemctl disable mongod 2>/dev/null || true
    fi
}

# Set trap for cleanup
trap cleanup ERR EXIT

# Check if running as root or with sudo
check_privileges() {
    if [[ $EUID -ne 0 ]]; then
        echo -e "${RED}[ERROR] This script must be run as root or with sudo${NC}"
        exit 1
    fi
}

# Detect distribution and set package manager
detect_distro() {
    echo -e "${BLUE}[1/8] Detecting distribution...${NC}"
    
    if [ -f /etc/os-release ]; then
        . /etc/os-release
        case "$ID" in
            ubuntu|debian)
                PKG_MGR="apt"
                PKG_UPDATE="apt update"
                PKG_INSTALL="apt install -y"
                DISTRO_FAMILY="debian"
                ;;
            almalinux|rocky|centos|rhel|ol)
                PKG_MGR="dnf"
                PKG_UPDATE="dnf check-update || true"
                PKG_INSTALL="dnf install -y"
                DISTRO_FAMILY="rhel"
                ;;
            fedora)
                PKG_MGR="dnf"
                PKG_UPDATE="dnf check-update || true"
                PKG_INSTALL="dnf install -y"
                DISTRO_FAMILY="rhel"
                ;;
            amzn)
                PKG_MGR="yum"
                PKG_UPDATE="yum check-update || true"
                PKG_INSTALL="yum install -y"
                DISTRO_FAMILY="rhel"
                ;;
            *)
                echo -e "${RED}[ERROR] Unsupported distribution: $ID${NC}"
                exit 1
                ;;
        esac
        echo -e "${GREEN}Detected: $PRETTY_NAME ($DISTRO_FAMILY family)${NC}"
    else
        echo -e "${RED}[ERROR] Cannot detect distribution${NC}"
        exit 1
    fi
}

# Install prerequisites
install_prerequisites() {
    echo -e "${BLUE}[2/8] Installing prerequisites...${NC}"
    
    $PKG_UPDATE
    
    if [ "$DISTRO_FAMILY" = "debian" ]; then
        $PKG_INSTALL curl gnupg lsb-release ca-certificates
    else
        $PKG_INSTALL curl gnupg2 redhat-lsb-core ca-certificates
    fi
    
    echo -e "${GREEN}Prerequisites installed successfully${NC}"
}

# Configure MongoDB repository
configure_mongodb_repo() {
    echo -e "${BLUE}[3/8] Configuring MongoDB repository...${NC}"
    
    if [ "$DISTRO_FAMILY" = "debian" ]; then
        # Add MongoDB GPG key
        curl -fsSL https://www.mongodb.org/static/pgp/server-${MONGODB_VERSION}.asc | gpg --dearmor -o /usr/share/keyrings/mongodb-server-${MONGODB_VERSION}.gpg
        
        # Detect Ubuntu version or use generic
        if [ "$ID" = "ubuntu" ]; then
            CODENAME=$(lsb_release -cs)
            # Use jammy for newer versions if not specifically supported
            case "$CODENAME" in
                focal|jammy|noble) ;;
                *) CODENAME="jammy" ;;
            esac
        else
            # Debian
            CODENAME="bookworm"
        fi
        
        echo "deb [ arch=amd64,arm64 signed-by=/usr/share/keyrings/mongodb-server-${MONGODB_VERSION}.gpg ] https://repo.mongodb.org/apt/ubuntu ${CODENAME}/mongodb-org/${MONGODB_VERSION} multiverse" > /etc/apt/sources.list.d/mongodb-org-${MONGODB_VERSION}.list
        
        $PKG_UPDATE
    else
        # RHEL family
        cat > /etc/yum.repos.d/mongodb-org-${MONGODB_VERSION}.repo << EOF
[mongodb-org-${MONGODB_VERSION}]
name=MongoDB Repository
baseurl=https://repo.mongodb.org/yum/redhat/\$releasever/mongodb-org/${MONGODB_VERSION}/x86_64/
gpgcheck=1
enabled=1
gpgkey=https://www.mongodb.org/static/pgp/server-${MONGODB_VERSION}.asc
EOF
    fi
    
    echo -e "${GREEN}MongoDB repository configured successfully${NC}"
}

# Install MongoDB
install_mongodb() {
    echo -e "${BLUE}[4/8] Installing MongoDB...${NC}"
    
    $PKG_INSTALL mongodb-org
    
    # Prevent automatic upgrades
    if [ "$DISTRO_FAMILY" = "debian" ]; then
        echo "mongodb-org hold" | dpkg --set-selections
        echo "mongodb-org-database hold" | dpkg --set-selections
        echo "mongodb-org-server hold" | dpkg --set-selections
        echo "mongodb-mongosh hold" | dpkg --set-selections
        echo "mongodb-org-mongos hold" | dpkg --set-selections
        echo "mongodb-org-tools hold" | dpkg --set-selections
    else
        # RHEL family - exclude from updates in yum.conf
        if ! grep -q "exclude=mongodb-org" /etc/yum.conf; then
            echo "exclude=mongodb-org,mongodb-org-database,mongodb-org-server,mongodb-mongosh,mongodb-org-mongos,mongodb-org-tools" >> /etc/yum.conf
        fi
    fi
    
    echo -e "${GREEN}MongoDB installed successfully${NC}"
}

# Configure MongoDB
configure_mongodb() {
    echo -e "${BLUE}[5/8] Configuring MongoDB for sharding...${NC}"
    
    # Create directories with proper permissions
    mkdir -p $DATA_DIR $LOG_DIR $CONFIG_DIR
    chown -R $MONGODB_USER:$MONGODB_USER $DATA_DIR $LOG_DIR
    chmod 755 $DATA_DIR $LOG_DIR
    
    # Create MongoDB configuration
    cat > /etc/mongod.conf << EOF
# mongod.conf for sharded cluster

# Storage configuration
storage:
  dbPath: $DATA_DIR
  journal:
    enabled: true
  wiredTiger:
    engineConfig:
      cacheSizeGB: 1
      directoryForIndexes: true

# Logging configuration
systemLog:
  destination: file
  logAppend: true
  path: $LOG_DIR/mongod.log
  logRotate: reopen

# Network configuration
net:
  port: 27017
  bindIp: 127.0.0.1,$(hostname -I | awk '{print $1}')

# Security configuration
security:
  authorization: enabled

# Replication configuration
replication:
  replSetName: $REPLICA_SET_NAME

# Sharding configuration
sharding:
  clusterRole: shardsvr

# Process management
processManagement:
  fork: true
  pidFilePath: /var/run/mongodb/mongod.pid

# Enable free monitoring
setParameter:
  enableLocalhostAuthBypass: false
EOF

    chmod 644 /etc/mongod.conf
    chown root:root /etc/mongod.conf
    
    echo -e "${GREEN}MongoDB configuration created${NC}"
}

# Configure firewall
configure_firewall() {
    echo -e "${BLUE}[6/8] Configuring firewall...${NC}"
    
    if command -v firewall-cmd &> /dev/null; then
        # firewalld (RHEL family)
        systemctl enable firewalld
        systemctl start firewalld
        firewall-cmd --permanent --add-port=27017/tcp
        firewall-cmd --permanent --add-port=27018/tcp
        firewall-cmd --permanent --add-port=27019/tcp
        firewall-cmd --reload
    elif command -v ufw &> /dev/null; then
        # UFW (Ubuntu/Debian)
        ufw --force enable
        ufw allow 27017/tcp
        ufw allow 27018/tcp
        ufw allow 27019/tcp
    else
        echo -e "${YELLOW}[WARNING] No supported firewall found. Please configure firewall manually to allow ports 27017-27019${NC}"
    fi
    
    echo -e "${GREEN}Firewall configured for MongoDB ports${NC}"
}

# Configure SELinux if present
configure_selinux() {
    if command -v getenforce &> /dev/null && [ "$(getenforce)" != "Disabled" ]; then
        echo -e "${BLUE}[7/8] Configuring SELinux for MongoDB...${NC}"
        
        # Install SELinux policy tools if not present
        if [ "$DISTRO_FAMILY" = "rhel" ]; then
            $PKG_INSTALL policycoreutils-python-utils 2>/dev/null || $PKG_INSTALL policycoreutils-python 2>/dev/null || true
        fi
        
        # Set SELinux contexts
        semanage fcontext -a -t mongod_exec_t "/usr/bin/mongod" 2>/dev/null || true
        semanage fcontext -a -t mongod_var_lib_t "$DATA_DIR(/.*)?" 2>/dev/null || true
        semanage fcontext -a -t mongod_log_t "$LOG_DIR(/.*)?" 2>/dev/null || true
        restorecon -R $DATA_DIR $LOG_DIR /usr/bin/mongod 2>/dev/null || true
        
        echo -e "${GREEN}SELinux configured for MongoDB${NC}"
    fi
}

# Start and enable MongoDB service
start_mongodb_service() {
    echo -e "${BLUE}[8/8] Starting MongoDB service...${NC}"
    
    systemctl daemon-reload
    systemctl enable mongod
    systemctl start mongod
    
    # Wait for MongoDB to start
    echo -n "Waiting for MongoDB to start"
    for i in {1..30}; do
        if systemctl is-active mongod &> /dev/null; then
            echo ""
            echo -e "${GREEN}MongoDB service started successfully${NC}"
            return 0
        fi
        echo -n "."
        sleep 2
    done
    
    echo ""
    echo -e "${RED}[ERROR] MongoDB failed to start within 60 seconds${NC}"
    systemctl status mongod
    exit 1
}

# Verify installation
verify_installation() {
    echo -e "${BLUE}Verifying MongoDB installation...${NC}"
    
    # Check service status
    if systemctl is-active mongod &> /dev/null; then
        echo -e "${GREEN}✓ MongoDB service is running${NC}"
    else
        echo -e "${RED}✗ MongoDB service is not running${NC}"
        exit 1
    fi
    
    # Check MongoDB connection
    if mongosh --eval "db.adminCommand('ping')" --quiet &> /dev/null; then
        echo -e "${GREEN}✓ MongoDB is accepting connections${NC}"
    else
        echo -e "${YELLOW}⚠ MongoDB connection test failed (normal for secured instances)${NC}"
    fi
    
    # Check configuration
    if [ -f /etc/mongod.conf ]; then
        echo -e "${GREEN}✓ Configuration file exists${NC}"
    else
        echo -e "${RED}✗ Configuration file missing${NC}"
        exit 1
    fi
    
    echo ""
    echo -e "${GREEN}=== MongoDB Sharding Installation Complete ===${NC}"
    echo -e "${BLUE}Next steps:${NC}"
    echo "1. Initialize replica set: mongosh --eval 'rs.initiate()'"
    echo "2. Create admin user for authentication"
    echo "3. Configure config servers and mongos routers"
    echo "4. Add shards to the cluster"
    echo "5. Configure zones and zone ranges for geographic distribution"
    echo ""
    echo -e "${YELLOW}Configuration file: /etc/mongod.conf${NC}"
    echo -e "${YELLOW}Data directory: $DATA_DIR${NC}"
    echo -e "${YELLOW}Log file: $LOG_DIR/mongod.log${NC}"
}

# Main execution
main() {
    echo -e "${GREEN}=== MongoDB Sharding Installation Script ===${NC}"
    
    # Validate arguments
    if [ "$#" -gt 2 ]; then
        usage
    fi
    
    check_privileges
    detect_distro
    install_prerequisites
    configure_mongodb_repo
    install_mongodb
    configure_mongodb
    configure_firewall
    configure_selinux
    start_mongodb_service
    verify_installation
    
    # Disable cleanup trap on success
    trap - ERR EXIT
}

# Execute main function
main "$@"