MongoDB automation, of a kind.

cluster.sh

#!/bin/bash
#
# Scroll down for juicy configuration and code.
#
# This work is licensed under the MIT license, which allows for commercial re-use,
# does not require back-contributions, isn't viral, and requires attribution.
# That's the plain english version of the full license text, below.  Additionally,
# the following std_disclaimer applies:
#
#   I do not accept responsibility for any effects, adverse or otherwise, 
#   that this code may have on you, your computer, your sanity, your dog, 
#   and anything else that you can think of. Use it at your own risk.
#
# The MIT License (MIT)
# 
# Copyright (c) 2014-2015 Alice Bevan-McGregor
# 
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# 
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
# 
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
# 

# Note: in production, some of these things would need to be tuned:
#  --smallfiles --objcheck --oplogSize --chunkSize
# These would also be farmed out to separate machines, no localhost.

# Set this to an empty value to use the default mongos port.
PORT=37017

# Base path for MongoDB file storage.
BASE=/Users/amcgregor/Projects/exocortex/var/data
RUN=/Users/amcgregor/Projects/exocortex/var/run
LOG=/Users/amcgregor/Projects/exocortex/var/log

# Basic configuration options.
# Do not set --oplogSize in production.
COMMON="--bind_ip 127.0.0.1 --fork --keyFile keyfile"
SHARD_COMMON="${COMMON} --objcheck --directoryperdb --smallfiles --shardsvr --oplogSize 200"
CONFIG_COMMON="${COMMON} --configsvr --smallfiles"
ROUTER_COMMON="${COMMON} --chunkSize 1"

# Admin password.  Pro tip: use a better one.
ADMIN="xyzzy"

mkdir -p $BASE
cd $BASE


function start_shard() {
	local rset=$1
	local replica=$2
	local port=$((10000+$rset*1000+$replica))
	
	mongod ${SHARD_COMMON} \
		--dbpath "./set-${rset}/replica-${replica}" \
		--port "${port}" \
		--replSet "rs${rset}" \
		--pidfilepath "${RUN}/set-${rset}-replica-${replica}.pid" \
		--logpath "${LOG}/set-${rset}-replica-${replica}.log"
}

function stop_shard() {
	kill $(cat "${RUN}/set-${1}-replica-${2}.pid")
}

function start_config() {
	local config=$1
	local port=$[10100 + $config]
	
	mongod ${CONFIG_COMMON} \
		--dbpath "./config-${config}" \
		--port "${port}" \
		--pidfilepath "${RUN}/config-${config}.pid" \
		--logpath "${LOG}/config-${config}.log"
}

function stop_config() {
	kill $(cat "${RUN}/config-${1}.pid")
}

function start_router() {
	: ${1:=27017}
	
	# TODO: Make --configdb variable.
	mongos ${ROUTER_COMMON} \
		--configdb localhost:10101,localhost:10102,localhost:10103 \
		--port ${1} \
		--pidfilepath "${RUN}/router.pid" \
		--logpath "${LOG}/router.log"
}

function stop_router() {
	kill $(cat "${RUN}/router.pid")
}

function add_user() {
	echo "db.createUser({user: 'admin', pwd: '', roles: ['readWriteAnyDatabase', 'userAdminAnyDatabase', 'dbAdminAnyDatabase', 'clusterAdmin']});"
	echo "db.createUser({user: 'admin', pwd: '${ADMIN}', roles: ['readWriteAnyDatabase', 'userAdminAnyDatabase', 'dbAdminAnyDatabase', 'clusterAdmin']});" | mongo "localhost:${1}/admin"
}

function configure() {
	# TODO: Make the number of members variable.
	echo "db.runCommand({replSetInitiate: {_id: 'rs${1}', members: [{_id: 1, host: 'localhost:1${1}001'}, {_id: 2, host: 'localhost:1${1}002'}, {_id: 3, host: 'localhost:1${1}003'}]}})"
	echo "db.runCommand({replSetInitiate: {_id: 'rs${1}', members: [{_id: 1, host: 'localhost:1${1}001'}, {_id: 2, host: 'localhost:1${1}002'}, {_id: 3, host: 'localhost:1${1}003'}]}})" | mongo "localhost:1${1}001/admin"
}

function construct() {
	echo "First-time startup.  This may take a while."
	
	mkdir -p set-{1,2}/replica-{1,2,3} config-{1,2,3}
	
	echo -e "\nGenerating key."
	openssl rand -base64 741 > keyfile
	chmod 600 keyfile
	
	for r in 1 2; do
		echo -e "\nConstructing replica set: rs${r}"
		
		for i in 1 2 3; do start_shard $r $i; done
		configure $r
	done
	
	echo -e "\nConstructing config servers."
	
	for i in 1 2 3; do start_config $i; done
	
	echo -e "\nSleeping a bit to let the replication settle."
	sleep 30  # Pro: block on the logs waiting for notification of election of a PRIMARY and SECONDARY announcements.
	
	# We assume the first host won the election.
	for i in 1 2; do
		echo "rs.status()" | mongo "localhost:1${i}001/admin"
		add_user "1${i}001"
	done
	
	echo -e "\nConstructing sharding router."
	
	start_router ${PORT}
	
	for r in 1 2; do
		# TODO: Make the shard members variable.
		echo "db.runCommand( { addshard : 'rs${r}/localhost:1${r}001,localhost:1${r}002,localhost:1${r}003' } )"
		echo "db.runCommand( { addshard : 'rs${r}/localhost:1${r}001,localhost:1${r}002,localhost:1${r}003' } )" | mongo "localhost:${PORT}/admin"
	done
	
	echo "sh.status()" | mongo "localhost:${PORT}/admin"
	
	add_user ${PORT}
	
	echo -e "\nDone.\nRemember to run db.runCommand({enablesharding: 'somedb'}) against your db of choice."
}

function start() {
	# Handle initial construction of the set, if data is not found.
	[[ ! -e "set-1" ]] && construct && exit
	
	echo "Starting up."
	
	for i in 1 2 3; do
		for r in 1 2; do start_shard $r $i; done
		start_config $i
	done
	
	start_router ${PORT}
	
	echo "Done."
}

function stop() {
	echo "Shutting down."
	
	stop_router
	
	for i in 1 2 3; do
		stop_config $i
		stop_shard 2 $i
		stop_shard 1 $i
	done
	
	echo "Done."
}


case "$1" in
	start)
		start
		;;
	stop)
		stop
		;;
	restart)
		stop
		start
		;;
	*)
		echo "Usage: $0 {start|stop|restart}"
		exit 1
esac