this script initialises an aiida environment (created with conda)

aiida_env

#!/usr/bin/env bash

# for help
# source script_name -h

# this script initialises an aiida environment (created with conda) by:
#    0. Reading the required variables from a yaml config file (by default .aiida_envs.yaml)
#    1a. `source activate env_name`
#    1b. (optionally) activate required branches of github repos (`git -C path checkout branch`)
#    2. `export PGDATA=/path/to/database`
#       `pg_ctl -l postgres_env_{env_name}.log start`
#    3. `rabbitmq-server -detached` for aiida_core >= v1.0.0
#    4. `export AIIDA_PATH=path/containing/.aiida`
#    5. `reentry scan -r aiida`
#    6. `verdi -p profile_name daemon start`
#    7. `verdi profile setdefault profile_name`
#    8. setup terminal tab completion of verdi sub commands
#    9. (optionally) backup database to AIIDA_PATH/.aiida/backups in the background

# The yaml file should look like this:
#
#    yaml_key:
#      conda_env: env_name
#      pq_server: path/to/database
#      aiida_path: path/containing/.aiida
#      aiida_profile: profile_name
#      git_repos:
#        - path: /path/to/repo1
#          branch: develop
#        - path: /path/to/repo2
#          branch: develop

if [ $_ == $0 ]; then
    echo "this script must be sourced: source $(basename $0)"
    exit
fi

print_help () {
    usage="source $(basename ${BASH_SOURCE[0]}) [-h --help] [-k str] [-p str] [-b]

initialises an aiida environment (created with conda), via a yaml config file

where:
    -h --help  show this help text
    -p --yamlpath set the path to the yaml file (default: .aiida_envs.yaml)
    -k --yamlkey set the key to find in the yaml file (default: default)
    -b --backupdb backup the postgre database, in the background (default: false)

The pyyaml pkg is required, and the yaml file should look like this:

yaml_key:
  conda_env: env_name
  pq_server: path/to/database
  aiida_path: path/containing/.aiida
  aiida_profile: profile_name
  git_repos:
    - path: /path/to/repo1
      branch: develop
"

    echo "$usage"
    # NB: must be run as: source {script_file} {yaml_key} {optional:yaml_path}
    # requires pyyaml to be installed
    #     echo "this script must be called with a variable: source {script_file} {yaml_key} {optional:yaml_path}"
    # backup this database in the background
}

read_args () {
    
    # set defaults
    local env_yaml=".aiida_envs.yaml"
    local env_key="default"
    local backup_db=false

    # read options (see https://stackoverflow.com/questions/192249/how-do-i-parse-command-line-arguments-in-bash)
    POSITIONAL=()
    while [[ $# -gt 0 ]]
    do
    key="$1"

    case $key in
        -k|--yamlkey)
        env_key="$2"
        shift # past argument
        shift # past value
        ;;
        -p|--yamlpath)
        env_yaml="$2"
        shift # past argument
        shift # past value
        ;;
        -b|--backupdb)
        backup_db=true
        shift # past argument
        ;;
        *)    # unknown option
        # POSITIONAL+=("$1") # save it in an array for later
        shift # past argument
        ;;
    esac
    done
    # set -- "${POSITIONAL[@]}" # restore positional parameters

    echo "$env_yaml,$env_key,$backup_db"

}

