set up cockroachdb

setup_db.sh

#!/bin/bash

if [ "$(id -u)" -ne 0 ]; then
    echo -e "\033[1;31mERROR: must be running as root, exit code 1\033[0m" >&2
    exit 1
fi

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

OS_ID=$(grep ^ID= /etc/os-release | cut -d= -f2 | tr -d '"')

if [[ ! -f "$SCRIPT_DIR/cockroach" && ! $(command -v cockroach) ]]; then
    echo -e "\033[1;31mERROR: cockroach not exist in $SCRIPT_DIR dir and \$PATH, exit code 2\033[0m" >&2
    exit 2
fi

if [[ -f "$SCRIPT_DIR/cockroach" ]]; then
    COCKROACH=$SCRIPT_DIR/cockroach
else
    COCKROACH=$(command -v cockroach)
fi

# 0. ASK IF CREATE NEW CLUSTER OR
# ADD NODE TO EXISTING CLUSTER
while true; do
    echo "1) create new cluster cert"
    echo "2) create node cert, setup and start node"
    echo "3) create client cert"
    echo "4) init cluster"
    echo "5) cockroach sql shell"

    read -rp "Choose an option [1-5]: " choice

    case "$choice" in
        1)
            echo "You selected: create new cluster cert"
            CHOICE=1
            break
            ;;
        2)
            echo "You selected: create node cert, setup and start node"
            CHOICE=2
            break
            ;;
        3)
            echo "You selected: create client cert"
            CHOICE=3
            break
            ;;
        4)
            echo "You selected: init cluster"
            CHOICE=4
            break
            ;;
        5)
            echo "You selected: cockroach sql shell"
            CHOICE=5
            break
            ;;
        *)
            echo "Invalid option. Please select 1,2,3,4 or 5."
            ;;
    esac
done

if [[ $CHOICE = 5 ]]; then
    $COCKROACH sql --certs-dir=$SCRIPT_DIR --host=127.0.0.1
fi

if [[ $CHOICE = 4 ]]; then
    $COCKROACH init --certs-dir=$SCRIPT_DIR --host=127.0.0.1
fi

if [[ $CHOICE = 3 ]]; then
    if [[ -f "$SCRIPT_DIR/client.root.crt" || -f "$SCRIPT_DIR/client.root.key" ]]; then
        echo "But client.root.crt or client.root.key exist in $SCRIPT_DIR, exit code 10"
        echo "(fix: delete client.root.crt and client.root.key in $SCRIPT_DIR)"
        exit 10
    fi

    $COCKROACH cert create-client root --certs-dir="$SCRIPT_DIR" --ca-key="$SCRIPT_DIR/ca.key"
fi

if [[ $CHOICE = 2 ]]; then
    if [[ ! -f "$SCRIPT_DIR/ca.crt" ]]; then
        echo "But ca.crt not exist in $SCRIPT_DIR, exit code 4"
        echo "(fix: copy ca.crt and ca.key to $SCRIPT_DIR)"
        exit 4
    fi

    if [[ -f "/etc/systemd/system/cockroachdb.service" ]]; then
        echo "/etc/systemd/system/cockraochdb.service already exist, exit code 9"
        echo "(fix: remove cockroachdb.service)"
        exit 9
    fi
fi

if [[ $CHOICE = 2 ]]; then
    echo -e ""
    echo "This current machine will be the one run the node?"

    select choice2 in "yes" "no"
    do
        case "$choice2" in
            "no")
                echo "This machine must be the one will run node, exit code 5"
                exit 5
                ;;
            "yes")
                echo "Start installing.."
                break
                ;;
            *)
                echo "Invalid option. Please select 1 or 2."
                ;;
        esac
    done
fi

if [[ $CHOICE = 1 ]]; then
    if [[ -f "$SCRIPT_DIR/ca.crt" ]]; then
        echo "But ca.crt already exist, exit code 3"
        echo "(fix: delete ca.crt and ca.key in $SCRIPT_DIR to reset)"
        exit 3
    fi
fi

# 1. SET UP NTP
if [[ $CHOICE = 2 ]]; then
    if [ $OS_ID = 'debian' ]; then
        export DEBIAN_FRONTEND=noninteractive

        apt install -y -qq ntpsec

        sed -i '/^pool /d' /etc/ntpsec/ntp.conf

        for s in time1.google.com time2.google.com time3.google.com time4.google.com; do
            line="server $s iburst"
            grep -qxF "$line" /etc/ntpsec/ntp.conf || echo "$line" >> /etc/ntpsec/ntp.conf
        done

        service ntpsec restart
    fi
