diff --git a/internal/testutil/env.go b/internal/testutil/env.go index 598229655..1ecbe6485 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. + NullEnvironment(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) { } } +// NullEnvironment sets up an empty environment with absolutely no environment variables set. +// The original environment is restored upon test completion. +// Note: use of this function is incompatible with parallel execution +func NullEnvironment(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 7413662b9..08282e463 100644 --- a/libs/auth/env.go +++ b/libs/auth/env.go @@ -1,6 +1,13 @@ package auth -import "github.com/databricks/databricks-sdk-go/config" +import ( + "fmt" + "os" + "slices" + "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. @@ -44,7 +51,7 @@ func GetEnvFor(name string) (string, bool) { // This is useful for spawning subprocesses since you can unset all auth environment // variables to clean up the environment before configuring authentication for the // child process. -func EnvVars() []string { +func envVars() []string { out := []string{} for _, attr := range config.ConfigAttributes { @@ -57,3 +64,52 @@ func EnvVars() []string { return out } + +// ProcessEnv generates the environment variables that should be set to authenticate +// downstream processes to use the same auth credentials as in cfg. +func ProcessEnv(cfg *config.Config) []string { + // We want child 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. + // Keeping the output deterministic helps with reproducibility and keeping the + // behavior consistent incase there are any issues. + slices.Sort(out) + + return out +} diff --git a/libs/auth/env_test.go b/libs/auth/env_test.go index 38dc1c6b7..5aa893e11 100644 --- a/libs/auth/env_test.go +++ b/libs/auth/env_test.go @@ -122,7 +122,7 @@ func TestAuthEnvVars(t *testing.T) { "ACTIONS_ID_TOKEN_REQUEST_TOKEN", } - out := EnvVars() + out := envVars() for _, v := range contains { assert.Contains(t, out, v) }