[VPEX][6b] local-env: address #5832 follow-up review comments

cmd/localenv/compute.go

@@ -6,6 +6,7 @@ import (
 	"strconv"
 
 	databricks "github.com/databricks/databricks-sdk-go"
+	"github.com/databricks/databricks-sdk-go/service/jobs"
 )
 
 // sdkCompute adapts the Databricks SDK to the localenv.ComputeClient interface.
@@ -57,7 +58,21 @@ func (c sdkCompute) GetJobSparkVersion(ctx context.Context, jobID string) (spark
 
 	// Serverless jobs have Environments populated; classic compute uses JobClusters.
 	if len(job.Settings.Environments) > 0 {
-		return "", true, "", nil
+		// The serverless environment version (e.g. "4") is recorded on the job's
+		// environment spec, unlike the bundle path where it is unavailable. Return
+		// it so ResolveTarget pins the matching serverless-vN instead of defaulting
+		// to v4. An empty version (older jobs) falls back to v4 in ResolveTarget.
+		version := environmentVersion(job.Settings.Environments[0])
+		// Tasks can reference any environment_key, so if the job's environments do
+		// not all share one version there is no single correct local environment
+		// (mirrors the job-cluster check below). Refuse rather than guess from the
+		// first. A pinned-vs-unpinned mix is also ambiguous, so compare raw values.
+		for _, e := range job.Settings.Environments[1:] {
+			if environmentVersion(e) != version {
+				return "", false, "", fmt.Errorf("job %d has serverless environments with differing environment_version; pass --serverless explicitly to disambiguate", id)
+			}
+		}
+		return "", true, version, nil
 	}
 
 	if len(job.Settings.JobClusters) > 0 {
@@ -78,3 +93,12 @@ func (c sdkCompute) GetJobSparkVersion(ctx context.Context, jobID string) (spark
 
 	return "", false, "", fmt.Errorf("could not determine compute for job %d from its environments or job clusters (task-level compute is not supported); pass --cluster or --serverless explicitly", id)
 }
+
+// environmentVersion returns the serverless environment_version recorded on a
+// job environment, or "" when the spec or version is absent.
+func environmentVersion(e jobs.JobEnvironment) string {
+	if e.Spec == nil {
+		return ""
+	}
+	return e.Spec.EnvironmentVersion
+}

cmd/localenv/sync.go

@@ -95,7 +95,14 @@ func runPipeline(cmd *cobra.Command) error {
 	}
 	cacheDir = filepath.Join(cacheDir, "databricks", "localenv")
 
-	bt := bundleTarget(cmd)
+	// The bundle is only a fallback: ResolveTarget consults it solely when no
+	// explicit --cluster/--serverless/--job flag is set. Skip the bundle load
+	// entirely when a flag is present — it would otherwise re-run TryConfigureBundle
+	// (a second full load) and re-print any bundle load-time diagnostics for nothing.
+	var bt libslocalenv.BundleTarget
+	if cluster == "" && serverless == "" && job == "" {
+		bt = bundleTarget(cmd)
+	}
 
 	w := cmdctx.WorkspaceClient(ctx)
 	p := &libslocalenv.Pipeline{

libs/localenv/target.go

@@ -93,9 +93,9 @@ func ResolveTarget(ctx context.Context, f TargetFlags, c ComputeClient, bt Bundl
 			return nil, NewError(ErrResolve, err, "resolving job %s", f.Job)
 		}
 		if isServerless {
-			// Default to v4 when the job is serverless; the serverless env version
-			// is not recorded in the bundle/project (documented stand-in from the
-			// original script, spec §4.3).
+			// Use the job's recorded serverless environment version when present;
+			// fall back to v4 when the job did not pin one (documented stand-in from
+			// the original script, spec §4.3).
 			v := version
 			if v == "" {
 				v = "v4"

libs/localenv/uv.go

@@ -12,27 +12,27 @@ import (
 	"runtime"
 	"strings"
 
+	"github.com/databricks/cli/libs/cmdio"
 	"github.com/databricks/cli/libs/env"
 	"github.com/databricks/cli/libs/log"
 	"github.com/databricks/cli/libs/process"
 )
 
+// EnvAutoInstallUv opts into installing uv without an interactive prompt. It
+// exists so non-interactive runs (CI, IDE integrations) can allow the install
+// that would otherwise be declined for lack of a TTY. Any truthy value enables it.
+const EnvAutoInstallUv = "DATABRICKS_LOCALENV_AUTO_INSTALL_UV"
+
 // uvManager implements PackageManager using the uv tool.
 // https://docs.astral.sh/uv/
 type uvManager struct {
 	bin string
 }
 
-// newUvManager returns a uvManager whose binary path is resolved lazily via
-// EnsureAvailable.
-func newUvManager() *uvManager {
-	return &uvManager{}
-}
-
-// NewUvManager returns a PackageManager backed by the uv tool.
-// This is the exported constructor for use outside this package.
+// NewUvManager returns a PackageManager backed by the uv tool. The binary path
+// is resolved lazily via EnsureAvailable.
 func NewUvManager() PackageManager {
-	return newUvManager()
+	return &uvManager{}
 }
 
 // Name returns "uv".
@@ -47,6 +47,10 @@ func (m *uvManager) Name() string {
 func (m *uvManager) EnsureAvailable(ctx context.Context) (string, error) {
 	bin, err := discoverUv(ctx)
 	if err != nil {
+		if !confirmUvInstall(ctx) {
+			return "", NewError(ErrUvMissing, nil,
+				"uv is required but not installed; install it (https://docs.astral.sh/uv/getting-started/installation/) or set %s=1 to let this command install it for you", EnvAutoInstallUv)
+		}
 		if installErr := installUv(ctx); installErr != nil {
 			return "", NewError(ErrUvMissing, installErr, "uv installation failed")
 		}
@@ -66,17 +70,24 @@ func (m *uvManager) EnsureAvailable(ctx context.Context) (string, error) {
 	return strings.TrimSpace(version), nil
 }
 
+// runUv runs the uv binary with args in dir, injecting UV_INDEX_URL from pip.conf
+// when appropriate. An empty dir runs in the current working directory
+// (process.WithDir("") is a no-op). The index-url is injected only when
+// resolveIndexURL returns non-empty; it returns "" when UV_INDEX_URL is already
+// set, so an explicit value in the environment is never clobbered.
+func (m *uvManager) runUv(ctx context.Context, args []string, dir string) error {
+	if indexURL := m.resolveIndexURL(ctx); indexURL != "" {
+		_, err := process.Background(ctx, args, process.WithDir(dir), process.WithEnv("UV_INDEX_URL", indexURL))
+		return err
+	}
+	_, err := process.Background(ctx, args, process.WithDir(dir))
+	return err
+}
+
 // EnsurePython installs the requested Python minor version via uv.
 func (m *uvManager) EnsurePython(ctx context.Context, minor string) error {
 	args := append([]string{m.bin}, m.pythonInstallArgs(minor)...)
-	indexURL := m.resolveIndexURL(ctx)
-	var err error
-	if indexURL != "" {
-		_, err = process.Background(ctx, args, process.WithEnv("UV_INDEX_URL", indexURL))
-	} else {
-		_, err = process.Background(ctx, args)
-	}
-	if err != nil {
+	if err := m.runUv(ctx, args, ""); err != nil {
 		return uvFailure(ErrPythonInstall, err, "uv python install "+minor)
 	}
 	return nil
@@ -85,14 +96,7 @@ func (m *uvManager) EnsurePython(ctx context.Context, minor string) error {
 // Provision runs `uv sync` inside projectDir to install project dependencies.
 func (m *uvManager) Provision(ctx context.Context, projectDir string) error {
 	args := append([]string{m.bin}, m.syncArgs()...)
-	indexURL := m.resolveIndexURL(ctx)
-	var err error
-	if indexURL != "" {
-		_, err = process.Background(ctx, args, process.WithDir(projectDir), process.WithEnv("UV_INDEX_URL", indexURL))
-	} else {
-		_, err = process.Background(ctx, args, process.WithDir(projectDir))
-	}
-	if err != nil {
+	if err := m.runUv(ctx, args, projectDir); err != nil {
 		return uvFailure(ErrProvision, err, "uv sync")
 	}
 	return nil
@@ -117,14 +121,7 @@ func venvPython(projectDir string) string {
 // activated.
 func (m *uvManager) PostProvision(ctx context.Context, projectDir string) error {
 	args := append([]string{m.bin}, m.pipSeedArgs(venvPython(projectDir))...)
-	indexURL := m.resolveIndexURL(ctx)
-	var err error
-	if indexURL != "" {
-		_, err = process.Background(ctx, args, process.WithDir(projectDir), process.WithEnv("UV_INDEX_URL", indexURL))
-	} else {
-		_, err = process.Background(ctx, args, process.WithDir(projectDir))
-	}
-	if err != nil {
+	if err := m.runUv(ctx, args, projectDir); err != nil {
 		return uvFailure(ErrProvision, err, "uv pip seed")
 	}
 	return nil
@@ -136,12 +133,16 @@ func (m *uvManager) PostProvision(ctx context.Context, projectDir string) error
 // error: PackageNotFoundError is caught so the probe never fails just because the
 // package is absent. The caller decides whether an empty version is acceptable.
 func (m *uvManager) Validate(ctx context.Context, projectDir string) (string, string, error) {
+	// Each value is printed with a unique prefix so parsing greps for the prefix
+	// rather than relying on line position: any stray line uv or the interpreter
+	// writes to stdout (e.g. a warning) would otherwise shift a positional parse.
+	// A missing databricks-connect prints an empty DBC: value, not an error.
 	pyCode := `import sys, importlib.metadata
-print(f"{sys.version_info.major}.{sys.version_info.minor}")
+print(f"` + validatePyPrefix + `{sys.version_info.major}.{sys.version_info.minor}")
 try:
-    print(importlib.metadata.version("databricks-connect"))
+    print("` + validateDBCPrefix + `" + importlib.metadata.version("databricks-connect"))
 except importlib.metadata.PackageNotFoundError:
-    print("")`
+    print("` + validateDBCPrefix + `")`
 	// --no-project runs the interpreter from the created .venv without re-resolving/syncing
 	// the project's declared dependencies, so validation observes exactly what was installed.
 	out, err := process.Background(ctx,
@@ -151,16 +152,32 @@ except importlib.metadata.PackageNotFoundError:
 	if err != nil {
 		return "", "", uvFailure(ErrValidate, err, "uv run python validation")
 	}
-	lines := strings.Split(strings.TrimSpace(out), "\n")
-	if len(lines) < 1 || strings.TrimSpace(lines[0]) == "" {
+	pyVer, ok := lineWithPrefix(out, validatePyPrefix)
+	if !ok || pyVer == "" {
 		return "", "", NewError(ErrValidate, nil, "unexpected output from uv run: %q", out)
 	}
-	// The databricks-connect line is empty when the package is not installed.
-	dbcVer := ""
-	if len(lines) >= 2 {
-		dbcVer = strings.TrimSpace(lines[len(lines)-1])
+	// The databricks-connect value is empty when the package is not installed.
+	dbcVer, _ := lineWithPrefix(out, validateDBCPrefix)
+	return pyVer, dbcVer, nil
+}
+
+// Validation output prefixes: uv run's stdout is grepped for these rather than
+// parsed positionally, so extra lines from uv or the interpreter don't break it.
+const (
+	validatePyPrefix  = "PYVER:"
+	validateDBCPrefix = "DBCVER:"
+)
+
+// lineWithPrefix returns the trimmed remainder of the first line in out that
+// starts with prefix, and whether such a line was found.
+func lineWithPrefix(out, prefix string) (string, bool) {
+	for line := range strings.SplitSeq(out, "\n") {
+		line = strings.TrimSpace(line)
+		if after, ok := strings.CutPrefix(line, prefix); ok {
+			return strings.TrimSpace(after), true
+		}
 	}
-	return strings.TrimSpace(lines[0]), dbcVer, nil
+	return "", false
 }
 
 // syncArgs returns the argument slice for `uv sync` (without the binary).
@@ -291,6 +308,25 @@ func uvFailure(code ErrorCode, err error, action string) *PipelineError {
 	return NewError(code, err, "%s", msg)
 }
 
+// confirmUvInstall reports whether the caller has consented to installUv running
+// a remote installer that mutates the machine. The EnvAutoInstallUv opt-in wins
+// outright (for CI / IDE integrations); otherwise an interactive session is
+// prompted, and a non-interactive session without the opt-in declines rather than
+// silently downloading and executing an installer.
+func confirmUvInstall(ctx context.Context) bool {
+	if optIn, ok := env.GetBool(ctx, EnvAutoInstallUv); ok && optIn {
+		return true
+	}
+	// EnsureAvailable is a library entry point reachable with a context that has
+	// no cmdio (e.g. Pipeline built with context.Background()); IsPromptSupported
+	// would panic there. Treat a missing cmdio as non-interactive and decline.
+	if !cmdio.HasIO(ctx) || !cmdio.IsPromptSupported(ctx) {
+		return false
+	}
+	ok, err := cmdio.AskYesOrNo(ctx, "uv is not installed. Download and run the official uv installer (https://astral.sh/uv/install.sh)?")
+	return err == nil && ok
+}
+
 // installUv runs the official uv installer for the current OS. Unix uses the
 // shell installer; Windows uses the PowerShell installer, because the Unix
 // `sh`/`curl` pipeline is not available in a default Windows shell.

libs/localenv/uv_test.go

@@ -8,6 +8,7 @@ import (
 	"runtime"
 	"testing"
 
+	"github.com/databricks/cli/libs/cmdio"
 	"github.com/databricks/cli/libs/env"
 	"github.com/databricks/cli/libs/process"
 	"github.com/stretchr/testify/assert"
@@ -221,3 +222,54 @@ func TestUvFailureIncludesStderr(t *testing.T) {
 		assert.Equal(t, "uv sync failed", pe.Msg)
 	})
 }
+
+func TestConfirmUvInstall(t *testing.T) {
+	t.Run("opt_in_env_var_consents_without_prompt", func(t *testing.T) {
+		// Non-interactive context, but the opt-in env var grants consent.
+		ctx := env.Set(t.Context(), EnvAutoInstallUv, "1")
+		assert.True(t, confirmUvInstall(ctx))
+	})
+
+	t.Run("non_interactive_without_opt_in_declines", func(t *testing.T) {
+		// SetupTest defaults to PromptSupported=false, i.e. non-interactive.
+		ctx, _ := cmdio.SetupTest(t.Context(), cmdio.TestOptions{})
+		assert.False(t, confirmUvInstall(ctx))
+	})
+
+	t.Run("missing_cmdio_declines_without_panic", func(t *testing.T) {
+		// A context with no cmdio (library entry point) must not panic in
+		// IsPromptSupported; it declines like any other non-interactive run.
+		assert.NotPanics(t, func() {
+			assert.False(t, confirmUvInstall(t.Context()))
+		})
+	})
+
+	t.Run("falsey_opt_in_does_not_consent_when_non_interactive", func(t *testing.T) {
+		ctx, _ := cmdio.SetupTest(t.Context(), cmdio.TestOptions{})
+		ctx = env.Set(ctx, EnvAutoInstallUv, "0")
+		assert.False(t, confirmUvInstall(ctx))
+	})
+}
+
+func TestLineWithPrefix(t *testing.T) {
+	// A stray leading line (as uv might emit) must not shift the parse: the
+	// value is located by prefix, not position.
+	out := "warning: something\nPYVER:3.12\nDBCVER:17.2.0\n"
+
+	pyVer, ok := lineWithPrefix(out, validatePyPrefix)
+	assert.True(t, ok)
+	assert.Equal(t, "3.12", pyVer)
+
+	dbcVer, ok := lineWithPrefix(out, validateDBCPrefix)
+	assert.True(t, ok)
+	assert.Equal(t, "17.2.0", dbcVer)
+
+	// An empty DBC value (databricks-connect absent) is found but blank.
+	empty, ok := lineWithPrefix("PYVER:3.12\nDBCVER:\n", validateDBCPrefix)
+	assert.True(t, ok)
+	assert.Empty(t, empty)
+
+	// A missing prefix reports not-found rather than a wrong line.
+	_, ok = lineWithPrefix("PYVER:3.12\n", validateDBCPrefix)
+	assert.False(t, ok)
+}