backup.mongo.plugin

#!/bin/bash
# =================================================================================================================
# Mongo Backup and Restore Functions:
# - Dynamically loaded as a plug-in
# -----------------------------------------------------------------------------------------------------------------
export serverDataDirectory="/data/db"

function onBackupDatabase(){
  (
    local OPTIND
    local unset flags
    while getopts : FLAG; do
      case $FLAG in
        ? ) flags+="-${OPTARG} ";;
      esac
    done
    shift $((OPTIND-1))

    _databaseSpec=${1}
    _backupFile=${2}

    _hostname=$(getHostname ${_databaseSpec})
    _database=$(getDatabaseName ${_databaseSpec})
    _port=$(getPort ${_databaseSpec})
    _portArg=${_port:+"--port=${_port}"}
    _username=$(getUsername ${_databaseSpec})
    _password=$(getPassword ${_databaseSpec})
    echoGreen "Backing up '${_hostname}${_port:+:${_port}}${_database:+/${_database}}' to '${_backupFile}' ..."

    _authDbArg=${MONGODB_AUTHENTICATION_DATABASE:+"--authenticationDatabase ${MONGODB_AUTHENTICATION_DATABASE}"}
    mongodump -h "${_hostname}" -d "${_database}" ${_authDbArg} ${_portArg} -u "${_username}" -p "${_password}" --quiet --gzip --archive=${_backupFile}
    return ${?}
  )
}

function onRestoreDatabase(){
  (
    local OPTIND
    local unset flags
    while getopts : FLAG; do
      case $FLAG in
        ? ) flags+="-${OPTARG} ";;
      esac
    done
    shift $((OPTIND-1))

    _databaseSpec=${1}
    _fileName=${2}
    _adminPassword=${3}

    _hostname=$(getHostname ${flags} ${_databaseSpec})
    _database=$(getDatabaseName ${_databaseSpec})
    _port=$(getPort ${flags} ${_databaseSpec})
    _portArg=${_port:+"--port=${_port}"}
    _username=$(getUsername ${_databaseSpec})
    _password=$(getPassword ${_databaseSpec})
    echo -e "Restoring '${_fileName}' to '${_hostname}${_port:+:${_port}}${_database:+/${_database}}' ...\n" >&2

    # ToDo:
    # - Add support for restoring to a different database.
    #   The following implementation only supports restoring to a database of the same name, 
    #   unlike the postgres implementation that allows the database to be restored to a database of a different
    #   name for testing.
    #   Ref: https://stackoverflow.com/questions/36321899/mongorestore-to-a-different-database

    _authDbArg=${MONGODB_AUTHENTICATION_DATABASE:+"--authenticationDatabase ${MONGODB_AUTHENTICATION_DATABASE}"}
    mongorestore --drop -h ${_hostname} -d "${_database}" ${_authDbArg} ${_portArg} -u "${_username}" -p "${_password}" --gzip --archive=${_fileName} --nsInclude="*"
    return ${?}
  )
}

function onDatabaseInit() {
  (
    local OPTIND
    local unset flags
    while getopts : FLAG; do
      case $FLAG in
        ? ) flags+="-${OPTARG} ";;
      esac
    done
    shift $((OPTIND-1))

    # Retrieve database details
    _databaseSpec=${1}
    _database=$(getDatabaseName "${_databaseSpec}")
    _username=$(getUsername "${_databaseSpec}")
    _password=$(getPassword "${_databaseSpec}")

    # Check if the database already exists
    result=$(mongosh --eval "db.getMongo().getDBNames().indexOf('$_database')")

    if [ ! -z "$result" ] && [ "$result" -ge 0 ]; then
      echoYellow "Database '$_database' already exists, skipping initialization."
      return 0
    fi

    # Initialize the database by creating the user with the roles
    mongosh "$_database" --quiet --eval "
      db.createUser({
        user: '$_username',
        pwd: '$_password',
        roles: [
          { role: 'dbOwner', db: '$_database' },
          { role: 'readWrite', db: '$_database' },
          { role: 'clusterAdmin', db: 'admin' }
        ]
      });
    " > /dev/null 2>&1 &

    # Check the exit status of the createUser command
    if [ $? -eq 0 ]; then
      echoGreen "Database '$_database' initialized successfully."
      return 0
    else
      echoRed "Failed to initialize database '$_database'."
      return 1
    fi
  )
}

function onStartServer(){
  (
    local OPTIND
    local unset flags
    while getopts : FLAG; do
      case $FLAG in
        ? ) flags+="-${OPTARG} ";;
      esac
    done
    shift $((OPTIND-1))

    _databaseSpec=${1}
    _hostname=$(getHostname ${flags} ${_databaseSpec})

    # Start a local MongoDb instance
    exec mongod --bind_ip $_hostname --quiet >/dev/null 2>&1 &
    
    # Wait for mongodb to be ready
    while ! mongosh --host $_hostname --eval "db.adminCommand('ping')" > /dev/null 2>&1; do
      sleep 1
    done

    # Initialize database if necessary
    onDatabaseInit "${_databaseSpec}"
  )
}

