From 807a37b36ab621802bffafd5ff3a7074404e9363 Mon Sep 17 00:00:00 2001 From: shreyas-goenka <88374338+shreyas-goenka@users.noreply.github.com> Date: Mon, 3 Mar 2025 20:28:43 +0530 Subject: [PATCH] Add the `auth.EnvVars` function (#2403) ## Changes This PR adds the auth.EnvVars function, which is a list of all environment variables that the SDK uses to read auth configuration. This is useful for spawning subprocesses since you can unset all auth environment variables to clean up the environment before configuring the auth. It's used in #2278 today and will also be useful for the `databricks bundle exec` command. ## Tests Unit test. --- libs/auth/env.go | 19 +++++++++++++++++ libs/auth/env_test.go | 48 +++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 67 insertions(+) diff --git a/libs/auth/env.go b/libs/auth/env.go index 5c0d21292..7413662b9 100644 --- a/libs/auth/env.go +++ b/libs/auth/env.go @@ -38,3 +38,22 @@ func GetEnvFor(name string) (string, bool) { return "", false } + +// EnvVars returns the list of environment variables that the SDK reads to configure +// authentication. +// 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 { + out := []string{} + + for _, attr := range config.ConfigAttributes { + if len(attr.EnvVars) == 0 { + continue + } + + out = append(out, attr.EnvVars...) + } + + return out +} diff --git a/libs/auth/env_test.go b/libs/auth/env_test.go index 850110b60..38dc1c6b7 100644 --- a/libs/auth/env_test.go +++ b/libs/auth/env_test.go @@ -79,3 +79,51 @@ func TestGetEnvFor(t *testing.T) { assert.False(t, ok) assert.Empty(t, out) } + +func TestAuthEnvVars(t *testing.T) { + // Few common environment variables that we expect the SDK to support. + contains := []string{ + // Generic attributes. + "DATABRICKS_HOST", + "DATABRICKS_CONFIG_PROFILE", + "DATABRICKS_AUTH_TYPE", + "DATABRICKS_METADATA_SERVICE_URL", + "DATABRICKS_CONFIG_FILE", + + // OAuth specific attributes. + "DATABRICKS_CLIENT_ID", + "DATABRICKS_CLIENT_SECRET", + "DATABRICKS_CLI_PATH", + + // Google specific attributes. + "DATABRICKS_GOOGLE_SERVICE_ACCOUNT", + "GOOGLE_CREDENTIALS", + + // Personal access token specific attributes. + "DATABRICKS_TOKEN", + + // Databricks password specific attributes. + "DATABRICKS_USERNAME", + "DATABRICKS_PASSWORD", + + // Account authentication attributes. + "DATABRICKS_ACCOUNT_ID", + + // Azure attributes + "DATABRICKS_AZURE_RESOURCE_ID", + "ARM_USE_MSI", + "ARM_CLIENT_SECRET", + "ARM_CLIENT_ID", + "ARM_TENANT_ID", + "ARM_ENVIRONMENT", + + // Github attributes + "ACTIONS_ID_TOKEN_REQUEST_URL", + "ACTIONS_ID_TOKEN_REQUEST_TOKEN", + } + + out := EnvVars() + for _, v := range contains { + assert.Contains(t, out, v) + } +}