fix: specific env instead of all required variables

ops/db-connect.sh

@@ -1,48 +1,96 @@
 #!/bin/sh
 #
-# This script is used to connect to the TDR Cloud SQL database in a GCP project.
+# Connect to a Terra Data Repo database in an environment.
 #
 # It uses the gcloud CLI to get the credentials for the cluster, kubectl to port
 # forward to the SQL proxy pod, and psql to connect to the database.
 #
 # You MUST have gcloud, kubectl, jq, and psql installed to run this script. In
 # addition, you MUST be connected to the VPN to access the database.
 #
-# Usage: ./db-connect.sh
-#
-# The following environment variables may be set to override the defaults:
-#
-#   PROJECT     the GCP project that contains the Cloud SQL database
-#
-#   NAMESPACE   the Kubernetes namespace for the SQL proxy pod
-#               these can be listed with `kubectl get namespaces`
-#
-#   SECRET      the secret path for the database password in GCP Secrets Manager
-#               the paths are listed at https://github.com/broadinstitute/vault-migration-tool/blob/main/mappings/tdr.yaml
-#
-#   PORT        the local port to forward to the SQL proxy pod
-#
-#   DATABASE    the name of the database to connect to
+# See usage section below for more details. All arguments are optional.
 #
 
-set -eu
+set -euo pipefail
 
-# Default values that may be overridden by environment variables
-PROJECT="${PROJECT:-broad-jade-dev}"
-NAMESPACE="${NAMESPACE:-dev}"
-SECRET="${SECRET:-helm-datarepodb}"
-PORT="${PORT:-5432}"
-DATABASE="${DATABASE:-stairway}"
+usage() {
+  cat <<EOF
+Usage: $0 [OPTION]...
+Connect to a Terra Data Repo database in an environment.
+
+  --env ENV             Environment to connect to (default: dev)
+  --port PORT           Local port to forward to the SQL proxy (default: 5432)
+  --database DATABASE   Database to connect to, either datarepo or stairway
+                        (default: datarepo)
+  --help                Display this help and exit
+EOF
+  exit 0
+}
 
 error() {
-    echo "ERROR: $1" >&2
-    exit 1
+  echo "ERROR: $1" >&2
+  exit 1
+}
+
+# default values that may be overridden by command line arguments or environment variables
+ENV="${ENV:-dev}"
+PORT="${PORT:-5432}"
+DATABASE="${DATABASE:-datarepo}"
+
+parse_cli_args() {
+  while [ $# -gt 0 ]; do
+    case "$1" in
+      --env)
+        ENV="$2"
+        shift 2
+        ;;
+      --port)
+        PORT="$2"
+        shift 2
+        ;;
+      --database)
+        if [ "$2" != "datarepo" ] && [ "$2" != "stairway" ]; then
+          error "Database must be one of 'datarepo' or 'stairway'"
+        fi
+        DATABASE="$2"
+        shift 2
+        ;;
+      --help)
+        usage
+        ;;
+      *)
+        error "Unknown option: $1. Try --help to see a list of all options."
+        ;;
+    esac
+  done
 }
 
 cleanup() {
   kill "$PID"
 }
 
+set_vars_from_env() {
+  case "$ENV" in
+    dev)
+      PROJECT="broad-jade-dev"
+      NAMESPACE="dev"
+      SECRET="helm-datarepodb"
+      ;;
+    staging)
+      PROJECT="terra-datarepo-staging"
+      NAMESPACE="terra-staging"
+      SECRET="sql-db"
+      ;;
+    *)
+      error "Unknown environment: $ENV"
+      ;;
+  esac
+}
+
+set_project_config() {
+  gcloud config set project "$PROJECT"
+}
+
 set_cluster_credentials() {
   CLUSTER_JSON=$(gcloud container clusters list --format="json")
 
@@ -64,6 +112,7 @@ port_forward_sqlproxy() {
   SQLPROXY_POD=$(kubectl get pods --namespace="$NAMESPACE" --output="json" | jq -r '[ .items[].metadata.name | select(. | contains("sqlproxy")) ].[0]')
   kubectl port-forward "$SQLPROXY_POD" --namespace "$NAMESPACE" "$PORT:5432" &
   PID=$!
+  trap cleanup EXIT
 }
 
 connect_cloud_sql_db() {
@@ -78,10 +127,9 @@ connect_cloud_sql_db() {
   psql "postgresql://$USERNAME:$PASSWORD@localhost:$PORT/$DATABASE"
 }
 
-trap cleanup EXIT
-
-gcloud config set project "$PROJECT"
-
+parse_cli_args "$@"
+set_vars_from_env
+set_project_config
 set_cluster_credentials
 port_forward_sqlproxy
 connect_cloud_sql_db