atlas

#!/usr/bin/env bash

ATLAS_HOME=${ATLAS_HOME:-"/opt/atlas"}

STORAGE_BACKEND_CONF="atlas.graph.storage.backend"
HBASE_STORAGE_LOCAL_CONF_ENTRY="atlas.graph.storage.hostname"
SOLR_INDEX_CONF_ENTRY="atlas.graph.index.search.backend"
SOLR_INDEX_MODE_CONF_ENTRY="atlas.graph.index.search.solr.mode"
SOLR_INDEX_LOCAL_STANDALONE_CONF_ENTRY="atlas.graph.index.search.solr.http-urls"
SOLR_INDEX_LOCAL_CLOUD_CONF_ENTRY="atlas.graph.index.search.solr.zookeeper-url"
SOLR_INDEX_ZK_URL="atlas.graph.index.search.solr.zookeeper-url"
ATLAS_MAIN_CLASS="org.apache.atlas.Atlas"

DEFAULT_JVM_HEAP_OPTS="-Xmx1024m"
DEFAULT_JVM_OPTS="-Dlog4j.configuration=atlas-log4j.xml -Djava.net.preferIPv4Stack=true -server"
DEFAULT_SOLR_PORT="9838"
DEFAULT_HBASE_CONF_DIR="${ATLAS_HOME}/hbase/conf"
DEFAULT_SOLR_CONF_DIR="${ATLAS_HOME}/solr/server/solr/configsets/_default/conf"

ATLAS_CONF_DIR=${ATLAS_CONF_DIR:-"${ATLAS_HOME}/conf"}
ATLAS_CONF_FILE=${ATLAS_CONF_FILE:-"atlas-application.properties"}
ATLAS_LOG_DIR=${ATLAS_LOG_DIR:-"${ATLAS_HOME}/logs"}
ATLAS_LOG_FILE=${ATLAS_LOG_FILE:-"application.log"}
ATLAS_DATA_DIR=${ATLAS_DATA_DIR:-"${ATLAS_HOME}/data"}
ATLAS_SERVER_HEAP=${ATLAS_SERVER_HEAP:-$DEFAULT_JVM_HEAP_OPTS}
ATLAS_OPTS=${ATLAS_OPTS:-$DEFAULT_JVM_OPTS}
ATLAS_CLASSPATH="${ATLAS_CONF_DIR}:${ATLAS_HOME}/server/webapp/atlas/WEB-INF/classes:${ATLAS_HOME}/server/webapp/atlas/WEB-INF/lib/*:${ATLAS_HOME}/libext/*"
HBASE_CONF_DIR=${HBASE_CONF_DIR:-$DEFAULT_HBASE_CONF_DIR}
ATLAS_CONF_OPTS="-Datlas.log.dir=$ATLAS_LOG_DIR -Datlas.log.file=$ATLAS_LOG_FILE -Datlas.home=$ATLAS_HOME -Datlas.conf=${ATLAS_CONF_DIR}"
ATLAS_EXPANDED_WEBAPP_DIR=${ATLAS_EXPANDED_WEBAPP_DIR:-"${ATLAS_HOME}/server/webapp"}
ATLAS_APP_METADATA_DIR=${ATLAS_APP_METADATA_DIR:-"${ATLAS_EXPANDED_WEBAPP_DIR}/atlas"}
SOLR_PORT=${SOLR_PORT:-$DEFAULT_SOLR_PORT}
SOLR_CONF_DIR=${SOLR_CONF_DIR:-$DEFAULT_SOLR_CONF_DIR}

MANAGE_LOCAL_HBASE=$(echo "${MANAGE_LOCAL_HBASE:-"False"}" | tr '[:upper:]' '[:lower:]')
MANAGE_EMBEDDED_CASSANDRA=$(echo "${MANAGE_EMBEDDED_CASSANDRA:-"False"}" | tr '[:upper:]' '[:lower:]')
MANAGE_LOCAL_SOLR=$(echo "${MANAGE_LOCAL_SOLR:-"False"}" | tr '[:upper:]' '[:lower:]')
MANAGE_LOCAL_ELASTICSEARCH=$(echo "${MANAGE_LOCAL_ELASTICSEARCH:-"False"}" | tr '[:upper:]' '[:lower:]')

function get_configuration_value {
  local __resultvar=$1
  local configuration_file=$2
  local configuration_key=$3
  eval $__resultvar="'$(grep ^${configuration_key}= "${configuration_file}" | sed s/${configuration_key}=//)'"
}

