mongos

#!/bin/bash

set -e           # exit on error
set -o pipefail  # trace ERR through pipes
set -o errtrace  # trace ERR through 'time command' and other functions
set -o errexit   # exit the script if any statement returns a non-true return value

function cleanup {
    echo "Cleaning up..."
    if [ -z ${TMPDIR+x} ] ; then
        echo -n
    else
        rm -rf $TMPDIR
        unset TMPDIR
    fi
    unset SOURCE_CREDENTIALS
    unset DEST_CREDENTIALS
}

function error {
    local parent_lineno="$1"
    local message="$2"
    local code="${3:-1}"
    if [[ -n "$message" ]] ; then
        echo "Error on or near line ${parent_lineno}: ${message}; exiting with status ${code}"
    else
        echo "Error on or near line ${parent_lineno}; exiting with status ${code}"
    fi
    cleanup
    exit "${code}"
}
trap 'error ${LINENO}' ERR

function usage_error {
    echo mongos: "$1"
    echo "usage: mongos push|pull|sync"
    echo ""
    exit
}

function config_not_found {
    echo "failed: '$1' not found, it needs to be in the same dir"
    echo "aborting..."
    echo ""
    exit
}

function parse_yaml {
   local prefix=$2
   local s='[[:space:]]*' w='[a-zA-Z0-9_]*' fs=$(echo @|tr @ '\034')
   sed -ne "s|^\($s\):|\1|" \
        -e "s|^\($s\)\($w\)$s:$s[\"']\(.*\)[\"']$s\$|\1$fs\2$fs\3|p" \
        -e "s|^\($s\)\($w\)$s:$s\(.*\)$s\$|\1$fs\2$fs\3|p"  "$1" |
   awk -F$fs '{
      indent = length($1)/2;
      vname[indent] = $2;
      for (i in vname) {if (i > indent) {delete vname[i]}}
      if (length($3) > 0) {
         vn=""; for (i=0; i<indent; i++) {vn=(vn)(vname[i])("_")}
         printf("%s%s%s=\"%s\"\n", "'$prefix'",vn, $2, $3);
      }
   }'
}

function get_script_dir {
    pushd . > /dev/null
    local SCRIPT_PATH="${BASH_SOURCE[0]}"

    if ([ -h "${SCRIPT_PATH}" ]) then
      while ([ -h "${SCRIPT_PATH}" ]) do cd `dirname "$SCRIPT_PATH"`; SCRIPT_PATH=`readlink "${SCRIPT_PATH}"`; done
    fi

    cd `dirname ${SCRIPT_PATH}` > /dev/null
    local SCRIPT_PATH=`pwd`;
    popd  > /dev/null

    echo $SCRIPT_PATH
}

function get_confirmation() {
    read -p "Are you sure you want to $1? Enter 'yes': " mongo_confr

    case $mongo_confr in
        [yY][Ee][Ss] )  ;;
        *) echo "Incorrect input, aborting..."; exit;;
    esac
}

function load_configs {
    echo "load_configs"
    if [ -z "$CONFIG_FILE" ] ; then
        CONFIG_FILE="config.yml"
    fi
    local config=$CONFIG_FILE
    echo "Parsing '$config'..."

    DIR=$(get_script_dir)
    local FILE="$DIR/$config"

    if [ -f "$FILE" ]; then
       eval $(parse_yaml "$FILE")
       success_msg
    else
       config_not_found $config
    fi

    # Loads:
    # - $source_db
    # - $source_host
    # - $source_username
    # - $source_password
    # - $source_authentication_db
    # - $dest_db
    # - $dest_uhost
    # - $dest_username
    # - $dest_password
    # - $dest_authentication_db

    SOURCE_CREDENTIALS=""
    if [[ ! -z $source_username ]] ; then
        SOURCE_CREDENTIALS="-u $source_username -p $source_password"
    fi

    DEST_CREDENTIALS=""
    if [[ ! -z $dest_username ]] ; then
        DEST_CREDENTIALS="-u $dest_username -p $dest_password"
    fi

    TMPDIR=/tmp/mongos/"$source_db"
}

function banner {
    echo mongos:
    echo -----------
}

function success_msg {
    echo "Success!"
    echo
}

function done_msg {
    echo "Done!"
    echo
}


function pull {
    banner
    if [ "$1" == false ] ; then
        get_confirmation 'pull'
    fi
    load_configs

    echo "Dumping Source DB to $TMPDIR... "
    mongodump \
        --ssl \
        --gzip \
        --db "$source_db" \
        --host "$source_host" \
        --username "$source_username" \
        --password "$source_password" \
        --authenticationDatabase "$source_authentication_db" \
        --out "$TMPDIR" > /dev/null

    success_msg
    done_msg
}

function push {
    banner
    if [ "$1" == false ] ; then
        get_confirmation 'push'
    fi
    load_configs

    echo "Overwriting Dest DB with Dump... "
    mongorestore \
        --ssl \
        --gzip \
        --drop \
        --db "$dest_db" \
        --host "$dest_host" \
        --username "$dest_username" \
        --password "$dest_password" \
        --authenticationDatabase "$dest_authentication_db" \
        --dir "$TMPDIR"/"$source_db" > /dev/null

    success_msg
    done_msg
}

function sync {
    banner
    if [ "$1" == false ] ; then
        get_confirmation 'sync'
    fi

    pull
    push

    cleanup
    success_msg
    done_msg
}


## MAIN
## ====

# TODO: Add --overwrite flag

if [[ $# -eq 0 ]] ; then
    usage_error "no arguments provided"
else
    action="$1"
    skip=false
    shift

    while [ "${1+defined}" ]; do
        if [ "$1" == '-c' ] || [ "$1" == '--config' ] ; then
            shift
            if [ -e "$1" ] ; then
                CONFIG_FILE="$1"
            else
                usage_error "No config file named \"$1\""
            fi
        elif [ "$1" == '-y' ] || [ "$1" == "--yes" ]  ; then
            skip=true
        else
            usage_error "unrecognized option \"$1\""
        fi
        shift
    done

    if [[ "$action" == 'push' ]] ; then
        push $skip
    elif [[ "$action" == 'pull' ]] ; then
        pull $skip
    elif [[ "$action" == 'sync' ]] ; then
        sync $skip
    elif [[ "$action" == 'clean' ]] ; then
        banner
        if [ $skip == false ] ; then
            get_confirmation 'clean'
        fi
        load_configs

        cleanup $skip

        success_msg
        done_msg
    else
        usage_error "unrecognized command \"$action\""
    fi
fi