Create tmp files in the cache dir in terraform command runs (#395)

## Changes Passes through tmp dir related env vars to the terraform process. Incase any of them are not set, we assign temp dir inside bundle cache dir as the location terraform should use. ## Tests Manually checked that these env vars do override location where os.CreateTemp files are created
databricks · May 23, 2023 · c53ad86 · c53ad86
1 parent ebb39cc
commit c53ad86
Show file tree

Hide file tree

Showing 2 changed files with 234 additions and 0 deletions.
diff --git a/bundle/deploy/terraform/init.go b/bundle/deploy/terraform/init.go
@@ -6,6 +6,7 @@ import (
 	"os"
 	"os/exec"
 	"path/filepath"
+	"runtime"
 	"strings"
 
 	"github.com/databricks/cli/bundle"
@@ -69,6 +70,44 @@ func (m *initialize) findExecPath(ctx context.Context, b *bundle.Bundle, tf *con
 	return tf.ExecPath, nil
 }
 
+// This function sets temp dir location for terraform to use. If user does not
+// specify anything here, we fall back to a `tmp` directory in the bundle's cache
+// directory
+//
+// This is necessary to avoid trying to create temporary files in directories
+// the CLI and its dependencies do not have access to.
+//
+// see: os.TempDir for more context
+func setTempDirEnvVars(env map[string]string, b *bundle.Bundle) error {
+	switch runtime.GOOS {
+	case "windows":
+		if v, ok := os.LookupEnv("TMP"); ok {
+			env["TMP"] = v
+		} else if v, ok := os.LookupEnv("TEMP"); ok {
+			env["TEMP"] = v
+		} else if v, ok := os.LookupEnv("USERPROFILE"); ok {
+			env["USERPROFILE"] = v
+		} else {
+			tmpDir, err := b.CacheDir("tmp")
+			if err != nil {
+				return err
+			}
+			env["TMP"] = tmpDir
+		}
+	default:
+		if v, ok := os.LookupEnv("TMPDIR"); ok {
+			env["TMPDIR"] = v
+		} else {
+			tmpDir, err := b.CacheDir("tmp")
+			if err != nil {
+				return err
+			}
+			env["TMPDIR"] = tmpDir
+		}
+	}
+	return nil
+}
+
 func (m *initialize) Apply(ctx context.Context, b *bundle.Bundle) ([]bundle.Mutator, error) {
 	tfConfig := b.Config.Bundle.Terraform
 	if tfConfig == nil {
@@ -102,6 +141,12 @@ func (m *initialize) Apply(ctx context.Context, b *bundle.Bundle) ([]bundle.Muta
 		env["HOME"] = home
 	}
 
+	// Set the temporary directory environment variables
+	err = setTempDirEnvVars(env, b)
+	if err != nil {
+		return nil, err
+	}
+
 	// Configure environment variables for auth for Terraform to use.
 	log.Debugf(ctx, "Environment variables for Terraform: %s", strings.Join(maps.Keys(env), ", "))
 	err = tf.SetEnv(env)

diff --git a/bundle/deploy/terraform/init_test.go b/bundle/deploy/terraform/init_test.go
@@ -2,14 +2,23 @@ package terraform
 
 import (
 	"context"
+	"os"
 	"os/exec"
+	"runtime"
 	"testing"
 
 	"github.com/databricks/cli/bundle"
 	"github.com/databricks/cli/bundle/config"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
+func unsetEnv(t *testing.T, name string) {
+	t.Setenv(name, "")
+	err := os.Unsetenv(name)
+	require.NoError(t, err)
+}
+
 func TestInitEnvironmentVariables(t *testing.T) {
 	_, err := exec.LookPath("terraform")
 	if err != nil {
@@ -37,3 +46,183 @@ func TestInitEnvironmentVariables(t *testing.T) {
 	_, err = Initialize().Apply(context.Background(), bundle)
 	require.NoError(t, err)
 }
+
+func TestSetTempDirEnvVarsForUnixWithTmpDirSet(t *testing.T) {
+	if runtime.GOOS != "darwin" && runtime.GOOS != "linux" {
+		t.SkipNow()
+	}
+
+	b := &bundle.Bundle{
+		Config: config.Root{
+			Path: t.TempDir(),
+			Bundle: config.Bundle{
+				Environment: "whatever",
+			},
+		},
+	}
+
+	// Set TMPDIR environment variable
+	t.Setenv("TMPDIR", "/foo/bar")
+
+	// compute env
+	env := make(map[string]string, 0)
+	err := setTempDirEnvVars(env, b)
+	require.NoError(t, err)
+
+	// assert that we pass through env var value
+	assert.Equal(t, map[string]string{
+		"TMPDIR": "/foo/bar",
+	}, env)
+}
+
+func TestSetTempDirEnvVarsForUnixWithTmpDirNotSet(t *testing.T) {
+	if runtime.GOOS != "darwin" && runtime.GOOS != "linux" {
+		t.SkipNow()
+	}
+
+	b := &bundle.Bundle{
+		Config: config.Root{
+			Path: t.TempDir(),
+			Bundle: config.Bundle{
+				Environment: "whatever",
+			},
+		},
+	}
+
+	// Unset TMPDIR environment variable confirm it's not set
+	unsetEnv(t, "TMPDIR")
+
+	// compute env
+	env := make(map[string]string, 0)
+	err := setTempDirEnvVars(env, b)
+	require.NoError(t, err)
+
+	// assert tmp dir is set to b.CacheDir("tmp")
+	tmpDir, err := b.CacheDir("tmp")
+	require.NoError(t, err)
+	assert.Equal(t, map[string]string{
+		"TMPDIR": tmpDir,
+	}, env)
+}
+
+func TestSetTempDirEnvVarsForWindowWithAllTmpDirEnvVarsSet(t *testing.T) {
+	if runtime.GOOS != "windows" {
+		t.SkipNow()
+	}
+
+	b := &bundle.Bundle{
+		Config: config.Root{
+			Path: t.TempDir(),
+			Bundle: config.Bundle{
+				Environment: "whatever",
+			},
+		},
+	}
+
+	// Set environment variables
+	t.Setenv("TMP", "c:\\foo\\a")
+	t.Setenv("TEMP", "c:\\foo\\b")
+	t.Setenv("USERPROFILE", "c:\\foo\\c")
+
+	// compute env
+	env := make(map[string]string, 0)
+	err := setTempDirEnvVars(env, b)
+	require.NoError(t, err)
+
+	// assert that we pass through the highest priority env var value
+	assert.Equal(t, map[string]string{
+		"TMP": "c:\\foo\\a",
+	}, env)
+}
+
+func TestSetTempDirEnvVarsForWindowWithUserProfileAndTempSet(t *testing.T) {
+	if runtime.GOOS != "windows" {
+		t.SkipNow()
+	}
+
+	b := &bundle.Bundle{
+		Config: config.Root{
+			Path: t.TempDir(),
+			Bundle: config.Bundle{
+				Environment: "whatever",
+			},
+		},
+	}
+
+	// Set environment variables
+	unsetEnv(t, "TMP")
+	t.Setenv("TEMP", "c:\\foo\\b")
+	t.Setenv("USERPROFILE", "c:\\foo\\c")
+
+	// compute env
+	env := make(map[string]string, 0)
+	err := setTempDirEnvVars(env, b)
+	require.NoError(t, err)
+
+	// assert that we pass through the highest priority env var value
+	assert.Equal(t, map[string]string{
+		"TEMP": "c:\\foo\\b",
+	}, env)
+}
+
+func TestSetTempDirEnvVarsForWindowWithUserProfileSet(t *testing.T) {
+	if runtime.GOOS != "windows" {
+		t.SkipNow()
+	}
+
+	b := &bundle.Bundle{
+		Config: config.Root{
+			Path: t.TempDir(),
+			Bundle: config.Bundle{
+				Environment: "whatever",
+			},
+		},
+	}
+
+	// Set environment variables
+	unsetEnv(t, "TMP")
+	unsetEnv(t, "TEMP")
+	t.Setenv("USERPROFILE", "c:\\foo\\c")
+
+	// compute env
+	env := make(map[string]string, 0)
+	err := setTempDirEnvVars(env, b)
+	require.NoError(t, err)
+
+	// assert that we pass through the user profile
+	assert.Equal(t, map[string]string{
+		"USERPROFILE": "c:\\foo\\c",
+	}, env)
+}
+
+func TestSetTempDirEnvVarsForWindowsWithoutAnyTempDirEnvVarsSet(t *testing.T) {
+	if runtime.GOOS != "windows" {
+		t.SkipNow()
+	}
+
+	b := &bundle.Bundle{
+		Config: config.Root{
+			Path: t.TempDir(),
+			Bundle: config.Bundle{
+				Environment: "whatever",
+			},
+		},
+	}
+
+	// unset all env vars
+	unsetEnv(t, "TMP")
+	unsetEnv(t, "TEMP")
+	unsetEnv(t, "USERPROFILE")
+
+	// compute env
+	env := make(map[string]string, 0)
+	err := setTempDirEnvVars(env, b)
+	require.NoError(t, err)
+
+	// assert TMP is set to b.CacheDir("tmp")
+	tmpDir, err := b.CacheDir("tmp")
+	require.NoError(t, err)
+	assert.Equal(t, map[string]string{
+		"TMP": tmpDir,
+	}, env)
+}