function is_hbase {
  get_configuration_value "storage_backend" "${ATLAS_CONF_DIR}/${ATLAS_CONF_FILE}" "$STORAGE_BACKEND_CONF"
  if [[ "$STORAGE_BACKEND" == "hbase" || "$STORAGE_BACKEND" == "hbase2" || "$STORAGE_BACKEND" == "" ]]; then
    return 0
  fi
  return 1
}

function is_hbase_local {
  if [[ "$MANAGE_LOCAL_HBASE" == "true" ]]; then
    get_configuration_value "hbase_storage_local" "${ATLAS_CONF_DIR}/${ATLAS_CONF_FILE}" "$HBASE_STORAGE_LOCAL_CONF_ENTRY"
    if  is_hbase && [[ -n "$hbase_storage_local" ]]; then
      return 0
    fi
  fi
  return 1
}

function is_cassandra_local {
  local manage_embedded_cassandra=${MANAGE_EMBEDDED_CASSANDRA}
  if [[ "$manage_embedded_cassandra" == "true" ]]; then
    return 0
  fi
  return 1
}

function is_solr_local {
  if [[ "$MANAGE_LOCAL_SOLR" == "true" ]]; then
    if is_solr_local_cloud_mode || is_solr_local_standalone_mode; then
      return 0
    fi
  fi
  return 1
}

function is_solr_local_cloud_mode {
  local configuration_file="${ATLAS_CONF_DIR}/${ATLAS_CONF_FILE}"
  get_configuration_value "solr_index" "$configuration_file" "$SOLR_INDEX_CONF_ENTRY"
  get_configuration_value "solr_index_mode" "$configuration_file" "$SOLR_INDEX_MODE_CONF_ENTRY"
  get_configuration_value "solr_index_local_cloud" "$configuration_file" "$SOLR_INDEX_LOCAL_CLOUD_CONF_ENTRY"
  if [[ -n "$solr_index" ]] && [[ "$solr_index_mode" == "cloud" ]] && [[ -n "$solr_index_local_cloud" ]]; then
    return 0
  fi
  return 1
}

function is_solr_local_standalone_mode {
  local configuration_file="${ATLAS_CONF_DIR}/${ATLAS_CONF_FILE}"
  get_configuration_value "solr_index" "$configuration_file" "$SOLR_INDEX_CONF_ENTRY"
  get_configuration_value "solr_index_mode" "$configuration_file" "$SOLR_INDEX_MODE_CONF_ENTRY"
  get_configuration_value "solr_index_local_standalone" "$configuration_file" "$SOLR_INDEX_LOCAL_STANDALONE_CONF_ENTRY"
  if [[ -n "$solr_index" ]] && [[ "$solr_index_mode" == "http" ]] && [[ -n "$solr_index_local_standalone" ]]; then
    return 0
  fi
  return 1
}

function is_elasticsearch_local {
  if [[ "$MANAGE_LOCAL_ELASTICSEARCH" == "true" ]]; then
    return 0
  fi
  return 1
}

function configure_hbase {
  if [[ "${HBASE_CONF_DIR}" == "${DEFAULT_HBASE_CONF_DIR}" ]]; then
    hbase_conf_file="hbase-site.xml"
    tmpl_file="${ATLAS_CONF_DIR}/hbase/${hbase_conf_file}.template"
    conf_file="${HBASE_CONF_DIR}/${hbase_conf_file}"

    if [[ -f "$tmpl_file" ]]; then
      printf "Configuring %s to %s" "$tmpl_file" "$conf_file"
      cp "$tmpl_file" "$conf_file"
      sed -i "s \${hbase_home} ${ATLAS_CONF_DIR} g" "$conf_file"
      sed -i "s \${atlas_data} ${ATLAS_DATA_DIR} g" "$conf_file"
      sed -i "s \${url_prefix} file:// g" "$conf_file"
      rm -f "$tmpl_file"
    fi
  fi
}

function configure_cassandra {
  cassandra_conf_file="cassandra.yml"
  tmpl_file="${ATLAS_CONF_DIR}/${cassandra_conf_file}.template"
  conf_file="${ATLAS_CONF_DIR}/${cassandra_conf_file}"

  if [[ -f "$tmpl_file" ]]; then
    printf "Configuring %s to %s" "$tmpl_file" "$conf_file"
    cp "$tmpl_file" "$conf_file"
    sed -i "s \${atlas_home} ${ATLAS_HOME} g" "$conf_file"
    rm -f "$tmpl_file"
  fi
}

