Clean up Python mutator temp dir after run

bundle/config/mutator/python/python_mutator.go

@@ -238,10 +238,11 @@ func (m *pythonMutator) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagno
 			return dyn.InvalidValue, fmt.Errorf("failed to get Python interpreter path: %w", err)
 		}
 
-		cacheDir, err := createCacheDir(ctx)
+		cacheDir, cleanup, err := createCacheDir(ctx)
 		if err != nil {
 			return dyn.InvalidValue, fmt.Errorf("failed to create cache dir: %w", err)
 		}
+		defer cleanup()
 
 		rightRoot, diags := m.runPythonMutator(ctx, leftRoot, runPythonMutatorOpts{
 			cacheDir:       cacheDir,
@@ -315,7 +316,8 @@ func (m *pythonMutator) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagno
 	return mutateDiags
 }
 
-func createCacheDir(ctx context.Context) (string, error) {
+// createCacheDir returns the directory for input/output files of the Python subprocess, and a cleanup function.
+func createCacheDir(ctx context.Context) (string, func(), error) {
 	// b.LocalStateDir doesn't work because target isn't yet selected
 
 	// support the same env variable as in b.LocalStateDir
@@ -325,13 +327,20 @@ func createCacheDir(ctx context.Context) (string, error) {
 
 		err := os.MkdirAll(cacheDir, 0o700)
 		if err != nil {
-			return "", err
+			return "", nil, err
 		}
 
-		return cacheDir, nil
+		// the user picked this location explicitly, keep the files for inspection
+		return cacheDir, func() {}, nil
+	}
+
+	cacheDir, err := os.MkdirTemp("", "-python")
+	if err != nil {
+		return "", nil, err
 	}
 
-	return os.MkdirTemp("", "-python")
+	// input.json contains the full serialized bundle configuration; don't leave it behind in the system temp dir
+	return cacheDir, func() { _ = os.RemoveAll(cacheDir) }, nil
 }
 
 func (m *pythonMutator) runPythonMutator(ctx context.Context, root dyn.Value, opts runPythonMutatorOpts) (dyn.Value, diag.Diagnostics) {

bundle/config/mutator/python/python_mutator_test.go

@@ -20,6 +20,7 @@ import (
 
 	"github.com/databricks/cli/bundle"
 	"github.com/databricks/cli/bundle/config"
+	"github.com/databricks/cli/internal/testutil"
 	"github.com/databricks/cli/libs/cmdio"
 	"github.com/databricks/cli/libs/process"
 	"github.com/stretchr/testify/assert"
@@ -471,6 +472,34 @@ func TestStrictNormalize(t *testing.T) {
 	assert.True(t, strictDiags.HasError())
 }
 
+func TestCreateCacheDir(t *testing.T) {
+	testutil.CleanupEnvironment(t)
+
+	t.Run("DATABRICKS_BUNDLE_TMP is set", func(t *testing.T) {
+		tempDir := t.TempDir()
+		t.Setenv(env.TempDirVariable, tempDir)
+
+		cacheDir, cleanup, err := createCacheDir(t.Context())
+		require.NoError(t, err)
+		require.Equal(t, filepath.Join(tempDir, "default", "python"), cacheDir)
+
+		cleanup()
+
+		// user-specified directories are kept for inspection
+		assert.DirExists(t, cacheDir)
+	})
+
+	t.Run("DATABRICKS_BUNDLE_TMP is not set", func(t *testing.T) {
+		cacheDir, cleanup, err := createCacheDir(t.Context())
+		require.NoError(t, err)
+		require.DirExists(t, cacheDir)
+
+		cleanup()
+
+		assert.NoDirExists(t, cacheDir)
+	})
+}
+
 func TestExplainProcessErr(t *testing.T) {
 	stderr := "/home/test/.venv/bin/python3: Error while finding module specification for 'databricks.bundles.build' (ModuleNotFoundError: No module named 'databricks')\n"
 	expected := `/home/test/.venv/bin/python3: Error while finding module specification for 'databricks.bundles.build' (ModuleNotFoundError: No module named 'databricks')
@@ -501,7 +530,7 @@ func withProcessStub(t *testing.T, args []string, output, diagnostics, locations
 	t.Setenv(env.TempDirVariable, t.TempDir())
 
 	// after we override env variable, we always get the same cache dir as mutator
-	cacheDir, err := createCacheDir(ctx)
+	cacheDir, _, err := createCacheDir(ctx)
 	require.NoError(t, err)
 
 	inputPath := filepath.Join(cacheDir, "input.json")