Pass CLI auth to Python mutator subprocess

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       map[string]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.Env(b.Config.Workspace.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.WithEnvs(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
 	})