diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index 9d26552d64..e17d3d725c 100644 --- a/NEXT_CHANGELOG.md +++ b/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 diff --git a/acceptance/bundle/python/auth-env-passthrough/databricks.yml b/acceptance/bundle/python/auth-env-passthrough/databricks.yml new file mode 100644 index 0000000000..4688f03069 --- /dev/null +++ b/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" diff --git a/acceptance/bundle/python/auth-env-passthrough/out.test.toml b/acceptance/bundle/python/auth-env-passthrough/out.test.toml new file mode 100644 index 0000000000..257a5e7686 --- /dev/null +++ b/acceptance/bundle/python/auth-env-passthrough/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-requirements requirements-latest.txt --no-cache"] diff --git a/acceptance/bundle/python/auth-env-passthrough/output.txt b/acceptance/bundle/python/auth-env-passthrough/output.txt new file mode 100644 index 0000000000..0cce75e748 --- /dev/null +++ b/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" +} diff --git a/acceptance/bundle/python/auth-env-passthrough/resources.py b/acceptance/bundle/python/auth-env-passthrough/resources.py new file mode 100644 index 0000000000..f6fd5950b3 --- /dev/null +++ b/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 diff --git a/acceptance/bundle/python/auth-env-passthrough/script b/acceptance/bundle/python/auth-env-passthrough/script new file mode 100644 index 0000000000..c03017aab4 --- /dev/null +++ b/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__ diff --git a/acceptance/bundle/python/auth-env-passthrough/test.toml b/acceptance/bundle/python/auth-env-passthrough/test.toml new file mode 100644 index 0000000000..ab77024283 --- /dev/null +++ b/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", +] 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..284be8e141 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 []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) 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 }) diff --git a/libs/process/opts.go b/libs/process/opts.go index dd06675168..70fdc7ebb2 100644 --- a/libs/process/opts.go +++ b/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