#!/usr/bin/env bash

set -euo pipefail

# TODO(cuongdo): add support for multiple stores on the same node
DATA_DIR="${DATA_DIR:-/mnt/data0}"

usage() {
    echo "usage: $0 [upload|download] google_cloud_storage_url"
    echo
    echo "google_cloud_storage_url is a gs:// URL to a directory where archived"
    echo "stores will be stored to and retrieved from."
    exit 1
}

# Backup all regular files in the data directory.
backup_data_dir() {
    local backup_dir
    backup_dir="${DATA_DIR}/backup-$(date +%Y%m%d-%H%M%S)"
    echo "Backing up files in ${DATA_DIR} to ${backup_dir}"
    mkdir "${backup_dir}"
    find "${DATA_DIR}" -maxdepth 1 -type f -exec mv {} "${backup_dir}" \;
}

if [ $# -ne 2 ]; then
    echo "expected 2 parameters"
    echo ""
    usage
fi

# ~/.config/gcloud may be erroneously owned by root, so fix that.
GCLOUD_DIR="${HOME}/.config/gcloud"
if [[ -d "${GCLOUD_DIR}" ]]; then
    sudo chown -R "${USER}" "${GCLOUD_DIR}"
fi

# CockroachDB nodes are named blah-cockroach-[0-9]*, so extract the final
# numeric part.
gcs_url="$2"
hostname="$(hostname)"
node_index=${NODE_INDEX-${hostname##*-}}
store_url="${gcs_url}/store${node_index}.tgz"

action=$1
case $action in
upload)
    echo "Uploading store ${DATA_DIR} to ${store_url}"
    cd "${DATA_DIR}"
    tar zcvf - * --exclude='logs' | gsutil cp - ${store_url}
    ;;
download)
    if [[ ! -d "${DATA_DIR}" ]]; then
        echo "${DATA_DIR} doesn't exist"
        exit 1
    fi
    if [[ ! -z $(compgen -G "${DATA_DIR}/*.sst" 2>/dev/null) ]]; then
        backup_data_dir
    fi
    echo "Downloading ${store_url} to store ${DATA_DIR}"
    cd "${DATA_DIR}"
    gsutil cat "${store_url}" | tar zxf -
    ;;
*)
    echo "Unknown action '${action}'"
    echo
    usage
esac