fi

# 2. GENERATE CLUSTER CERTIFICATE
if [[ $CHOICE = 1 ]]; then
    echo "-- CREATE ca.crt AND ca.key --"
    chmod +x "$SCRIPT_DIR/cockroach"
    $COCKROACH cert create-ca --certs-dir="$SCRIPT_DIR" --ca-key="$SCRIPT_DIR/ca.key"
    echo "NOTE: backup ca.key in a safe place"
fi

# 3. GENERATE NODE CERTIFICATE
if [[ $CHOICE = 2 ]]; then
    echo "-- CREATE node.crt AND node.key --"
    chmod +x "$SCRIPT_DIR/cockroach"

    if [[ -f "$SCRIPT_DIR/node.crt" || -f "$SCRIPT_DIR/node.key" ]]; then
        echo "node.crt exists, exit 6"
        echo "(fix: delete node.crt)"
        exit 6
    fi

    if [[ -f "$SCRIPT_DIR/node.key" ]]; then
        echo "node.key exists, exit 7"
        echo "(fix: delete node.key)"
        exit 7
    fi

    echo
    echo "Please add domain node will serve from eg: example.com, using comma as separator,"
    echo "for example: 192.168.10.9,load_balancer1,lb2,1.2.3.4,2.3.4.5,10.9.8.7,example.com"
    echo "note: script already detect private IP and loopback and hostname, avoid add dynamic public IP"

    read -r -p "Domain:" lb_input

    lb_output="${lb_input//,/ }"

    ip_v4=$(ip route get 1.1.1.1 | awk '{for(i=1;i<=NF;i++) if($i=="src") print $(i+1)}')
    ip_v6=$(ip -6 route get 2606:4700:4700::1111 | awk '{for(i=1;i<=NF;i++) if($i=="src") print $(i+1)}')

    lb_output="$lb_output $ip_v4 $ip_v6 127.0.0.1 $HOST localhost)"

    $COCKROACH cert create-node --certs-dir="$SCRIPT_DIR" --ca-key="$SCRIPT_DIR/ca.key" $lb_output
fi

# 4. GENERATE SERVICE FILE
if [[ $CHOICE = 2 ]]; then
    advertise_port=26257
    default_addr="${lb_input}:${advertise_port}"

    echo
    echo "Advertise address?"
    read -e -i "$default_addr" -p "--advertise-addr=" advertise_addr

    echo $advertise_addr

    crdb_rdv="crdb.rdv.example.com:26257"

    echo
    echo "Join node? (should use a domain for DDNS, eg: $crdb_rdv)"
    read -r -p "--join=" crdb_rdv

    echo
    echo "Locality (eg: continent=NA,country=US,region=US-EAST-1,provider=AWS,zone=A)"
    read -e -i "continent=NA,country=US,region=US-EAST-1,provider=AWS,zone=A" -p "--cache=" locality

    echo $locality > $SCRIPT_DIR/locality_file

    echo
    echo "Max cache (0.1-1)?"
    read -e -i "0.25" -p "--cache=" cache_perc

    echo
    echo "Max sql memory (0.1-1)?"
    read -e -i "0.25" -p "--cache=" max_sql_memory

    if [ $OS_ID = 'debian' ]; then
        if [[ -f "/etc/systemd/system/cockroachdb.service" ]]; then
            echo "cockroachdb.service already exist, exit code 8"
            exit 8
        fi

        touch /etc/systemd/system/cockroachdb.service
        cat << EOF > /etc/systemd/system/cockroachdb.service
[Unit]
Description=Storkspace Cockroach Database Node
Requires=network.target
[Service]
Type=notify
WorkingDirectory=$SCRIPT_DIR
ExecStart=$SCRIPT_DIR/cockroach start --certs-dir=$SCRIPT_DIR --advertise-addr=$advertise_addr --join=$crdb_rdv --cache=$cache_perc --max-sql-memory=$max_sql_memory --locality-file=$SCRIPT_DIR/locality_file
TimeoutStopSec=300
Restart=always
RestartSec=10
User=root
[Install]
WantedBy=default.target
EOF

        systemctl enable cockroachdb
        systemctl start cockroachdb
    fi
fi