riak-admin

#!/bin/bash

RUNNER_SCRIPT_DIR={{runner_script_dir}}
RUNNER_SCRIPT=${0##*/}

RUNNER_BASE_DIR={{runner_base_dir}}
RUNNER_ETC_DIR={{runner_etc_dir}}
RUNNER_LOG_DIR={{runner_log_dir}}
RUNNER_USER={{runner_user}}

# Make sure this script is running as the appropriate user
if [ "$RUNNER_USER" -a "x$LOGNAME" != "x$RUNNER_USER" ]; then
    type -p sudo > /dev/null 2>&1
    if [ $? -ne 0 ]; then
        echo "sudo doesn't appear to be installed and your EUID isn't $RUNNER_USER" 1>&2
        exit 1
    fi
    echo "Attempting to restart script through sudo -u $RUNNER_USER"
    exec sudo -u $RUNNER_USER -i $RUNNER_SCRIPT_DIR/$RUNNER_SCRIPT $@
fi

# Make sure CWD is set to runner base dir
cd $RUNNER_BASE_DIR

# Extract the target node name from node.args
NAME_ARG=`grep '\-[s]*name' $RUNNER_ETC_DIR/vm.args`
if [ -z "$NAME_ARG" ]; then
    echo "vm.args needs to have either -name or -sname parameter."
    exit 1
fi

# Learn how to specify node name for connection from remote nodes
echo "$NAME_ARG" | grep '^-sname' > /dev/null 2>&1
if [ "X$?" = "X0" ]; then
    NAME_PARAM="-sname"
    NAME_HOST=""
else
    NAME_PARAM="-name"
    echo "$NAME_ARG" | grep '@.*' > /dev/null 2>&1
    if [ "X$?" = "X0" ]; then
        NAME_HOST=`echo "${NAME_ARG}" | sed -e 's/.*\(@.*\)$/\1/'`
    else
        NAME_HOST=""
    fi
fi

# Extract the target cookie
COOKIE_ARG=`grep '\-setcookie' $RUNNER_ETC_DIR/vm.args`
if [ -z "$COOKIE_ARG" ]; then
    echo "vm.args needs to have a -setcookie parameter."
    exit 1
fi

# Identify the script name
SCRIPT=`basename $0`

# Parse out release and erts info
START_ERL=`cat $RUNNER_BASE_DIR/releases/start_erl.data`
ERTS_VSN=${START_ERL% *}
APP_VSN=${START_ERL#* }

# Add ERTS bin dir to our path
ERTS_PATH=$RUNNER_BASE_DIR/erts-$ERTS_VSN/bin

# Setup command to control the node
NODETOOL="$ERTS_PATH/escript $ERTS_PATH/nodetool $NAME_ARG $COOKIE_ARG"

# Check the first argument for instructions
case "$1" in
    join)
        # Make sure the local node IS running
        RES=`$NODETOOL ping`
        if [ "$RES" != "pong" ]; then
            echo "Node is not running!"
            exit 1
        fi

        shift

        $NODETOOL rpc riak_kv_console join $@
        ;;

    leave)
        # Make sure the local node is running
        RES=`$NODETOOL ping`
        if [ "$RES" != "pong" ]; then
            echo "Node is not running!"
            exit 1
        fi

        shift
        $NODETOOL rpc riak_kv_console leave $@
        ;;

    remove)
        if [ $# -ne 2 ]; then
            echo "Usage: $SCRIPT remove <node>"
            exit 1
        fi

        RES=`$NODETOOL ping`
        if [ "$RES" != "pong" ]; then
            echo "Node is not running!"
            exit 1
        fi

        shift
        $NODETOOL rpc riak_kv_console remove $@
        ;;

    status)
        # Make sure the local node IS running
        RES=`$NODETOOL ping`
        if [ "$RES" != "pong" ]; then
            echo "Node is not running!"
            exit 1
        fi
        shift

        $NODETOOL rpc riak_kv_console status $@
        ;;

    ringready)
        # Make sure the local node IS running
        RES=`$NODETOOL ping`
        if [ "$RES" != "pong" ]; then
            echo "Node is not running!"
            exit 1
        fi
        shift

        $NODETOOL rpc riak_kv_console ringready $@
        ;;

    transfers)
        # Make sure the local node IS running
        RES=`$NODETOOL ping`
        if [ "$RES" != "pong" ]; then
            echo "Node is not running!"
            exit 1
        fi
        shift

        $NODETOOL rpc riak_kv_console transfers $@
        ;;

    cluster_info)
        # Make sure the local node IS running
        RES=`$NODETOOL ping`
        if [ "$RES" != "pong" ]; then
            echo "Node is not running!"
            exit 1
        fi
        shift

        $NODETOOL rpc riak_kv_console cluster_info $@
        ;;

    services)
        $NODETOOL rpcterms riak_core_node_watcher services ''
        ;;

    wait-for-service)
        SVC=$2
        TARGETNODE=$3
        if [ $# -lt 3 ]; then
            echo "Usage: $SCRIPT wait-for-service <service_name> <target_node>"
            exit 1
        fi

        while (true); do
            # Make sure riak_core_node_watcher is up and running locally before trying to query it
            # to avoid ugly (but harmless) error messages
            NODEWATCHER=`$NODETOOL rpcterms erlang whereis "'riak_core_node_watcher'."`
            if [ "$NODEWATCHER" = "undefined" ]; then
                echo "$SVC is not up: node watcher is not running"
                continue
            fi

            # Get the list of services that are available on the requested noe
            SERVICES=`$NODETOOL rpcterms riak_core_node_watcher services "'${TARGETNODE}'."`
            echo "$SERVICES" | grep "[[,]$SVC[],]" > /dev/null 2>&1
            if [ "X$?" = "X0" ]; then
                echo "$SVC is up"
                exit 0
            else
                echo "$SVC is not up: $SERVICES"
            fi
            sleep 3
        done
        ;;

    js_reload)
        # Reload all Javascript VMs
        RES=`$NODETOOL ping`
        if [ "$RES" != "pong" ]; then
            echo "Node is not running!"
            exit 1
        fi

        $NODETOOL rpc riak_kv_js_manager reload
        ;;

    reip)
        ACTION=$1
        shift
        if [ $# -lt 2 ]; then
            echo "Usage $SCRIPT $ACTION <old_nodename> <new_nodename>"
            exit 1
        fi
        # Make sure the local node IS not running
        RES=`$NODETOOL ping`
        if [ "$RES" = "pong" ]; then
            echo "Node must be down to re-ip."
            exit 1
        fi
        OLDNODE=$1
        NEWNODE=$2
        $ERTS_PATH/erl -noshell \
            -config $RUNNER_ETC_DIR/app.config \
            -eval "riak_kv_console:$ACTION(['$OLDNODE', '$NEWNODE'])" \
            -s init stop
        ;;

    restore)
        ACTION=$1
        shift

        if [ $# -lt 3 ]; then
            echo "Usage: $SCRIPT $ACTION <node> <cookie> <filename>"
            exit 1
        fi

        NODE=$1
        COOKIE=$2
        FILENAME=$3

        $ERTS_PATH/erl -noshell $NAME_PARAM riak_kv_backup$NAME_HOST -setcookie $COOKIE \
                       -eval "riak_kv_backup:$ACTION('$NODE', \"$FILENAME\")" -s init stop
        ;;

    backup)
        ACTION=$1
        shift

        if [ $# -lt 4 ]; then
            echo "Usage: $SCRIPT $ACTION <node> <cookie> <filename> [node|all]"
            exit 1
        fi

        NODE=$1
        COOKIE=$2
        FILENAME=$3
        TYPE=$4

        $ERTS_PATH/erl -noshell $NAME_PARAM riak_kv_backup$NAME_HOST -setcookie $COOKIE \
                       -eval "riak_kv_backup:$ACTION('$NODE', \"$FILENAME\", \"$TYPE\")" -s init stop
        ;;

    test)
        # Make sure the local node IS running
        RES=`$NODETOOL ping`
        if [ "$RES" != "pong" ]; then
            echo "Node is not running!"
            exit 1
        fi

        shift

        # Parse out the node name to pass to the client
        NODE_NAME=${NAME_ARG#* }

        $ERTS_PATH/erl -noshell $NAME_PARAM riak_test$NAME_HOST $COOKIE_ARG \
                       -eval "case catch(riak:client_test(\"$NODE_NAME\")) of \
                               ok -> init:stop();                             \
                               _  -> init:stop(1)                             \
                              end."

        ;;
    *)
        echo "Usage: $SCRIPT { join | leave | backup | restore | test | status | "
        echo "                    reip | js_reload | wait-for-service | ringready | "
        echo "                    transfers | remove | cluster_info }"
        exit 1
        ;;
esac