From 74f1989006cf9536f412fccaa8983f184c7756b9 Mon Sep 17 00:00:00 2001
From: Ben Phegan <BenPhegan@users.noreply.github.com>
Date: Tue, 26 Mar 2024 15:32:04 +1100
Subject: [PATCH] Allow setting the token cache directory without modifying
 HOME env via the DATABRICKS_TOKEN_CACHE_DIR env var

---
 libs/auth/cache/cache.go      |  7 +++++++
 libs/auth/cache/cache_test.go | 10 ++++++++++
 2 files changed, 17 insertions(+)

diff --git a/libs/auth/cache/cache.go b/libs/auth/cache/cache.go
index 5511c192..ac6180a4 100644
--- a/libs/auth/cache/cache.go
+++ b/libs/auth/cache/cache.go
@@ -74,6 +74,13 @@ func (c *TokenCache) Lookup(key string) (*oauth2.Token, error) {
 }
 
 func (c *TokenCache) location() (string, error) {
+	// Allow users to override the location of the token cache file
+	tokenCacheDirName, exists := os.LookupEnv("DATABRICKS_TOKEN_CACHE_DIR")
+	if exists {
+		return filepath.Join(tokenCacheDirName, tokenCacheFile), nil
+	}
+
+	// Otherwise, default to using the home directory
 	home, err := os.UserHomeDir()
 	if err != nil {
 		return "", fmt.Errorf("home: %w", err)
diff --git a/libs/auth/cache/cache_test.go b/libs/auth/cache/cache_test.go
index 6529882c..7c1b3cf0 100644
--- a/libs/auth/cache/cache_test.go
+++ b/libs/auth/cache/cache_test.go
@@ -103,3 +103,13 @@ func TestStoreOnDev(t *testing.T) {
 	// macOS: read-only file system
 	assert.Error(t, err)
 }
+
+func TestOverrideTokenCacheLocationUsingEnvVar(t *testing.T) {
+	t.Setenv("DATABRICKS_TOKEN_CACHE_DIR", os.TempDir())
+	c := &TokenCache{}
+	err := c.Store("x", &oauth2.Token{
+		AccessToken: "abc",
+	})
+	require.NoError(t, err)
+	assert.Contains(t, c.fileLocation, os.TempDir())
+}