Using auth.ProcessEnv()

Signed-off-by: Varun Deep Saini <varun.23bcs10048@ms.sst.scaler.com>

Author: Varun Deep Saini

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/out.test.toml

@@ -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/output.txt

@@ -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/resources.py

@@ -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/script

@@ -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/auth-env-passthrough/test.toml

@@ -105,7 +105,7 @@ type runPythonMutatorOpts struct {
 	bundleRootPath string
 	pythonPath     string
 	loadLocations  bool
-	authEnvs       map[string]string
+	authEnvs       []string
 }
 
 // getOpts adapts deprecated PyDABs and upcoming Python configuration
@@ -224,7 +224,7 @@ 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())
+	authEnvs := auth.ProcessEnv(b.WorkspaceClient().Config)
 
 	err = b.Config.Mutate(func(leftRoot dyn.Value) (dyn.Value, error) {
 		pythonPath, err := detectExecutable(ctx, opts.venvPath)
@@ -369,7 +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),
+		process.WithEnviron(opts.authEnvs),
 	)
 	if processErr != nil {
 		logger.Debugf(ctx, "python mutator process failed: %s", processErr)

bundle/config/mutator/python/python_mutator.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