diff --git a/acceptance/bundle/python/workspace-client-auth/databricks.yml b/acceptance/bundle/python/workspace-client-auth/databricks.yml new file mode 100644 index 0000000000..edd9f1b947 --- /dev/null +++ b/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" diff --git a/acceptance/bundle/python/workspace-client-auth/mutators.py b/acceptance/bundle/python/workspace-client-auth/mutators.py new file mode 100644 index 0000000000..a7fa05c2b1 --- /dev/null +++ b/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}") diff --git a/acceptance/bundle/python/workspace-client-auth/out.test.toml b/acceptance/bundle/python/workspace-client-auth/out.test.toml new file mode 100644 index 0000000000..f5f882b98b --- /dev/null +++ b/acceptance/bundle/python/workspace-client-auth/out.test.toml @@ -0,0 +1,6 @@ +Local = true +Cloud = false + +[EnvMatrix] + DATABRICKS_BUNDLE_ENGINE = ["terraform", "direct"] + UV_ARGS = ["--with databricks-bundles==0.266.0 --with databricks-sdk", "--with-requirements requirements-latest.txt --with databricks-sdk --no-cache"] diff --git a/acceptance/bundle/python/workspace-client-auth/output.txt b/acceptance/bundle/python/workspace-client-auth/output.txt new file mode 100644 index 0000000000..3d84c0b3b3 --- /dev/null +++ b/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]" +} diff --git a/acceptance/bundle/python/workspace-client-auth/script b/acceptance/bundle/python/workspace-client-auth/script new file mode 100644 index 0000000000..9828e8c332 --- /dev/null +++ b/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__ diff --git a/acceptance/bundle/python/workspace-client-auth/test.toml b/acceptance/bundle/python/workspace-client-auth/test.toml new file mode 100644 index 0000000000..df18324a55 --- /dev/null +++ b/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", +] diff --git a/bundle/config/mutator/python/python_mutator.go b/bundle/config/mutator/python/python_mutator.go index c20e172c00..d4aec4a84b 100644 --- a/bundle/config/mutator/python/python_mutator.go +++ b/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) diff --git a/bundle/config/mutator/python/python_mutator_test.go b/bundle/config/mutator/python/python_mutator_test.go index 285d1b3b87..2d596c66dd 100644 --- a/bundle/config/mutator/python/python_mutator_test.go +++ b/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 })