#!/bin/bash

set -e

# Use credentials from environment
if [ -n "$GCLOUD_SERVICE_KEY" ]; then
    echo "DEBUG: SETTING UP CREDS BLOCK HIT"  # TODO: once the issue is solved this can be removed.
    # 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
    if [ -z "$GOOGLE_APPLICATION_CREDENTIALS" ]; then
        # Create a unique credentials file since this script can be called in parallel by multiple tests.
        export GOOGLE_APPLICATION_CREDENTIALS=$(mktemp /tmp/gcp.json.XXXXXXXXXX)
    fi
    echo "$GCLOUD_SERVICE_KEY" > "$GOOGLE_APPLICATION_CREDENTIALS"
fi

if [ -n "$GOOGLE_APPLICATION_CREDENTIALS" ] && which gcloud > /dev/null; then
    gcloud --verbosity=debug --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 8 workers because most of the tests are io-bound
    # When using pytest-xdist, total worker initialization time scales
    # linearly to the number of workers. 8 workers seems to bring the best
    # performance. See https://github.com/pytest-dev/pytest-xdist/issues/346
    exec pytest --black --flake8 --isort --mypy-ignore-missing-imports --pydocstyle -n 8
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/bqetl query run "${@: -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