Pass CLI auth to Python mutator subprocess

NEXT_CHANGELOG.md

@@ -7,6 +7,7 @@
 ### CLI
 
 ### Bundles
+* Pass authentication environment variables to Python mutator process ([#4172](https://github.com/databricks/cli/pull/4172))
 
 ### Dependency updates
 

acceptance/bundle/python/auth-env-passthrough/databricks.yml

@@ -0,0 +1,8 @@
+bundle:
+  name: auth_env_test
+
+sync: {paths: []} # don't need to copy files
+
+python:
+  resources:
+    - "resources:load_resources"

acceptance/bundle/python/auth-env-passthrough/output.txt

@@ -0,0 +1,5 @@
+
+>>> uv run [UV_ARGS] -q [CLI] bundle validate --output json
+{
+  "job_name": "auth_env_test_job"
+}

acceptance/bundle/python/auth-env-passthrough/resources.py

@@ -0,0 +1,32 @@
+import os
+
+from databricks.bundles.core import Bundle, Resources
+
+
+def load_resources(bundle: Bundle) -> Resources:
+    host = os.environ.get("DATABRICKS_HOST", "NOT_SET")
+
+    if host == "NOT_SET":
+        raise ValueError("DATABRICKS_HOST was not passed to Python subprocess")
+
+    if not host.startswith("http://") and not host.startswith("https://"):
+        raise ValueError(f"DATABRICKS_HOST has invalid format: {host}")
+
+    home = os.environ.get("HOME", os.environ.get("USERPROFILE", "NOT_SET"))
+    if home == "NOT_SET":
+        raise ValueError("HOME/USERPROFILE was not passed to Python subprocess")
+
+    profile = os.environ.get("DATABRICKS_CONFIG_PROFILE")
+    if profile is not None:
+        raise ValueError(
+            f"DATABRICKS_CONFIG_PROFILE should have been removed but was: {profile}. "
+            "Conflicting auth env vars must be cleaned to prevent auth issues."
+        )
+
+    resources = Resources()
+    resources.add_job(
+        resource_name="test_job",
+        job={"name": "auth_env_test_job"},
+    )
+
+    return resources

acceptance/bundle/python/auth-env-passthrough/script

@@ -0,0 +1,10 @@
+echo "$DATABRICKS_BUNDLES_WHEEL" > "requirements-latest.txt"
+
+# Verify that auth env vars (DATABRICKS_HOST) are correctly passed to the
+# Python subprocess via auth.ProcessEnv. The Python code validates that
+# both DATABRICKS_HOST and HOME are present in the environment.
+
+trace uv run $UV_ARGS -q $CLI bundle validate --output json | \
+  jq '{job_name: .resources.jobs.test_job.name}'
+
+rm -fr .databricks __pycache__

acceptance/bundle/python/auth-env-passthrough/test.toml

@@ -0,0 +1,19 @@
+# This test verifies that auth env vars from the bundle config are correctly
+# passed to the Python subprocess via auth.ProcessEnv, and that conflicting
+# auth env vars are REMOVED to prevent auth issues.
+
+Local = true
+Cloud = false
+
+Ignore = [
+  "requirements-latest.txt",
+]
+
+[Env]
+DATABRICKS_CONFIG_PROFILE = "conflicting_profile"
+
+[EnvMatrix]
+UV_ARGS = [
+  "--with databricks-bundles==0.266.0",
+  "--with-requirements requirements-latest.txt --no-cache",
+]

acceptance/bundle/python/workspace-client-auth/databricks.yml

@@ -0,0 +1,13 @@
+bundle:
+  name: workspace_client_auth_test
+
+sync: {paths: []}
+
+python:
+  mutators:
+    - "mutators:test_workspace_client"
+
+resources:
+  jobs:
+    test_job:
+      name: "Test Job"

acceptance/bundle/python/workspace-client-auth/mutators.py

@@ -0,0 +1,11 @@
+from dataclasses import replace
+from databricks.bundles.jobs import Job
+from databricks.bundles.core import job_mutator, Bundle
+from databricks.sdk import WorkspaceClient
+
+
+@job_mutator
+def test_workspace_client(bundle: Bundle, job: Job) -> Job:
+    w = WorkspaceClient()
+    user = w.current_user.me()
+    return replace(job, description=f"Validated by user: {user.user_name}")

acceptance/bundle/python/workspace-client-auth/output.txt

@@ -0,0 +1,5 @@
+
+>>> uv run [UV_ARGS] -q [CLI] bundle validate --output json
+{
+  "description": "Validated by user: [USERNAME]"
+}

acceptance/bundle/python/workspace-client-auth/script

@@ -0,0 +1,6 @@
+echo "$DATABRICKS_BUNDLES_WHEEL" > "requirements-latest.txt"
+
+trace uv run $UV_ARGS -q $CLI bundle validate --output json | \
+  jq '{description: .resources.jobs.test_job.description}'
+
+rm -fr .databricks __pycache__

acceptance/bundle/python/workspace-client-auth/test.toml

@@ -0,0 +1,5 @@
+[EnvMatrix]
+UV_ARGS = [
+  "--with databricks-bundles==0.266.0 --with databricks-sdk",
+  "--with-requirements requirements-latest.txt --with databricks-sdk --no-cache",
+]

bundle/config/mutator/python/python_mutator.go

@@ -15,6 +15,7 @@ import (
 
 	"github.com/databricks/cli/bundle/config/mutator/resourcemutator"
 
+	"github.com/databricks/cli/libs/auth"
 	"github.com/databricks/cli/libs/log"
 	"github.com/databricks/cli/libs/logdiag"
 
@@ -104,6 +105,7 @@ type runPythonMutatorOpts struct {
 	bundleRootPath string
 	pythonPath     string
 	loadLocations  bool
+	authEnvs       []string
 }
 
 // getOpts adapts deprecated PyDABs and upcoming Python configuration
@@ -222,6 +224,8 @@ func (m *pythonMutator) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagno
 	var result applyPythonOutputResult
 	mutateDiagsHasError := errors.New("unexpected error")
 
+	authEnvs := auth.ProcessEnv(b.WorkspaceClient().Config)
+
 	err = b.Config.Mutate(func(leftRoot dyn.Value) (dyn.Value, error) {
 		pythonPath, err := detectExecutable(ctx, opts.venvPath)
 		if err != nil {
@@ -238,6 +242,7 @@ func (m *pythonMutator) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagno
 			bundleRootPath: b.BundleRootPath,
 			pythonPath:     pythonPath,
 			loadLocations:  opts.loadLocations,
+			authEnvs:       authEnvs,
 		})
 		mutateDiags = diags
 		if diags.HasError() {
@@ -364,6 +369,7 @@ func (m *pythonMutator) runPythonMutator(ctx context.Context, root dyn.Value, op
 		process.WithDir(opts.bundleRootPath),
 		process.WithStderrWriter(stderrWriter),
 		process.WithStdoutWriter(stdoutWriter),
+		process.WithEnviron(opts.authEnvs),
 	)
 	if processErr != nil {
 		logger.Debugf(ctx, "python mutator process failed: %s", processErr)

bundle/config/mutator/python/python_mutator_test.go

@@ -493,7 +493,73 @@ or activate the environment before running CLI commands:
 	assert.Equal(t, expected, out)
 }
 
+func TestPythonMutator_authEnvVarsPassedToSubprocess(t *testing.T) {
+	withFakeVEnv(t, ".venv")
+
+	b := loadYaml("databricks.yml", `
+experimental:
+  python:
+    venv_path: .venv
+    resources: ["resources:load_resources"]
+resources:
+  jobs:
+    job0:
+      name: job_0
+workspace:
+  host: https://acme.databricks.com`)
+
+	ctx := withProcessStubWithEnvCheck(
+		t,
+		[]string{
+			interpreterPath(".venv"),
+			"-m",
+			"databricks.bundles.build",
+			"--phase",
+			"load_resources",
+		},
+		`{
+			"experimental": {
+				"python": {
+					"resources": ["resources:load_resources"],
+					"venv_path": ".venv"
+				}
+			},
+			"resources": {
+				"jobs": {
+					"job0": {
+						name: "job_0"
+					}
+				}
+			},
+			"workspace": {
+				"host": "https://acme.databricks.com"
+			}
+		}`,
+		"",
+		"",
+		func(cmd *exec.Cmd) {
+			foundHost := false
+			for _, envVar := range cmd.Env {
+				if envVar == "DATABRICKS_HOST=https://acme.databricks.com" {
+					foundHost = true
+					break
+				}
+			}
+			assert.True(t, foundHost, "DATABRICKS_HOST should be passed to subprocess")
+		},
+	)
+
+	mutator := PythonMutator(PythonMutatorPhaseLoadResources)
+	diags := bundle.Apply(ctx, b, mutator)
+
+	assert.NoError(t, diags.Error())
+}
+
 func withProcessStub(t *testing.T, args []string, output, diagnostics, locations string) context.Context {
+	return withProcessStubWithEnvCheck(t, args, output, diagnostics, locations, nil)
+}
+
+func withProcessStubWithEnvCheck(t *testing.T, args []string, output, diagnostics, locations string, envCheck func(*exec.Cmd)) context.Context {
 	ctx := context.Background()
 	ctx, stub := process.WithStub(ctx)
 
@@ -535,6 +601,10 @@ func withProcessStub(t *testing.T, args []string, output, diagnostics, locations
 		err = os.WriteFile(diagnosticsPath, []byte(diagnostics), 0o600)
 		require.NoError(t, err)
 
+		if envCheck != nil {
+			envCheck(actual)
+		}
+
 		return nil
 	})
 

libs/process/opts.go

@@ -31,6 +31,13 @@ func WithEnvs(envs map[string]string) execOption {
 	}
 }
 
+func WithEnviron(envs []string) execOption {
+	return func(ctx context.Context, c *exec.Cmd) error {
+		c.Env = envs
+		return nil
+	}
+}
+
 func WithDir(dir string) execOption {
 	return func(_ context.Context, c *exec.Cmd) error {
 		c.Dir = dir