[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 versions; pass --serverless explicitly to disambiguate", id)
+			}
+		}
+		return "", true, version, nil
 	}
 
 	if len(job.Settings.JobClusters) > 0 {
@@ -78,3 +93,22 @@ 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.
+//
+// The version can arrive in either of two fields. environment_version is the
+// current one; client is its deprecated predecessor ("Use environment_version
+// instead") and is still what some jobs pin. Reading both means the v4 fallback
+// and the divergence guard observe whichever field actually carries the pin,
+// rather than treating a client-pinned job as unversioned. base_environment is
+// deliberately ignored: it is a path/ID, not a version.
+func environmentVersion(e jobs.JobEnvironment) string {
+	if e.Spec == nil {
+		return ""
+	}
+	if e.Spec.EnvironmentVersion != "" {
+		return e.Spec.EnvironmentVersion
+	}
+	return e.Spec.Client
+}

cmd/localenv/compute_test.go

@@ -0,0 +1,32 @@
+package localenv
+
+import (
+	"testing"
+
+	"github.com/databricks/databricks-sdk-go/service/compute"
+	"github.com/databricks/databricks-sdk-go/service/jobs"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestEnvironmentVersion(t *testing.T) {
+	cases := []struct {
+		name string
+		env  jobs.JobEnvironment
+		want string
+	}{
+		{"nil spec", jobs.JobEnvironment{}, ""},
+		{"environment_version", jobs.JobEnvironment{Spec: &compute.Environment{EnvironmentVersion: "3"}}, "3"},
+		// client is the deprecated predecessor of environment_version; some jobs
+		// still pin via it, so it must be read when environment_version is empty.
+		{"client fallback", jobs.JobEnvironment{Spec: &compute.Environment{Client: "2"}}, "2"},
+		// environment_version wins when both are present.
+		{"environment_version wins", jobs.JobEnvironment{Spec: &compute.Environment{EnvironmentVersion: "4", Client: "2"}}, "4"},
+		// base_environment is a path/ID, not a version, and is ignored.
+		{"base_environment ignored", jobs.JobEnvironment{Spec: &compute.Environment{BaseEnvironment: "/Workspace/env.yaml"}}, ""},
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			assert.Equal(t, tc.want, environmentVersion(tc.env))
+		})
+	}
+}

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/target_test.go

@@ -90,6 +90,27 @@ func TestResolveJobClassicUsesSparkVersionReturn(t *testing.T) {
 	assert.Equal(t, "dbr/15.4.x-scala2.12", ti.EnvKey)
 }
 
+func TestResolveJobServerlessUsesRecordedVersion(t *testing.T) {
+	// A serverless job (isServerless=true) pins its serverless version via the
+	// third "recorded version" return; ResolveTarget must map it to the matching
+	// serverless-vN rather than the classic dbr path.
+	c := jobStubCompute{isServerless: true, version: "3"}
+	ti, err := ResolveTarget(t.Context(), TargetFlags{Job: "42"}, c, BundleTarget{})
+	require.NoError(t, err)
+	assert.Equal(t, "job", ti.Source)
+	assert.Empty(t, ti.SparkVersion)
+	assert.Equal(t, "serverless/serverless-v3", ti.EnvKey)
+}
+
+func TestResolveJobServerlessEmptyVersionFallsBackToV4(t *testing.T) {
+	// When the job records no serverless version, ResolveTarget defaults to v4
+	// (documented stand-in), matching the bundle serverless path.
+	c := jobStubCompute{isServerless: true, version: ""}
+	ti, err := ResolveTarget(t.Context(), TargetFlags{Job: "42"}, c, BundleTarget{})
+	require.NoError(t, err)
+	assert.Equal(t, "serverless/serverless-v4", ti.EnvKey)
+}
+
 func TestValidateTargetFlagsMutuallyExclusive(t *testing.T) {
 	assert.Error(t, ValidateTargetFlags(TargetFlags{Cluster: "a", Serverless: "v4"}))
 	assert.NoError(t, ValidateTargetFlags(TargetFlags{Cluster: "a"}))

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,18 +308,46 @@ 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
+	}
+	// Name the OS-specific installer URL that installUv will actually fetch, so
+	// the prompt is transparent about what runs (install.ps1 on Windows).
+	ok, err := cmdio.AskYesOrNo(ctx, "uv is not installed. Download and run the official uv installer ("+uvInstallerURL()+")?")
+	return err == nil && ok
+}
+
+// uvInstallerURL returns the URL of the official uv installer script that
+// installUv fetches for the current OS.
+func uvInstallerURL() string {
+	if runtime.GOOS == "windows" {
+		return "https://astral.sh/uv/install.ps1"
+	}
+	return "https://astral.sh/uv/install.sh"
+}
+
 // 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.
 // https://docs.astral.sh/uv/getting-started/installation/
 func installUv(ctx context.Context) error {
 	var cmd []string
 	if runtime.GOOS == "windows" {
-		// https://astral.sh/uv/install.ps1
-		cmd = []string{"powershell", "-ExecutionPolicy", "ByPass", "-Command", "irm https://astral.sh/uv/install.ps1 | iex"}
+		cmd = []string{"powershell", "-ExecutionPolicy", "ByPass", "-Command", "irm " + uvInstallerURL() + " | iex"}
 	} else {
-		// https://astral.sh/uv/install.sh
-		cmd = []string{"sh", "-c", "curl -LsSf https://astral.sh/uv/install.sh | sh"}
+		cmd = []string{"sh", "-c", "curl -LsSf " + uvInstallerURL() + " | sh"}
 	}
 	// This downloads and runs a remote installer that mutates the user's machine
 	// (~/.local/bin), so record exactly what ran before it fires — visible under