Add the `auth.ProcessEnv` function (#2404)

## Changes
This function provides all environment variables necessary to
authenticate the downstream applications to the same credentials as the
parent process.

It's used in https://github.com/databricks/cli/pull/2278 and will also
be useful for the `databricks bundle exec` command.

## Tests
Unit test.
This commit is contained in:
shreyas-goenka 2025-03-04 20:17:11 +05:30 committed by GitHub
parent c0f5436a28
commit 549b226cbc
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 79 additions and 14 deletions

View File

@ -13,19 +13,11 @@ import (
// The original environment is restored upon test completion. // The original environment is restored upon test completion.
// Note: use of this function is incompatible with parallel execution. // Note: use of this function is incompatible with parallel execution.
func CleanupEnvironment(t TestingT) { 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") path := os.Getenv("PATH")
pwd := os.Getenv("PWD") pwd := os.Getenv("PWD")
os.Clearenv()
// Clear all environment variables.
NullEnvironment(t)
// We use t.Setenv instead of os.Setenv because the former actively // We use t.Setenv instead of os.Setenv because the former actively
// prevents a test being run with t.Parallel. Modifying the environment // 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. // Changes into specified directory for the duration of the test.
// Returns the current working directory. // Returns the current working directory.
func Chdir(t TestingT, dir string) string { func Chdir(t TestingT, dir string) string {

View File

@ -1,6 +1,13 @@
package auth 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 // Env generates the authentication environment variables we need to set for
// downstream applications from the CLI to work correctly. // 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 // This is useful for spawning subprocesses since you can unset all auth environment
// variables to clean up the environment before configuring authentication for the // variables to clean up the environment before configuring authentication for the
// child process. // child process.
func EnvVars() []string { func envVars() []string {
out := []string{} out := []string{}
for _, attr := range config.ConfigAttributes { for _, attr := range config.ConfigAttributes {
@ -57,3 +64,52 @@ func EnvVars() []string {
return out 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: <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
}

View File

@ -122,7 +122,7 @@ func TestAuthEnvVars(t *testing.T) {
"ACTIONS_ID_TOKEN_REQUEST_TOKEN", "ACTIONS_ID_TOKEN_REQUEST_TOKEN",
} }
out := EnvVars() out := envVars()
for _, v := range contains { for _, v := range contains {
assert.Contains(t, out, v) assert.Contains(t, out, v)
} }