function onStopServer(){
  (
    local OPTIND
    local unset flags
    while getopts : FLAG; do
      case $FLAG in
        ? ) flags+="-${OPTARG} ";;
      esac
    done
    shift $((OPTIND-1))

    _databaseSpec=${1}
    _hostname=$(getHostname ${flags} ${_databaseSpec})

    local startTime=$(date +%s%N)
    local timeout=30

    mongosh --host $_hostname --eval "db.getSiblingDB('admin').shutdownServer()" &
    mongosh_pid=$!

    # Wait for mongosh to finish or timeout
    while kill -0 $mongosh_pid 2>/dev/null; do
      if [ $(($(date +%s) - startTime)) -ge $timeout ]; then
        kill $mongosh_pid
        break
      fi
      sleep 1
    done

    # Ensure mongod is stopped
    pkill -f "mongod --bind_ip $_hostname"
  )
}

function onCleanup(){
  (
    if ! dirIsEmpty ${serverDataDirectory}; then
      # Delete the database files and configuration
      echo -e "Cleaning up ...\n" >&2
      rm -rf ${serverDataDirectory}/*
    else
      echo -e "Already clean ...\n" >&2
    fi
  )
}

function onPingDbServer(){
  (
    local OPTIND
    local unset flags
    while getopts : FLAG; do
      case $FLAG in
        ? ) flags+="-${OPTARG} ";;
      esac
    done
    shift $((OPTIND-1))

    _databaseSpec=${1}
    _hostname=$(getHostname ${flags} ${_databaseSpec})
    
    mongosh --host $_hostname --eval "db.adminCommand('ping')" > /dev/null 2>&1

    local mongoshExitCode=$?
    
    if (( ${mongoshExitCode} != 0 )); then
      return 1
    else
      return 0
    fi
  )
}

function onVerifyBackup(){
  (
    local OPTIND
    local unset flags
    while getopts : FLAG; do
      case $FLAG in
        ? ) flags+="-${OPTARG} ";;
      esac
    done
    shift $((OPTIND-1))

    _databaseSpec=${1}

    _hostname=$(getHostname -l ${_databaseSpec})
    _database=$(getDatabaseName ${_databaseSpec})
    _port=$(getPort -l ${_databaseSpec})
    _portArg=${_port:+"--port ${_port}"}
    _username=$(getUsername ${_databaseSpec})
    _password=$(getPassword ${_databaseSpec})

    _dbAddressArg=${_hostname}${_port:+:${_port}}${_database:+/${_database}}
    _authDbArg=${MONGODB_AUTHENTICATION_DATABASE:+"--authenticationDatabase ${MONGODB_AUTHENTICATION_DATABASE}"}
    collections=$(mongosh ${_dbAddressArg} ${_authDbArg} -u "${_username}" -p "${_password}" --quiet --eval 'var dbs = [];dbs = db.getCollectionNames();for (i in dbs){ print(db.dbs[i]);}';)
    rtnCd=${?}

    # Get the size of the restored database
    if (( ${rtnCd} == 0 )); then
      size=$(getDbSize -l "${_databaseSpec}")
      rtnCd=${?}
    fi

    if (( ${rtnCd} == 0 )); then
      numResults=$(echo "${collections}"| wc -l)
      if [[ ! -z "${collections}" ]] && (( numResults >= 1 )); then
        # All good
        verificationLog="\nThe restored database contained ${numResults} collections, and is ${size} in size."
      else
        # Not so good
        verificationLog="\nNo collections were found in the restored database ${_database}."
        rtnCd="3"
      fi
    fi

    echo ${verificationLog}
    return ${rtnCd}
  )
}

function onGetDbSize(){
  (
    local OPTIND
    local unset flags
    while getopts : FLAG; do
      case $FLAG in
        ? ) flags+="-${OPTARG} ";;
      esac
    done
    shift $((OPTIND-1))

    _databaseSpec=${1}

    _hostname=$(getHostname ${flags} ${_databaseSpec})
    _database=$(getDatabaseName ${_databaseSpec})
    _port=$(getPort ${flags} ${_databaseSpec})
    _portArg=${_port:+"--port ${_port}"}
    _username=$(getUsername ${_databaseSpec})
    _password=$(getPassword ${_databaseSpec})

    _dbAddressArg=${_hostname}${_port:+:${_port}}${_database:+/${_database}}
    _authDbArg=${MONGODB_AUTHENTICATION_DATABASE:+"--authenticationDatabase ${MONGODB_AUTHENTICATION_DATABASE}"}
    size=$(mongosh ${_dbAddressArg} ${_authDbArg} -u "${_username}" -p "${_password}" --quiet --eval 'printjson(db.stats().fsTotalSize)')
    rtnCd=${?}

    echo ${size}
    return ${rtnCd}
  )
}
# =================================================================================================================