diff --git a/internal/testutil/env.go b/internal/testutil/env.go index 598229655..0aa97d779 100644 --- a/internal/testutil/env.go +++ b/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 { diff --git a/libs/auth/env.go b/libs/auth/env.go index f58f29ef7..6b53d22a6 100644 --- a/libs/auth/env.go +++ b/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. @@ -39,7 +47,7 @@ func GetEnvFor(name string) (string, bool) { return "", false } -func EnvVars() []string { +func envVars() []string { out := []string{} for _, attr := range config.ConfigAttributes { @@ -52,3 +60,52 @@ func EnvVars() []string { 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: ` 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 +} diff --git a/libs/auth/env_test.go b/libs/auth/env_test.go index d7efd0ddc..cf771cce4 100644 --- a/libs/auth/env_test.go +++ b/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" ) @@ -113,6 +114,39 @@ func TestAuthEnvVars(t *testing.T) { "DATABRICKS_RATE_LIMIT", } - out := EnvVars() + 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 process env: + 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) }