#!/bin/bash

set -e

# Use credentials from environment
if [ -n "$GCLOUD_SERVICE_KEY" ]; then
    # Google's client libraries will check for GOOGLE_APPLICATION_CREDENTIALS
    # and use a file in that location for credentials if present;
    # See https://cloud.google.com/docs/authentication/production
    export GOOGLE_APPLICATION_CREDENTIALS="${GOOGLE_APPLICATION_CREDENTIALS:-/tmp/gcp.json}"
    echo "$GCLOUD_SERVICE_KEY" > "$GOOGLE_APPLICATION_CREDENTIALS"
fi

if [ -n "$GOOGLE_APPLICATION_CREDENTIALS" ] && which gcloud > /dev/null; then
    gcloud --quiet auth activate-service-account --key-file "$GOOGLE_APPLICATION_CREDENTIALS"
    python -W ignore -c 'import google.auth; print("project = ", google.auth.default()[1])' >> ~/.config/gcloud/configurations/config_default
    sed "1i--project_id=$(gcloud config get-value project)" -i ~/.bigqueryrc
fi

if [ "$#" = 0 ]; then
    # Default to 20 workers because most of the tests are io-bound
    exec pytest --black --flake8 --mypy-ignore-missing-imports --pydocstyle --numprocesses 20
elif [ "${1:0:1}" = - ]; then
    # First argument is a flag, assume intended executable is pytest
    exec pytest "$@"
elif [ "$1" = "query" ]; then
    # For an invocation like:
    #     query [options] FILE
    # we dispatch to a script that inspects metadata and emulates the following call:
    #     bq query [options] < FILE 
    exec script/run_query --query_file="${@: -1}" "${@:1:$#-1}"
elif [ "$XCOM_PUSH" = "true" ]; then
    # KubernetesPodOperator will extract the contents of /airflow/xcom/return.json as an xcom
    # if the xcom_push parameter is true
    mkdir -p /airflow/xcom/ && touch /airflow/xcom/return.json
    exec "$@" | tee /airflow/xcom/return.json
else
    exec "$@"
fi