activate_aiida () {

    if [[ $1 == "-h" ]] || [[ $1 == "--help" ]]; then
        print_help
        return
    fi

    local args=$(read_args "$@")
    local env_yaml
    local env_key
    local backup_db
    IFS=',' read -r yaml_file yaml_key backup_db <<< "$args"

    # NB: if you are using the echo command, be sure to use the -e flag to allow backslash escapes.
    local COLORRED='\033[0;31m'
    local COLORORANGE='\033[0;33m'
    local COLORGREEN='\033[0;32m'
    local COLORNONE='\033[0m' # No Color

    echo -e "${COLORGREEN}- Reading variables from key '${yaml_key}' of ${yaml_file}${COLORNONE}"

    # TODO create temp files in more secure way
    local yfile="__read_yaml_key.py"
    cat <<EOF > $yfile
import sys, os
try:
    import yaml
except ImportError:
    sys.stderr.write("${COLORRED}ERROR: pyyaml not installed (pip install pyyaml) ${COLORNONE}\n")
    sys.exit(1)
fpath = sys.argv[1]
profile = sys.argv[2]
if not os.path.exists(fpath):
    sys.stderr.write("${COLORRED}ERROR: could not find path {}${COLORNONE}\n".format(fpath))
    sys.exit(1)
with open(fpath) as f:
    data = yaml.load(f)
if not data.get(profile, False):
    sys.stderr.write("${COLORRED}ERROR: could not find key '{}' in yaml${COLORNONE}\n".format(profile))
    sys.exit(1)
pdata = data[profile]

for key in ["conda_env", "pq_server", "aiida_path", "aiida_profile"]:
    if key not in pdata:
        sys.stderr.write("${COLORRED}ERROR: could not find required key '{0}/{1}' in yaml${COLORNONE}\n".format(profile,key))
        sys.exit(1)

outstring = "{0},{1},{2},{3}".format(pdata["conda_env"], pdata["pq_server"],
pdata["aiida_path"], pdata["aiida_profile"])

for repo in pdata.get("git_repos", []):
    if "path" in repo and "branch" in repo:
        outstring += ",git -C {0} checkout {1}".format(repo["path"],repo["branch"])

sys.stdout.write(outstring)
EOF
    local output1=$(python $yfile "${yaml_file}" $yaml_key)
    rm -f $yfile

    if [[ -z $output1 ]]; then
        echo -e "${COLORRED}QUITTING PROCESS${COLORNONE}"
        return
    fi

    # comma delimited
    local outarray
    IFS=',' read -ra outarray <<< "$output1"

    local ENV=${outarray[0]}
    local SERVER=${outarray[1]}
    local AIIDAPATH=${outarray[2]}
    local PROFILE=${outarray[3]}

    echo -e "Using inputs for '${yaml_key}':\n\tconda_env = $ENV\n\tpq_server = $SERVER\n\taiida_path = $AIIDAPATH\n\tprofile = $PROFILE"

    # activate conda environment
    if [[ $PATH == *"conda/envs/$ENV/bin"* ]]; then
        echo -e "${COLORORANGE}Conda env '$ENV' already activated.${COLORNONE}"
    else
        echo -e "${COLORGREEN}- Activating Conda environment: '$ENV'${COLORNONE}"
        source deactivate
        source activate $ENV
    fi

    # select correct branches of git repos
    echo -e "${COLORGREEN}- Activating Github Branches:${COLORNONE}"
    local i
    for i in "${outarray[@]:4}"; do
        echo "$i"
        eval "$i" >/dev/null
    done

    # stop any running server
    pkill postgres
    # activate server
    if [[ -z `pg_ctl -D $SERVER status | grep "server is running"` ]]; then
        echo -e "${COLORGREEN}- Activating Postgres server: $SERVER${COLORNONE}"
        # close any other active server (from: https://askubuntu.com/questions/547434/how-to-nicely-stop-all-postgres-processes)
        psql -Xtc 'show data_directory' &>/dev/null && pg_ctl -D $(psql -Xtc 'show data_directory') stop
        pg_ctl -D $SERVER start -l "postgres_env_$ENV.log"
        echo -e "${COLORORANGE}Logging Postgres server to: postgres_env_$ENV.log${COLORNONE}"
    else
        echo -e "${COLORORANGE}Postgres server already running: $SERVER${COLORNONE}"
    fi

    echo -e "${COLORGREEN}- Setting PGDATA='${SERVER}'${COLORNONE}"
    export PGDATA="$SERVER"

    # start rabbitmq (aiida_core >= v1.0.0)
    # RabbitMQ is a message queue application that allows AiiDA to send messages to the daemon
    # it should start automatically (after system reboot), but just in case
    if hash rabbitmq-server 2>/dev/null; then
        # TODO check if its already running. use `rabbitmqctl status`, but what to grep for?
        echo -e "${COLORGREEN}- Ensuring rabbitmq is running${COLORNONE}"
        rabbitmq-server -detached >/dev/null 2>&1
    else
        echo -e "${COLORORANGE}Warning: rabbitmq-server not available (required for aiida >= v1.0.0).${COLORNONE}"
        echo -e "${COLORORANGE}To install: conda install rabbitmq-server${COLORNONE}"
    fi

    # use correct .aiida path
    echo -e "${COLORGREEN}- Setting AIIDA_PATH='${AIIDAPATH}'${COLORNONE}"
    export AIIDA_PATH="${AIIDAPATH}"

    # check profile is available
    if [[ -z `verdi profile list | grep -w $PROFILE` ]]; then
        echo -e "${COLORRED}Profile: $PROFILE, was not found.$COLORNONE"
        echo "available profiles:"
        verdi profile list
        return
    fi

    # ensure plugins are up to date
    echo -e "${COLORGREEN}- Rescanning aiida plugins${COLORNONE}"
    reentry scan -r aiida

    # start aiida daemon
    if [[ -z `verdi -p $PROFILE daemon status | grep "Daemon is running "` ]]; then
        echo -e "${COLORGREEN}- Activating daemon for profile: $PROFILE${COLORNONE}"
        verdi -p $PROFILE daemon start
    else
    #    echo -e "${COLORORANGE}Daemon already running for profile: $PROFILE${COLORNONE}"
        echo -e "${COLORGREEN}- Restarting daemon for profile: $PROFILE${COLORNONE}"
        verdi -p $PROFILE daemon restart
    fi

    echo -e "${COLORGREEN}- Setting default profile: $PROFILE${COLORNONE}"
    # different for v0.12 and v1
    verdi profile setdefault $PROFILE &>/dev/null || verdi profile setdefault verdi $PROFILE

    # setup terminal tab completion of verdi sub commands
    # `verdi completioncommand` only for aiida_core < v1.0.0
    echo -e "${COLORGREEN}- Activating verdi tab completion ${COLORNONE}"
    local compcommand=false
    verdi completioncommand >/dev/null 2>&1 && compcommand=true
    if [ $compcommand == true ]; then
        # we are in aiida <v1
        verdi completioncommand > __run_completioncommand.sh
        source __run_completioncommand.sh
        rm -f __run_completioncommand.sh
    else
        # we are in aiida >=v1
        eval "$(_VERDI_COMPLETE=source verdi)"
    fi

    # backup database

    if [ ${backup_db} == true ] ;then

        echo -e "${COLORGREEN}- Backing up database${COLORNONE}"

        # these are in .aiida/config.json
        # make python file here to keep script contained
        # TODO comma delimited parsing (like for yaml)
        local jfname="__read_config_json.py"
        cat <<EOF > $jfname
import sys, os, json
fpath = sys.argv[1]
profile = sys.argv[2]
if not os.path.exists(fpath):
    sys.stderr.write("${COLORRED}ERROR: could not find {}${COLORNONE}\n".format(fpath))
    sys.exit(1)
with open(fpath) as f:
    data = json.load(f)
if not data.get("profiles",{}).get(profile, False):
    sys.stderr.write("${COLORRED}ERROR: could not find profile; {} in config.json${COLORNONE}\n".format(profile))
    sys.exit(1)
pdata = data["profiles"][profile]

for key in ["AIIDADB_HOST", "AIIDADB_PORT", "AIIDADB_USER", "AIIDADB_NAME"]:
    if key not in pdata:
        sys.stderr.write("${COLORRED}ERROR: could not find required key 'profiles/{0}/{1}' in config.json${COLORNONE}\n".format(profile,key))
        sys.exit(1)

outstring = "{0} {1} {2} {3}".format(pdata["AIIDADB_HOST"], pdata["AIIDADB_PORT"],
pdata["AIIDADB_USER"], pdata["AIIDADB_NAME"])

sys.stdout.write(outstring)
EOF
        local output=($(python $jfname "${AIIDAPATH}/.aiida/config.json" $PROFILE))
        rm -f $jfname

        if [[ ! -z $output ]]; then

            local AIIIDAHOST=${output[0]}
            local AIIDAPORT=${output[1]}
            local AIIDAUSER=${output[2]}
            local AIIDADB_NAME=${output[3]}

            echo "Using profile settings from config.json: host=$AIIIDAHOST, port=$AIIDAPORT, user=$AIIDAUSER, dbname=$AIIDADB_NAME"

            local AIIDALOCALTMPDUMPFILE="${AIIDAPATH}/.aiida/backups/${AIIDADB_NAME}-backup.psql.gz"
            local AIIDALOCALTMPDUMPLOG="${AIIDAPATH}/.aiida/backups/${AIIDADB_NAME}-backup.psql.log"

            if [ ! -d "${AIIDAPATH}/.aiida/backups" ]; then
                mkdir "${AIIDAPATH}/.aiida/backups"
            fi
            if [ -e ${AIIDALOCALTMPDUMPFILE} ]
            then
                mv ${AIIDALOCALTMPDUMPFILE} ${AIIDALOCALTMPDUMPFILE}~
            fi

            # NOTE: password stored in ~/.pgpass, where pg_dump will read it automatically
            local cmd="pg_dump -h $AIIIDAHOST -p $AIIDAPORT -U $AIIDAUSER $AIIDADB_NAME"
            nohup sh -c "$cmd | gzip > $AIIDALOCALTMPDUMPFILE; echo "completed" $(date) || rm $AIIDALOCALTMPDUMPFILE" >> $AIIDALOCALTMPDUMPLOG 2>&1 & disown
            echo -e "${COLORORANGE}Postgre DB Backup running in background (PID=$!) to: $AIIDALOCALTMPDUMPFILE${COLORNONE}"
            echo -e "${COLORORANGE}log file: $AIIDALOCALTMPDUMPLOG${COLORNONE}"
        fi
    fi

    # TODO backup repo? https://aiida-core.readthedocs.io/en/latest/backup/index.html#setup-repository-backup

}

activate_aiida "$@"