From 31b7ad071c4b1f836a1170e369045f9ac4911f84 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Fri, 28 Feb 2025 14:12:15 +0100 Subject: [PATCH 1/9] Add the auth.EnvVars function --- libs/auth/env.go | 14 ++++++++++++++ libs/auth/env_test.go | 37 +++++++++++++++++++++++++++++++++++++ 2 files changed, 51 insertions(+) diff --git a/libs/auth/env.go b/libs/auth/env.go index 5c0d2129297..f58f29ef7d8 100644 --- a/libs/auth/env.go +++ b/libs/auth/env.go @@ -38,3 +38,17 @@ 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 +} diff --git a/libs/auth/env_test.go b/libs/auth/env_test.go index 850110b6029..d7efd0ddc4f 100644 --- a/libs/auth/env_test.go +++ b/libs/auth/env_test.go @@ -79,3 +79,40 @@ 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) +} From bf234c4f5018bb3096aba920d8e281a5e53c5715 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Fri, 28 Feb 2025 14:41:37 +0100 Subject: [PATCH 2/9] Add the auth.ProcessEnv function --- internal/testutil/env.go | 31 ++++++++++++-------- libs/auth/env.go | 61 ++++++++++++++++++++++++++++++++++++++-- libs/auth/env_test.go | 36 +++++++++++++++++++++++- 3 files changed, 114 insertions(+), 14 deletions(-) diff --git a/internal/testutil/env.go b/internal/testutil/env.go index 598229655fd..0aa97d77991 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 f58f29ef7d8..6b53d22a606 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 d7efd0ddc4f..cf771cce4dd 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) } From 42f6ecf6d72ca8cc312ae7ec56adc681c69dab78 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Fri, 28 Feb 2025 14:47:16 +0100 Subject: [PATCH 3/9] - --- libs/auth/env_test.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/libs/auth/env_test.go b/libs/auth/env_test.go index cf771cce4dd..e04153ed308 100644 --- a/libs/auth/env_test.go +++ b/libs/auth/env_test.go @@ -125,7 +125,7 @@ func TestAuthProcessEnv(t *testing.T) { t.Setenv("HOME", "/home/user") t.Setenv("HTTPS_PROXY", "http://proxy.com") - // Environment variables that should be cleaned up by process env: + // 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") From f559322e8672fc45581d2f3b39a858b52d29832c Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Mon, 3 Mar 2025 10:41:31 +0100 Subject: [PATCH 4/9] address comments --- internal/testutil/env.go | 6 +++--- libs/auth/env_test.go | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/internal/testutil/env.go b/internal/testutil/env.go index 0aa97d77991..1ecbe6485cd 100644 --- a/internal/testutil/env.go +++ b/internal/testutil/env.go @@ -17,7 +17,7 @@ func CleanupEnvironment(t TestingT) { pwd := os.Getenv("PWD") // Clear all environment variables. - ClearEnvironment(t) + 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 @@ -30,10 +30,10 @@ func CleanupEnvironment(t TestingT) { } } -// ClearEnvironment sets up an empty environment with no environment variables set. +// 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 ClearEnvironment(t TestingT) { +func NullEnvironment(t TestingT) { // Restore environment when test finishes. environ := os.Environ() t.Cleanup(func() { diff --git a/libs/auth/env_test.go b/libs/auth/env_test.go index e04153ed308..9c2bf9a1969 100644 --- a/libs/auth/env_test.go +++ b/libs/auth/env_test.go @@ -119,7 +119,7 @@ func TestAuthEnvVars(t *testing.T) { } func TestAuthProcessEnv(t *testing.T) { - testutil.ClearEnvironment(t) + testutil.NullEnvironment(t) // Environment variables that should be inherited by child processes. t.Setenv("HOME", "/home/user") From a4269302838e827418871f581944786519d44eb8 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Mon, 3 Mar 2025 14:43:06 +0100 Subject: [PATCH 5/9] address some comments --- libs/auth/env.go | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/libs/auth/env.go b/libs/auth/env.go index 6b53d22a606..ddff3168fec 100644 --- a/libs/auth/env.go +++ b/libs/auth/env.go @@ -64,7 +64,7 @@ func envVars() []string { // 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 + // We want child processes to inherit environment variables like $HOME or $HTTPS_PROXY // because they influence auth resolution. base := os.Environ() @@ -102,7 +102,8 @@ func ProcessEnv(cfg *config.Config) []string { out = append(out, fmt.Sprintf("%s=%s", k, v)) } - // Sort the environment variables so that the output is deterministic. + // Sort the environment variables so that the output is deterministic for + // unit tests. sort.Slice(out, func(i, j int) bool { return out[i] < out[j] }) From 55745635064e93a19ea1b4a9371089eb79c9781f Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Mon, 3 Mar 2025 14:54:09 +0100 Subject: [PATCH 6/9] - --- libs/auth/env.go | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/libs/auth/env.go b/libs/auth/env.go index ddff3168fec..f85817346ca 100644 --- a/libs/auth/env.go +++ b/libs/auth/env.go @@ -102,8 +102,7 @@ func ProcessEnv(cfg *config.Config) []string { out = append(out, fmt.Sprintf("%s=%s", k, v)) } - // Sort the environment variables so that the output is deterministic for - // unit tests. + // Sort the environment variables so that the output is deterministic. sort.Slice(out, func(i, j int) bool { return out[i] < out[j] }) From 8283d61b921d4f2a9873eef397df6951ecab5685 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Mon, 3 Mar 2025 15:08:22 +0100 Subject: [PATCH 7/9] address some comments --- libs/auth/env.go | 2 ++ 1 file changed, 2 insertions(+) diff --git a/libs/auth/env.go b/libs/auth/env.go index f85817346ca..aa2136e4332 100644 --- a/libs/auth/env.go +++ b/libs/auth/env.go @@ -103,6 +103,8 @@ func ProcessEnv(cfg *config.Config) []string { } // 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. sort.Slice(out, func(i, j int) bool { return out[i] < out[j] }) From fe11a8f6f6b6df12d7e24b86d525aa0258577812 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Mon, 3 Mar 2025 15:09:30 +0100 Subject: [PATCH 8/9] use better sort function --- libs/auth/env.go | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/libs/auth/env.go b/libs/auth/env.go index aa2136e4332..a630374f5c7 100644 --- a/libs/auth/env.go +++ b/libs/auth/env.go @@ -4,7 +4,6 @@ import ( "fmt" "os" "slices" - "sort" "strings" "github.com/databricks/databricks-sdk-go/config" @@ -105,9 +104,7 @@ func ProcessEnv(cfg *config.Config) []string { // 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. - sort.Slice(out, func(i, j int) bool { - return out[i] < out[j] - }) + slices.Sort(out) return out } From 32de7089bdc4f2158e68feec18f389515b3b7966 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Mon, 3 Mar 2025 16:37:54 +0100 Subject: [PATCH 9/9] - --- libs/auth/env.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/libs/auth/env.go b/libs/auth/env.go index a8d5cc913f0..08282e46349 100644 --- a/libs/auth/env.go +++ b/libs/auth/env.go @@ -65,8 +65,8 @@ 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. +// 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.