From 6c3ddbd921c568c3937d52a7167ae63fa3cbf533 Mon Sep 17 00:00:00 2001
From: shreyas-goenka <88374338+shreyas-goenka@users.noreply.github.com>
Date: Wed, 22 Jan 2025 17:44:54 +0530
Subject: [PATCH] Add `auth.Env` function (#2204)

## Changes
`auth.Env` is a generic function that we can use for authenticated tools
downstream to the CLI.

## Tests
Unit test.
---
 bundle/bundle.go      | 20 ++------------------
 libs/auth/env.go      | 26 ++++++++++++++++++++++++++
 libs/auth/env_test.go | 42 ++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 70 insertions(+), 18 deletions(-)
 create mode 100644 libs/auth/env.go
 create mode 100644 libs/auth/env_test.go

diff --git a/bundle/bundle.go b/bundle/bundle.go
index 3bf4ffb62..e715b8b2c 100644
--- a/bundle/bundle.go
+++ b/bundle/bundle.go
@@ -17,6 +17,7 @@ import (
 	"github.com/databricks/cli/bundle/config"
 	"github.com/databricks/cli/bundle/env"
 	"github.com/databricks/cli/bundle/metadata"
+	"github.com/databricks/cli/libs/auth"
 	"github.com/databricks/cli/libs/fileset"
 	"github.com/databricks/cli/libs/locker"
 	"github.com/databricks/cli/libs/log"
@@ -24,7 +25,6 @@ import (
 	"github.com/databricks/cli/libs/terraform"
 	"github.com/databricks/cli/libs/vfs"
 	"github.com/databricks/databricks-sdk-go"
-	sdkconfig "github.com/databricks/databricks-sdk-go/config"
 	"github.com/hashicorp/terraform-exec/tfexec"
 )
 
@@ -242,21 +242,5 @@ func (b *Bundle) AuthEnv() (map[string]string, error) {
 	}
 
 	cfg := b.client.Config
-	out := make(map[string]string)
-	for _, attr := range sdkconfig.ConfigAttributes {
-		// Ignore profile so that downstream tools don't try and reload
-		// the profile even though we know the current configuration is valid.
-		if attr.Name == "profile" {
-			continue
-		}
-		if len(attr.EnvVars) == 0 {
-			continue
-		}
-		if attr.IsZero(cfg) {
-			continue
-		}
-		out[attr.EnvVars[0]] = attr.GetString(cfg)
-	}
-
-	return out, nil
+	return auth.Env(cfg), nil
 }
diff --git a/libs/auth/env.go b/libs/auth/env.go
new file mode 100644
index 000000000..c58cc53e3
--- /dev/null
+++ b/libs/auth/env.go
@@ -0,0 +1,26 @@
+package auth
+
+import "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.
+func Env(cfg *config.Config) map[string]string {
+	out := make(map[string]string)
+	for _, attr := range config.ConfigAttributes {
+		// Ignore profile so that downstream tools don't try and reload
+		// the profile. We know the current configuration is already valid since
+		// otherwise the CLI would have thrown an error when loading it.
+		if attr.Name == "profile" {
+			continue
+		}
+		if len(attr.EnvVars) == 0 {
+			continue
+		}
+		if attr.IsZero(cfg) {
+			continue
+		}
+		out[attr.EnvVars[0]] = attr.GetString(cfg)
+	}
+
+	return out
+}
diff --git a/libs/auth/env_test.go b/libs/auth/env_test.go
new file mode 100644
index 000000000..be1cfc7ac
--- /dev/null
+++ b/libs/auth/env_test.go
@@ -0,0 +1,42 @@
+package auth
+
+import (
+	"testing"
+
+	"github.com/databricks/databricks-sdk-go/config"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestAuthEnv(t *testing.T) {
+	in := &config.Config{
+		Profile:            "thisshouldbeignored",
+		Host:               "https://test.com",
+		Token:              "test-token",
+		Password:           "test-password",
+		MetadataServiceURL: "http://somurl.com",
+
+		AzureUseMSI:       true,
+		AzureTenantID:     "test-tenant-id",
+		AzureClientID:     "test-client-id",
+		AzureClientSecret: "test-client-secret",
+
+		ActionsIDTokenRequestToken: "test-actions-id-token-request-token",
+	}
+
+	expected := map[string]string{
+		"DATABRICKS_HOST":                 "https://test.com",
+		"DATABRICKS_TOKEN":                "test-token",
+		"DATABRICKS_PASSWORD":             "test-password",
+		"DATABRICKS_METADATA_SERVICE_URL": "http://somurl.com",
+
+		"ARM_USE_MSI":       "true",
+		"ARM_TENANT_ID":     "test-tenant-id",
+		"ARM_CLIENT_ID":     "test-client-id",
+		"ARM_CLIENT_SECRET": "test-client-secret",
+
+		"ACTIONS_ID_TOKEN_REQUEST_TOKEN": "test-actions-id-token-request-token",
+	}
+
+	out := Env(in)
+	assert.Equal(t, expected, out)
+}