Consolidate environment variable interaction (databricks#747)

## Changes

There are a couple places throughout the code base where interaction
with environment variables takes place. Moreover, more than one of these
would try to read a value from more than one environment variable as
fallback (for backwards compatibility). This change consolidates those
accesses.

The majority of diffs in this change are mechanical (i.e. add an
argument or replace a call).

This change:
* Moves common environment variable lookups for bundles to
`bundles/env`.
* Adds a `libs/env` package that wraps `os.LookupEnv` and `os.Getenv`
and allows for overrides to take place in a `context.Context`. By
scoping overrides to a `context.Context` we can avoid `t.Setenv` in
testing and unlock parallel test execution for integration tests.
* Updates call sites to pass through a `context.Context` where needed.
* For bundles, introduces `DATABRICKS_BUNDLE_ROOT` as new primary
variable instead of `BUNDLE_ROOT`. This was the last environment
variable that did not use the `DATABRICKS_` prefix.

## Tests

Unit tests pass.

bundle/bundle.go

@@ -14,6 +14,7 @@ import (
 	"sync"
 
 	"github.com/databricks/cli/bundle/config"
+	"github.com/databricks/cli/bundle/env"
 	"github.com/databricks/cli/folders"
 	"github.com/databricks/cli/libs/git"
 	"github.com/databricks/cli/libs/locker"
@@ -51,8 +52,6 @@ type Bundle struct {
 	AutoApprove bool
 }
 
-const ExtraIncludePathsKey string = "DATABRICKS_BUNDLE_INCLUDES"
-
 func Load(ctx context.Context, path string) (*Bundle, error) {
 	bundle := &Bundle{}
 	stat, err := os.Stat(path)
@@ -61,9 +60,9 @@ func Load(ctx context.Context, path string) (*Bundle, error) {
 	}
 	configFile, err := config.FileNames.FindInPath(path)
 	if err != nil {
-		_, hasIncludePathEnv := os.LookupEnv(ExtraIncludePathsKey)
-		_, hasBundleRootEnv := os.LookupEnv(envBundleRoot)
-		if hasIncludePathEnv && hasBundleRootEnv && stat.IsDir() {
+		_, hasRootEnv := env.Root(ctx)
+		_, hasIncludesEnv := env.Includes(ctx)
+		if hasRootEnv && hasIncludesEnv && stat.IsDir() {
 			log.Debugf(ctx, "No bundle configuration; using bundle root: %s", path)
 			bundle.Config = config.Root{
 				Path: path,
@@ -86,7 +85,7 @@ func Load(ctx context.Context, path string) (*Bundle, error) {
 // MustLoad returns a bundle configuration.
 // It returns an error if a bundle was not found or could not be loaded.
 func MustLoad(ctx context.Context) (*Bundle, error) {
-	root, err := mustGetRoot()
+	root, err := mustGetRoot(ctx)
 	if err != nil {
 		return nil, err
 	}
@@ -98,7 +97,7 @@ func MustLoad(ctx context.Context) (*Bundle, error) {
 // It returns an error if a bundle was found but could not be loaded.
 // It returns a `nil` bundle if a bundle was not found.
 func TryLoad(ctx context.Context) (*Bundle, error) {
-	root, err := tryGetRoot()
+	root, err := tryGetRoot(ctx)
 	if err != nil {
 		return nil, err
 	}
@@ -124,13 +123,12 @@ func (b *Bundle) WorkspaceClient() *databricks.WorkspaceClient {
 
 // CacheDir returns directory to use for temporary files for this bundle.
 // Scoped to the bundle's target.
-func (b *Bundle) CacheDir(paths ...string) (string, error) {
+func (b *Bundle) CacheDir(ctx context.Context, paths ...string) (string, error) {
 	if b.Config.Bundle.Target == "" {
 		panic("target not set")
 	}
 
-	cacheDirName, exists := os.LookupEnv("DATABRICKS_BUNDLE_TMP")
-
+	cacheDirName, exists := env.TempDir(ctx)
 	if !exists || cacheDirName == "" {
 		cacheDirName = filepath.Join(
 			// Anchor at bundle root directory.
@@ -163,8 +161,8 @@ func (b *Bundle) CacheDir(paths ...string) (string, error) {
 
 // This directory is used to store and automaticaly sync internal bundle files, such as, f.e
 // notebook trampoline files for Python wheel and etc.
-func (b *Bundle) InternalDir() (string, error) {
-	cacheDir, err := b.CacheDir()
+func (b *Bundle) InternalDir(ctx context.Context) (string, error) {
+	cacheDir, err := b.CacheDir(ctx)
 	if err != nil {
 		return "", err
 	}
@@ -181,8 +179,8 @@ func (b *Bundle) InternalDir() (string, error) {
 // GetSyncIncludePatterns returns a list of user defined includes
 // And also adds InternalDir folder to include list for sync command
 // so this folder is always synced
-func (b *Bundle) GetSyncIncludePatterns() ([]string, error) {
-	internalDir, err := b.InternalDir()
+func (b *Bundle) GetSyncIncludePatterns(ctx context.Context) ([]string, error) {
+	internalDir, err := b.InternalDir(ctx)
 	if err != nil {
 		return nil, err
 	}

bundle/bundle_test.go

@@ -6,6 +6,7 @@ import (
 	"path/filepath"
 	"testing"
 
+	"github.com/databricks/cli/bundle/env"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -23,12 +24,13 @@ func TestLoadExists(t *testing.T) {
 }
 
 func TestBundleCacheDir(t *testing.T) {
+	ctx := context.Background()
 	projectDir := t.TempDir()
 	f1, err := os.Create(filepath.Join(projectDir, "databricks.yml"))
 	require.NoError(t, err)
 	f1.Close()
 
-	bundle, err := Load(context.Background(), projectDir)
+	bundle, err := Load(ctx, projectDir)
 	require.NoError(t, err)
 
 	// Artificially set target.
@@ -38,21 +40,22 @@ func TestBundleCacheDir(t *testing.T) {
 	// unset env variable in case it's set
 	t.Setenv("DATABRICKS_BUNDLE_TMP", "")
 
-	cacheDir, err := bundle.CacheDir()
+	cacheDir, err := bundle.CacheDir(ctx)
 
 	// format is <CWD>/.databricks/bundle/<target>
 	assert.NoError(t, err)
 	assert.Equal(t, filepath.Join(projectDir, ".databricks", "bundle", "default"), cacheDir)
 }
 
 func TestBundleCacheDirOverride(t *testing.T) {
+	ctx := context.Background()
 	projectDir := t.TempDir()
 	bundleTmpDir := t.TempDir()
 	f1, err := os.Create(filepath.Join(projectDir, "databricks.yml"))
 	require.NoError(t, err)
 	f1.Close()
 
-	bundle, err := Load(context.Background(), projectDir)
+	bundle, err := Load(ctx, projectDir)
 	require.NoError(t, err)
 
 	// Artificially set target.
@@ -62,22 +65,22 @@ func TestBundleCacheDirOverride(t *testing.T) {
 	// now we expect to use 'bundleTmpDir' instead of CWD/.databricks/bundle
 	t.Setenv("DATABRICKS_BUNDLE_TMP", bundleTmpDir)
 
-	cacheDir, err := bundle.CacheDir()
+	cacheDir, err := bundle.CacheDir(ctx)
 
 	// format is <DATABRICKS_BUNDLE_TMP>/<target>
 	assert.NoError(t, err)
 	assert.Equal(t, filepath.Join(bundleTmpDir, "default"), cacheDir)
 }
 
 func TestBundleMustLoadSuccess(t *testing.T) {
-	t.Setenv(envBundleRoot, "./tests/basic")
+	t.Setenv(env.RootVariable, "./tests/basic")
 	b, err := MustLoad(context.Background())
 	require.NoError(t, err)
 	assert.Equal(t, "tests/basic", filepath.ToSlash(b.Config.Path))
 }
 
 func TestBundleMustLoadFailureWithEnv(t *testing.T) {
-	t.Setenv(envBundleRoot, "./tests/doesntexist")
+	t.Setenv(env.RootVariable, "./tests/doesntexist")
 	_, err := MustLoad(context.Background())
 	require.Error(t, err, "not a directory")
 }
@@ -89,14 +92,14 @@ func TestBundleMustLoadFailureIfNotFound(t *testing.T) {
 }
 
 func TestBundleTryLoadSuccess(t *testing.T) {
-	t.Setenv(envBundleRoot, "./tests/basic")
+	t.Setenv(env.RootVariable, "./tests/basic")
 	b, err := TryLoad(context.Background())
 	require.NoError(t, err)
 	assert.Equal(t, "tests/basic", filepath.ToSlash(b.Config.Path))
 }
 
 func TestBundleTryLoadFailureWithEnv(t *testing.T) {
-	t.Setenv(envBundleRoot, "./tests/doesntexist")
+	t.Setenv(env.RootVariable, "./tests/doesntexist")
 	_, err := TryLoad(context.Background())
 	require.Error(t, err, "not a directory")
 }

bundle/config/mutator/override_compute.go

@@ -3,11 +3,11 @@ package mutator
 import (
 	"context"
 	"fmt"
-	"os"
 
 	"github.com/databricks/cli/bundle"
 	"github.com/databricks/cli/bundle/config"
 	"github.com/databricks/cli/bundle/config/resources"
+	"github.com/databricks/cli/libs/env"
 )
 
 type overrideCompute struct{}
@@ -39,8 +39,8 @@ func (m *overrideCompute) Apply(ctx context.Context, b *bundle.Bundle) error {
 		}
 		return nil
 	}
-	if os.Getenv("DATABRICKS_CLUSTER_ID") != "" {
-		b.Config.Bundle.ComputeID = os.Getenv("DATABRICKS_CLUSTER_ID")
+	if v := env.Get(ctx, "DATABRICKS_CLUSTER_ID"); v != "" {
+		b.Config.Bundle.ComputeID = v
 	}
 
 	if b.Config.Bundle.ComputeID == "" {

bundle/config/mutator/process_root_includes.go

@@ -10,11 +10,12 @@ import (
 
 	"github.com/databricks/cli/bundle"
 	"github.com/databricks/cli/bundle/config"
+	"github.com/databricks/cli/bundle/env"
 )
 
 // Get extra include paths from environment variable
-func GetExtraIncludePaths() []string {
-	value, exists := os.LookupEnv(bundle.ExtraIncludePathsKey)
+func getExtraIncludePaths(ctx context.Context) []string {
+	value, exists := env.Includes(ctx)
 	if !exists {
 		return nil
 	}
@@ -48,7 +49,7 @@ func (m *processRootIncludes) Apply(ctx context.Context, b *bundle.Bundle) error
 	var files []string
 
 	// Converts extra include paths from environment variable to relative paths
-	for _, extraIncludePath := range GetExtraIncludePaths() {
+	for _, extraIncludePath := range getExtraIncludePaths(ctx) {
 		if filepath.IsAbs(extraIncludePath) {
 			rel, err := filepath.Rel(b.Config.Path, extraIncludePath)
 			if err != nil {

bundle/config/mutator/process_root_includes_test.go

@@ -2,16 +2,17 @@ package mutator_test
 
 import (
 	"context"
-	"fmt"
 	"os"
 	"path"
 	"path/filepath"
 	"runtime"
+	"strings"
 	"testing"
 
 	"github.com/databricks/cli/bundle"
 	"github.com/databricks/cli/bundle/config"
 	"github.com/databricks/cli/bundle/config/mutator"
+	"github.com/databricks/cli/bundle/env"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -129,10 +130,7 @@ func TestProcessRootIncludesExtrasFromEnvVar(t *testing.T) {
 	rootPath := t.TempDir()
 	testYamlName := "extra_include_path.yml"
 	touch(t, rootPath, testYamlName)
-	os.Setenv(bundle.ExtraIncludePathsKey, path.Join(rootPath, testYamlName))
-	t.Cleanup(func() {
-		os.Unsetenv(bundle.ExtraIncludePathsKey)
-	})
+	t.Setenv(env.IncludesVariable, path.Join(rootPath, testYamlName))
 
 	bundle := &bundle.Bundle{
 		Config: config.Root{
@@ -149,7 +147,13 @@ func TestProcessRootIncludesDedupExtrasFromEnvVar(t *testing.T) {
 	rootPath := t.TempDir()
 	testYamlName := "extra_include_path.yml"
 	touch(t, rootPath, testYamlName)
-	t.Setenv(bundle.ExtraIncludePathsKey, fmt.Sprintf("%s%s%s", path.Join(rootPath, testYamlName), string(os.PathListSeparator), path.Join(rootPath, testYamlName)))
+	t.Setenv(env.IncludesVariable, strings.Join(
+		[]string{
+			path.Join(rootPath, testYamlName),
+			path.Join(rootPath, testYamlName),
+		},
+		string(os.PathListSeparator),
+	))
 
 	bundle := &bundle.Bundle{
 		Config: config.Root{

bundle/config/mutator/set_variables.go

@@ -3,10 +3,10 @@ package mutator
 import (
 	"context"
 	"fmt"
-	"os"
 
 	"github.com/databricks/cli/bundle"
 	"github.com/databricks/cli/bundle/config/variable"
+	"github.com/databricks/cli/libs/env"
 )
 
 const bundleVarPrefix = "BUNDLE_VAR_"
@@ -21,15 +21,15 @@ func (m *setVariables) Name() string {
 	return "SetVariables"
 }
 
-func setVariable(v *variable.Variable, name string) error {
+func setVariable(ctx context.Context, v *variable.Variable, name string) error {
 	// case: variable already has value initialized, so skip
 	if v.HasValue() {
 		return nil
 	}
 
 	// case: read and set variable value from process environment
 	envVarName := bundleVarPrefix + name
-	if val, ok := os.LookupEnv(envVarName); ok {
+	if val, ok := env.Lookup(ctx, envVarName); ok {
 		err := v.Set(val)
 		if err != nil {
 			return fmt.Errorf(`failed to assign value "%s" to variable %s from environment variable %s with error: %w`, val, name, envVarName, err)
@@ -54,7 +54,7 @@ func setVariable(v *variable.Variable, name string) error {
 
 func (m *setVariables) Apply(ctx context.Context, b *bundle.Bundle) error {
 	for name, variable := range b.Config.Variables {
-		err := setVariable(variable, name)
+		err := setVariable(ctx, variable, name)
 		if err != nil {
 			return err
 		}

bundle/config/mutator/set_variables_test.go

@@ -21,7 +21,7 @@ func TestSetVariableFromProcessEnvVar(t *testing.T) {
 	// set value for variable as an environment variable
 	t.Setenv("BUNDLE_VAR_foo", "process-env")
 
-	err := setVariable(&variable, "foo")
+	err := setVariable(context.Background(), &variable, "foo")
 	require.NoError(t, err)
 	assert.Equal(t, *variable.Value, "process-env")
 }
@@ -33,7 +33,7 @@ func TestSetVariableUsingDefaultValue(t *testing.T) {
 		Default:     &defaultVal,
 	}
 
-	err := setVariable(&variable, "foo")
+	err := setVariable(context.Background(), &variable, "foo")
 	require.NoError(t, err)
 	assert.Equal(t, *variable.Value, "default")
 }
@@ -49,7 +49,7 @@ func TestSetVariableWhenAlreadyAValueIsAssigned(t *testing.T) {
 
 	// since a value is already assigned to the variable, it would not be overridden
 	// by the default value
-	err := setVariable(&variable, "foo")
+	err := setVariable(context.Background(), &variable, "foo")
 	require.NoError(t, err)
 	assert.Equal(t, *variable.Value, "assigned-value")
 }
@@ -68,7 +68,7 @@ func TestSetVariableEnvVarValueDoesNotOverridePresetValue(t *testing.T) {
 
 	// since a value is already assigned to the variable, it would not be overridden
 	// by the value from environment
-	err := setVariable(&variable, "foo")
+	err := setVariable(context.Background(), &variable, "foo")
 	require.NoError(t, err)
 	assert.Equal(t, *variable.Value, "assigned-value")
 }
@@ -79,7 +79,7 @@ func TestSetVariablesErrorsIfAValueCouldNotBeResolved(t *testing.T) {
 	}
 
 	// fails because we could not resolve a value for the variable
-	err := setVariable(&variable, "foo")
+	err := setVariable(context.Background(), &variable, "foo")
 	assert.ErrorContains(t, err, "no value assigned to required variable foo. Assignment can be done through the \"--var\" flag or by setting the BUNDLE_VAR_foo environment variable")
 }
 

bundle/config/mutator/trampoline.go

@@ -43,16 +43,16 @@ func (m *trampoline) Name() string {
 func (m *trampoline) Apply(ctx context.Context, b *bundle.Bundle) error {
 	tasks := m.functions.GetTasks(b)
 	for _, task := range tasks {
-		err := m.generateNotebookWrapper(b, task)
+		err := m.generateNotebookWrapper(ctx, b, task)
 		if err != nil {
 			return err
 		}
 	}
 	return nil
 }
 
-func (m *trampoline) generateNotebookWrapper(b *bundle.Bundle, task TaskWithJobKey) error {
-	internalDir, err := b.InternalDir()
+func (m *trampoline) generateNotebookWrapper(ctx context.Context, b *bundle.Bundle, task TaskWithJobKey) error {
+	internalDir, err := b.InternalDir(ctx)
 	if err != nil {
 		return err
 	}

bundle/config/mutator/trampoline_test.go

@@ -83,7 +83,7 @@ func TestGenerateTrampoline(t *testing.T) {
 	err := bundle.Apply(ctx, b, trampoline)
 	require.NoError(t, err)
 
-	dir, err := b.InternalDir()
+	dir, err := b.InternalDir(ctx)
 	require.NoError(t, err)
 	filename := filepath.Join(dir, "notebook_test_to_trampoline.py")
 

bundle/deploy/files/sync.go

@@ -9,12 +9,12 @@ import (
 )
 
 func getSync(ctx context.Context, b *bundle.Bundle) (*sync.Sync, error) {
-	cacheDir, err := b.CacheDir()
+	cacheDir, err := b.CacheDir(ctx)
 	if err != nil {
 		return nil, fmt.Errorf("cannot get bundle cache directory: %w", err)
 	}
 
-	includes, err := b.GetSyncIncludePatterns()
+	includes, err := b.GetSyncIncludePatterns(ctx)
 	if err != nil {
 		return nil, fmt.Errorf("cannot get list of sync includes: %w", err)
 	}