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_LIB_DIR={{platform_lib_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)
        if [ $# -ne 2 ]; then
            echo "Usage: $SCRIPT join <node>"
            exit 1
        fi

        # 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)
        if [ $# -ne 1 ]; then
            echo "Usage: $SCRIPT leave"
            exit 1
        fi

        # 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)
        echo "The 'remove' command no longer exists. If you want a node to"
        echo "safely leave the cluster (handoff its data before exiting),"
        echo "then execute 'leave' on the desired node. If a node is down and"
        echo "unrecoverable (and therefore cannot be safely removed), then"
        echo "use the 'force-remove' command. A force removal drops all data"
        echo "owned by the removed node. Read-repair can be used to restore"
        echo "lost replicas."
        exit 1
        ;;

    force-remove)
        if [ $# -ne 2 ]; then
            echo "Usage: $SCRIPT force-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 $@
        ;;

    down)
        if [ $# -ne 2 ]; then
            echo "Usage: $SCRIPT down <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 down $@
        ;;

    status)
        if [ $# -ne 1 ]; then
            echo "Usage: $SCRIPT status"
            exit 1
        fi

        # 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 $@
        ;;

    vnode-status)
        if [ $# -ne 1 ]; then
            echo "Usage: $SCRIPT vnode-status"
            exit 1
        fi

        # 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 vnode_status $@
        ;;

    ringready)
        if [ $# -ne 1 ]; then
            echo "Usage: $SCRIPT ringready"
            exit 1
        fi

        # 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)
        if [ $# -ne 1 ]; then
            echo "Usage: $SCRIPT transfers"
            exit 1
        fi

        # 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 $@
        ;;

    member_status)
        if [ $# -ne 1 ]; then
            echo "Usage: $SCRIPT member_status"
            exit 1
        fi

        # 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_core_console member_status $@
        ;;

    ring_status)
        if [ $# -ne 1 ]; then
            echo "Usage: $SCRIPT ring_status"
            exit 1
        fi

        # 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_core_console ring_status $@
        ;;

    cluster_info)
        if [ $# -lt 2 ]; then
            echo "Usage: $SCRIPT cluster_info <output_file> ['local' | <node> ['local' | <node>] [...]]"
            exit 1
        fi

        # 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_infinity 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

        shift #optional names come after 'js_reload'
        $NODETOOL rpc riak_kv_js_manager reload $@
        ;;

    erl_reload)
        # Reload user Erlang code
        RES=`$NODETOOL ping`
        if [ "$RES" != "pong" ]; then
            echo "Node is not running!"
            exit 1
        fi

        $NODETOOL rpc riak_kv_console reload_code
        ;;

    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 \
            -pa $RUNNER_LIB_DIR/basho-patches \
            -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 \
                       -pa $RUNNER_LIB_DIR/basho-patches \
                       -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 \
                       -pa $RUNNER_LIB_DIR/basho-patches \
                       -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 \
                       -pa $RUNNER_LIB_DIR/basho-patches \
                       -eval "case catch(riak:client_test(\"$NODE_NAME\")) of \
                               ok -> init:stop();                             \
                               _  -> init:stop(1)                             \
                              end."

        ;;
    diag)
        # Riaknostic location
        RIAKNOSTIC_LOC="$ERTS_PATH/../../lib/riaknostic/riaknostic"

        # Riaknostic user
        RIAKNOSTIC_USER=$RUNNER_USER
        if [ -z $RUNNER_USER ]; then
            RIAKNOSTIC_USER=$LOGNAME
        fi

        # Setup command to run riaknostic
        RIAKNOSTIC="$ERTS_PATH/escript $RIAKNOSTIC_LOC --user $RIAKNOSTIC_USER \
            --etc $RUNNER_ETC_DIR --base $RUNNER_BASE_DIR"

        # URL for Riaknostic download instructions
        RIAKNOSTIC_URL="http://riaknostic.basho.com/"

        shift

        if [ -f "$RIAKNOSTIC_LOC" ]; then
            $RIAKNOSTIC "$@"
        else
            echo "Riak diagnostics utility is not present!"
            echo "Visit $RIAKNOSTIC_URL for instructions."
            exit 1
        fi

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

        MYPID=$$
        NODE_NAME=${NAME_ARG#* }
        $ERTS_PATH/erl -noshell -noinput \
            -pa $RUNNER_LIB_DIR/basho-patches \
            -hidden $NAME_PARAM riak_etop$MYPID$NAME_HOST $COOKIE_ARG \
            -s etop -s erlang halt -output text \
            -node $NODE_NAME \
            $* -tracing off
        ;;
    *)
        echo "Usage: $SCRIPT { join | leave | backup | restore | test | status | "
        echo "                    reip | js_reload | erl_reload | wait-for-service | "
        echo "                    ringready | transfers | force-remove | down | "
        echo "                    cluster_info | member_status | ring_status | vnode-status |"
        echo "                    diag |"
        echo "                    top [-interval N] [-sort reductions|memory|msg_q} [-lines N]"
        exit 1
        ;;
esac