Merge #145610

145610: drtprod: only set roachprod environment variables if not already set r=nameisbhaskar a=shailendra-patel

This change modifies the environment variable initialization in the drtprod CLI to only set ROACHPROD_DNS, ROACHPROD_GCE_DNS_DOMAIN, ROACHPROD_GCE_DNS_ZONE, and ROACHPROD_GCE_DEFAULT_PROJECT if they are not already set in the environment.

Previously, these variables were always overwritten, which could cause issues when users needed to use custom values. This change follows the same pattern already in place for the DD_API_KEY environment variable, providing more flexibility while maintaining the default configuration for most users.

Epic: none

Release note: None

Co-authored-by: Shailendra Patel <[email protected]>

Author: craig[bot]

pkg/cmd/drtprod/cli/handlers.go

@@ -16,12 +16,19 @@ import (
 	"github.com/spf13/cobra"
 )
 
+// setEnvIfNotExists sets an environment variable only if it doesn't already exist.
+func setEnvIfNotExists(key, value string) {
+	if _, exists := os.LookupEnv(key); !exists {
+		_ = os.Setenv(key, value)
+	}
+}
+
 func init() {
-	// Set environment variables for the GCE project and DNS configurations.
-	_ = os.Setenv("ROACHPROD_DNS", "drt.crdb.io")
-	_ = os.Setenv("ROACHPROD_GCE_DNS_DOMAIN", "drt.crdb.io")
-	_ = os.Setenv("ROACHPROD_GCE_DNS_ZONE", "drt")
-	_ = os.Setenv("ROACHPROD_GCE_DEFAULT_PROJECT", "cockroach-drt")
+	// Set environment variables for the GCE project and DNS configurations if not already set.
+	setEnvIfNotExists("ROACHPROD_DNS", "drt.crdb.io")
+	setEnvIfNotExists("ROACHPROD_GCE_DNS_DOMAIN", "drt.crdb.io")
+	setEnvIfNotExists("ROACHPROD_GCE_DNS_ZONE", "drt")
+	setEnvIfNotExists("ROACHPROD_GCE_DEFAULT_PROJECT", "cockroach-drt")
 
 	if _, exists := os.LookupEnv("DD_API_KEY"); !exists {
 		// set the DD_API_KEY if we are able to fetch it from the secrets.

pkg/cmd/drtprod/scripts/generate_tpcc_run.sh

@@ -99,6 +99,7 @@ for ((NODE=0; NODE<WORKLOAD_NODES; NODE++)); do
   cat <<EOF >/tmp/tpcc_run_${suffix}.sh
 #!/usr/bin/env bash
 
+export ROACHPROD_GCE_DEFAULT_PROJECT=$ROACHPROD_GCE_DEFAULT_PROJECT
 ./drtprod sync
 $([ "$execute_script" = "true" ] && [ "$NODE" -eq 0 ] && echo "${pwd}/tpcc_init_${suffix}.sh")
 PGURLS=\$(./drtprod load-balancer pgurl $CLUSTER | sed s/\'//g)