Add the auth.ProcessEnv function

internal/testutil/env.go

@@ -13,19 +13,11 @@ import (
 // The original environment is restored upon test completion.
 // Note: use of this function is incompatible with parallel execution.
 func CleanupEnvironment(t TestingT) {
-	// Restore environment when test finishes.
-	environ := os.Environ()
-	t.Cleanup(func() {
-		// Restore original environment.
-		for _, kv := range environ {
-			kvs := strings.SplitN(kv, "=", 2)
-			os.Setenv(kvs[0], kvs[1])
-		}
-	})
-
 	path := os.Getenv("PATH")
 	pwd := os.Getenv("PWD")
-	os.Clearenv()
+
+	// Clear all environment variables.
+	ClearEnvironment(t)
 
 	// We use t.Setenv instead of os.Setenv because the former actively
 	// prevents a test being run with t.Parallel. Modifying the environment
@@ -38,6 +30,23 @@ func CleanupEnvironment(t TestingT) {
 	}
 }
 
+// ClearEnvironment sets up an empty environment with no environment variables set.
+// The original environment is restored upon test completion.
+// Note: use of this function is incompatible with parallel execution
+func ClearEnvironment(t TestingT) {
+	// Restore environment when test finishes.
+	environ := os.Environ()
+	t.Cleanup(func() {
+		// Restore original environment.
+		for _, kv := range environ {
+			kvs := strings.SplitN(kv, "=", 2)
+			os.Setenv(kvs[0], kvs[1])
+		}
+	})
+
+	os.Clearenv()
+}
+
 // Changes into specified directory for the duration of the test.
 // Returns the current working directory.
 func Chdir(t TestingT, dir string) string {

libs/auth/env.go

@@ -1,6 +1,14 @@
 package auth
 
-import "github.com/databricks/databricks-sdk-go/config"
+import (
+	"fmt"
+	"os"
+	"slices"
+	"sort"
+	"strings"
+
+	"github.com/databricks/databricks-sdk-go/config"
+)
 
 // Env generates the authentication environment variables we need to set for
 // downstream applications from the CLI to work correctly.
@@ -38,3 +46,66 @@ func GetEnvFor(name string) (string, bool) {
 
 	return "", false
 }
+
+func envVars() []string {
+	out := []string{}
+
+	for _, attr := range config.ConfigAttributes {
+		if len(attr.EnvVars) == 0 {
+			continue
+		}
+
+		out = append(out, attr.EnvVars[0])
+	}
+
+	return out
+}
+
+// ProcessEnv generates the environment variables can be set to authenticate downstream
+// processes to use the same auth credentials as in cfg.
+func ProcessEnv(cfg *config.Config) []string {
+	// We want child telemetry processes to inherit environment variables like $HOME or $HTTPS_PROXY
+	// because they influence auth resolution.
+	base := os.Environ()
+
+	out := []string{}
+	authEnvVars := envVars()
+
+	// Remove any existing auth environment variables. This is done because
+	// the CLI offers multiple modalities of configuring authentication like
+	// `--profile` or `DATABRICKS_CONFIG_PROFILE` or `profile: <profile>` in the
+	// bundle config file.
+	//
+	// Each of these modalities have different priorities and thus we don't want
+	// any auth configuration to piggyback into the child process environment.
+	//
+	// This is a precaution to avoid conflicting auth configurations being passed
+	// to the child telemetry process.
+	//
+	// Normally this should be unnecessary because the SDK should error if multiple
+	// authentication methods have been configured. But there is no harm in doing this
+	// as a precaution.
+	for _, v := range base {
+		k, _, found := strings.Cut(v, "=")
+		if !found {
+			continue
+		}
+		if slices.Contains(authEnvVars, k) {
+			continue
+		}
+		out = append(out, v)
+	}
+
+	// Now add the necessary authentication environment variables.
+	newEnv := Env(cfg)
+	for k, v := range newEnv {
+		out = append(out, fmt.Sprintf("%s=%s", k, v))
+	}
+
+	// Sort the environment variables so that the output is deterministic.
+	sort.Slice(out, func(i, j int) bool {
+		return out[i] < out[j]
+	})
+
+	return out
+}

libs/auth/env_test.go

@@ -3,6 +3,7 @@ package auth
 import (
 	"testing"
 
+	"github.com/databricks/cli/internal/testutil"
 	"github.com/databricks/databricks-sdk-go/config"
 	"github.com/stretchr/testify/assert"
 )
@@ -79,3 +80,73 @@ func TestGetEnvFor(t *testing.T) {
 	assert.False(t, ok)
 	assert.Empty(t, out)
 }
+
+func TestAuthEnvVars(t *testing.T) {
+	expected := []string{
+		"DATABRICKS_HOST",
+		"DATABRICKS_CLUSTER_ID",
+		"DATABRICKS_WAREHOUSE_ID",
+		"DATABRICKS_SERVERLESS_COMPUTE_ID",
+		"DATABRICKS_METADATA_SERVICE_URL",
+		"DATABRICKS_ACCOUNT_ID",
+		"DATABRICKS_TOKEN",
+		"DATABRICKS_USERNAME",
+		"DATABRICKS_PASSWORD",
+		"DATABRICKS_CONFIG_PROFILE",
+		"DATABRICKS_CONFIG_FILE",
+		"DATABRICKS_GOOGLE_SERVICE_ACCOUNT",
+		"GOOGLE_CREDENTIALS",
+		"DATABRICKS_AZURE_RESOURCE_ID",
+		"ARM_USE_MSI",
+		"ARM_CLIENT_SECRET",
+		"ARM_CLIENT_ID",
+		"ARM_TENANT_ID",
+		"ACTIONS_ID_TOKEN_REQUEST_URL",
+		"ACTIONS_ID_TOKEN_REQUEST_TOKEN",
+		"ARM_ENVIRONMENT",
+		"DATABRICKS_AZURE_LOGIN_APP_ID",
+		"DATABRICKS_CLIENT_ID",
+		"DATABRICKS_CLIENT_SECRET",
+		"DATABRICKS_CLI_PATH",
+		"DATABRICKS_AUTH_TYPE",
+		"DATABRICKS_DEBUG_TRUNCATE_BYTES",
+		"DATABRICKS_DEBUG_HEADERS",
+		"DATABRICKS_RATE_LIMIT",
+	}
+
+	out := envVars()
+	assert.Equal(t, expected, out)
+}
+
+func TestAuthProcessEnv(t *testing.T) {
+	testutil.ClearEnvironment(t)
+
+	// Environment variables that should be inherited by child processes.
+	t.Setenv("HOME", "/home/user")
+	t.Setenv("HTTPS_PROXY", "http://proxy.com")
+
+	// Environment variables that should be cleaned up by ProcessEnv():
+	t.Setenv("DATABRICKS_HOST", "https://test.com")
+	t.Setenv("DATABRICKS_TOKEN", "test-token")
+	t.Setenv("DATABRICKS_PASSWORD", "test-password")
+	t.Setenv("DATABRICKS_METADATA_SERVICE_URL", "http://somurl.com")
+	t.Setenv("ARM_USE_MSI", "true")
+	t.Setenv("ARM_TENANT_ID", "test-tenant-id")
+	t.Setenv("ARM_CLIENT_ID", "test-client-id")
+	t.Setenv("ARM_CLIENT_SECRET", "test-client-secret")
+
+	in := &config.Config{
+		Host:  "https://newhost.com",
+		Token: "new-token",
+	}
+
+	expected := []string{
+		"DATABRICKS_HOST=https://newhost.com",
+		"DATABRICKS_TOKEN=new-token",
+		"HOME=/home/user",
+		"HTTPS_PROXY=http://proxy.com",
+	}
+
+	out := ProcessEnv(in)
+	assert.Equal(t, expected, out)
+}