function configure_zookeeper {
  zookeeper_conf_file="zoo.cfg"
  tmpl_file="${ATLAS_CONF_DIR}/${zookeeper_conf_file}.template"
  conf_file="${ATLAS_CONF_DIR}/${zookeeper_conf_file}"

  if [[ -f "$tmpl_file" ]]; then
    printf "Configuring %s to %s" "$tmpl_file" "$conf_file"
    cp "$tmpl_file" "$conf_file"
    sed -i "s \${atlas_home} ${ATLAS_HOME} g" "$conf_file"
    rm -f "$tmpl_file"
  fi
}

function run_zookeeper {
  printf "Executing : %s" "${ATLAS_HOME}/zkServer.sh start ${ATLAS_CONF_DIR}/${zookeeper_conf_file}"
  "${ATLAS_HOME}/zk/bin/zkServer.sh" start "${ATLAS_CONF_DIR}/${zookeeper_conf_file}"
}

function run_hbase_action {
  printf "Executing : %s" "${ATLAS_HOME}/hbase/bin/hbase-daemon.sh --config ${HBASE_CONF_DIR} start master\n"
  "${ATLAS_HOME}/hbase/bin/hbase-daemon.sh" --config "${HBASE_CONF_DIR}" start master
}

function run_solr {
  get_configuration_value "solr_zk_url" "${ATLAS_CONF_DIR}/${ATLAS_CONF_FILE}" "$SOLR_INDEX_ZK_URL"
  if [[ -n "${solr_zk_url}" ]]; then
    ZK_URL="-z ${solr_zk_url}"
  fi
  "${ATLAS_HOME}/solr/bin/solr" start -p "${SOLR_PORT}" ${ZK_URL}
}

function create_solr_collection {
  local index=$1
  "${ATLAS_HOME}/solr/bin/solr" create -c "${index}" -d "${SOLR_CONF_DIR}" -shards 1 -replicationFactor 1
}

function start_elasticsearch {
  printf "Executing : %s" "${ATLAS_HOME}/elasticsearch/bin/elasticsearch -d -p ${ATLAS_HOME}/logs/elasticsearch.pid\n"
  "${ATLAS_HOME}/elasticsearch/bin/elasticsearch" -d -p "${ATLAS_HOME}/logs/elasticsearch.pid"
}

if [[ "$DEBUG" == "true" ]]; then
  set -x
fi

if is_hbase; then
    if [[ -f "${HBASE_CONF_DIR}/hbase-site.xml" ]]; then
        ATLAS_CLASSPATH="${ATLAS_CLASSPATH}:${HBASE_CONF_DIR}"
    else
        printf "Could not find hbase-site.xml in %s. Please set env var HBASE_CONF_DIR to the hbase client conf dir\n" "$HBASE_CONF_DIR"
        exit 1
    fi
fi

if is_hbase_local; then
  printf "configured for local hbase\n"
  configure_hbase
  run_hbase_action
  printf "hbase started\n"
fi

if is_solr_local; then
  printf "configured for local solr\n"

  if is_cassandra_local; then
    printf "Cassandra embedded configured\n"
    configure_cassandra
    configure_zookeeper
    run_zookeeper
    printf "zookeeper started\n"
  fi

  run_solr
  printf "solr started\n"

  printf "setting up solr collections...\n"
  create_solr_collection vertex_index
  create_solr_collection edge_index
  create_solr_collection fulltext_index
fi

if is_elasticsearch_local; then
  printf "configured for local elasticsearch\n"
  start_elasticsearch
  printf "elasticsearch started\n"
fi

# Which jar to use
if [ -z "$JAVA_HOME" ]; then
  JAR="jar"
else
  JAR="$JAVA_HOME/bin/jar"
fi

# Which java to use
if [ -z "$JAVA_HOME" ]; then
  JAVA="java"
else
  JAVA="$JAVA_HOME/bin/java"
fi

if [[ ! -d "${ATLAS_APP_METADATA_DIR}/WEB-INF" ]]; then
  mkdir -p "${ATLAS_APP_METADATA_DIR}/WEB-INF"
  pushd "${ATLAS_APP_METADATA_DIR}" || exit
  "$JAR" -xf "${ATLAS_EXPANDED_WEBAPP_DIR}/atlas.war"
  popd || exit
fi

"$JAVA" $ATLAS_CONF_OPTS $ATLAS_SERVER_HEAP $ATLAS_OPTS -classpath "$ATLAS_CLASSPATH" "$ATLAS_MAIN_CLASS" -app "${ATLAS_HOME}/server/webapp/atlas"