Skip to content

Commit 8e7ad0e

Browse files
authored
[VPEX][6b] local-env: address #5832 follow-up review comments (#5936)
Follow-up refinements from the [#5832](#5832) review (all were left as non-blocking there and deferred with a "will address in a follow-up" reply). The `local-env` command remains hidden until #5835, so there is no user-visible changelog entry. ## Changes - **uv install consent** (`libs/localenv/uv.go`) — `EnsureAvailable` no longer runs the remote uv installer (`curl … | sh` / `irm … | iex`) silently. It now requires consent: a truthy `DATABRICKS_LOCALENV_AUTO_INSTALL_UV` opt-in for non-interactive runs (CI/IDE), or an interactive `y/N` prompt via `cmdio.AskYesOrNo`. A non-interactive run without the opt-in returns an actionable error instead of downloading and executing an installer. (The `--debug` log of the exact installer command from #5832 is kept.) - **serverless job version** (`cmd/localenv/compute.go`) — `GetJobSparkVersion` now reads `Environments[0].Spec.EnvironmentVersion`, so a serverless `--job` resolves to its actual `serverless-vN` instead of always defaulting to v4. Empty still falls back to v4. - **double bundle load** (`cmd/localenv/sync.go`) — skip `bundleTarget` when an explicit `--cluster/--serverless/--job` flag is set. `ResolveTarget` only consults the bundle as a fallback, so the second `TryConfigureBundle` load (and its re-printed load-time diagnostics) was wasted for the explicit-flag case. - **robust `Validate` parse** (`libs/localenv/uv.go`) — the `uv run` probe now prints `PYVER:` / `DBCVER:` sentinels and parses by prefix instead of line position, so a stray stdout line from uv doesn't shift the parse. - **cleanup** (`libs/localenv/uv.go`) — fold single-caller `newUvManager` into `NewUvManager`; collapse the triplicated `resolveIndexURL` + conditional `WithEnv` into one `runUv` helper (the conditional stays so an already-set `UV_INDEX_URL` isn't clobbered). ## Tests Adds unit tests for the install consent gate (`TestConfirmUvInstall`) and the sentinel parse (`TestLineWithPrefix`). Full `libs/localenv` suite passes, including under a CI-like environment with `UV_INDEX_URL`/`PIP_INDEX_URL` set. This pull request and its description were written by Isaac.
1 parent dce783d commit 8e7ad0e

7 files changed

Lines changed: 243 additions & 52 deletions

File tree

cmd/localenv/compute.go

Lines changed: 35 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import (
66
"strconv"
77

88
databricks "github.com/databricks/databricks-sdk-go"
9+
"github.com/databricks/databricks-sdk-go/service/jobs"
910
)
1011

1112
// sdkCompute adapts the Databricks SDK to the localenv.ComputeClient interface.
@@ -57,7 +58,21 @@ func (c sdkCompute) GetJobSparkVersion(ctx context.Context, jobID string) (spark
5758

5859
// Serverless jobs have Environments populated; classic compute uses JobClusters.
5960
if len(job.Settings.Environments) > 0 {
60-
return "", true, "", nil
61+
// The serverless environment version (e.g. "4") is recorded on the job's
62+
// environment spec, unlike the bundle path where it is unavailable. Return
63+
// it so ResolveTarget pins the matching serverless-vN instead of defaulting
64+
// to v4. An empty version (older jobs) falls back to v4 in ResolveTarget.
65+
version := environmentVersion(job.Settings.Environments[0])
66+
// Tasks can reference any environment_key, so if the job's environments do
67+
// not all share one version there is no single correct local environment
68+
// (mirrors the job-cluster check below). Refuse rather than guess from the
69+
// first. A pinned-vs-unpinned mix is also ambiguous, so compare raw values.
70+
for _, e := range job.Settings.Environments[1:] {
71+
if environmentVersion(e) != version {
72+
return "", false, "", fmt.Errorf("job %d has serverless environments with differing versions; pass --serverless explicitly to disambiguate", id)
73+
}
74+
}
75+
return "", true, version, nil
6176
}
6277

6378
if len(job.Settings.JobClusters) > 0 {
@@ -78,3 +93,22 @@ func (c sdkCompute) GetJobSparkVersion(ctx context.Context, jobID string) (spark
7893

7994
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)
8095
}
96+
97+
// environmentVersion returns the serverless environment version recorded on a
98+
// job environment, or "" when the spec or version is absent.
99+
//
100+
// The version can arrive in either of two fields. environment_version is the
101+
// current one; client is its deprecated predecessor ("Use environment_version
102+
// instead") and is still what some jobs pin. Reading both means the v4 fallback
103+
// and the divergence guard observe whichever field actually carries the pin,
104+
// rather than treating a client-pinned job as unversioned. base_environment is
105+
// deliberately ignored: it is a path/ID, not a version.
106+
func environmentVersion(e jobs.JobEnvironment) string {
107+
if e.Spec == nil {
108+
return ""
109+
}
110+
if e.Spec.EnvironmentVersion != "" {
111+
return e.Spec.EnvironmentVersion
112+
}
113+
return e.Spec.Client
114+
}

cmd/localenv/compute_test.go

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
package localenv
2+
3+
import (
4+
"testing"
5+
6+
"github.com/databricks/databricks-sdk-go/service/compute"
7+
"github.com/databricks/databricks-sdk-go/service/jobs"
8+
"github.com/stretchr/testify/assert"
9+
)
10+
11+
func TestEnvironmentVersion(t *testing.T) {
12+
cases := []struct {
13+
name string
14+
env jobs.JobEnvironment
15+
want string
16+
}{
17+
{"nil spec", jobs.JobEnvironment{}, ""},
18+
{"environment_version", jobs.JobEnvironment{Spec: &compute.Environment{EnvironmentVersion: "3"}}, "3"},
19+
// client is the deprecated predecessor of environment_version; some jobs
20+
// still pin via it, so it must be read when environment_version is empty.
21+
{"client fallback", jobs.JobEnvironment{Spec: &compute.Environment{Client: "2"}}, "2"},
22+
// environment_version wins when both are present.
23+
{"environment_version wins", jobs.JobEnvironment{Spec: &compute.Environment{EnvironmentVersion: "4", Client: "2"}}, "4"},
24+
// base_environment is a path/ID, not a version, and is ignored.
25+
{"base_environment ignored", jobs.JobEnvironment{Spec: &compute.Environment{BaseEnvironment: "/Workspace/env.yaml"}}, ""},
26+
}
27+
for _, tc := range cases {
28+
t.Run(tc.name, func(t *testing.T) {
29+
assert.Equal(t, tc.want, environmentVersion(tc.env))
30+
})
31+
}
32+
}

cmd/localenv/sync.go

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -95,7 +95,14 @@ func runPipeline(cmd *cobra.Command) error {
9595
}
9696
cacheDir = filepath.Join(cacheDir, "databricks", "localenv")
9797

98-
bt := bundleTarget(cmd)
98+
// The bundle is only a fallback: ResolveTarget consults it solely when no
99+
// explicit --cluster/--serverless/--job flag is set. Skip the bundle load
100+
// entirely when a flag is present — it would otherwise re-run TryConfigureBundle
101+
// (a second full load) and re-print any bundle load-time diagnostics for nothing.
102+
var bt libslocalenv.BundleTarget
103+
if cluster == "" && serverless == "" && job == "" {
104+
bt = bundleTarget(cmd)
105+
}
99106

100107
w := cmdctx.WorkspaceClient(ctx)
101108
p := &libslocalenv.Pipeline{

libs/localenv/target.go

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -93,9 +93,9 @@ func ResolveTarget(ctx context.Context, f TargetFlags, c ComputeClient, bt Bundl
9393
return nil, NewError(ErrResolve, err, "resolving job %s", f.Job)
9494
}
9595
if isServerless {
96-
// Default to v4 when the job is serverless; the serverless env version
97-
// is not recorded in the bundle/project (documented stand-in from the
98-
// original script, spec §4.3).
96+
// Use the job's recorded serverless environment version when present;
97+
// fall back to v4 when the job did not pin one (documented stand-in from
98+
// the original script, spec §4.3).
9999
v := version
100100
if v == "" {
101101
v = "v4"

libs/localenv/target_test.go

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -90,6 +90,27 @@ func TestResolveJobClassicUsesSparkVersionReturn(t *testing.T) {
9090
assert.Equal(t, "dbr/15.4.x-scala2.12", ti.EnvKey)
9191
}
9292

93+
func TestResolveJobServerlessUsesRecordedVersion(t *testing.T) {
94+
// A serverless job (isServerless=true) pins its serverless version via the
95+
// third "recorded version" return; ResolveTarget must map it to the matching
96+
// serverless-vN rather than the classic dbr path.
97+
c := jobStubCompute{isServerless: true, version: "3"}
98+
ti, err := ResolveTarget(t.Context(), TargetFlags{Job: "42"}, c, BundleTarget{})
99+
require.NoError(t, err)
100+
assert.Equal(t, "job", ti.Source)
101+
assert.Empty(t, ti.SparkVersion)
102+
assert.Equal(t, "serverless/serverless-v3", ti.EnvKey)
103+
}
104+
105+
func TestResolveJobServerlessEmptyVersionFallsBackToV4(t *testing.T) {
106+
// When the job records no serverless version, ResolveTarget defaults to v4
107+
// (documented stand-in), matching the bundle serverless path.
108+
c := jobStubCompute{isServerless: true, version: ""}
109+
ti, err := ResolveTarget(t.Context(), TargetFlags{Job: "42"}, c, BundleTarget{})
110+
require.NoError(t, err)
111+
assert.Equal(t, "serverless/serverless-v4", ti.EnvKey)
112+
}
113+
93114
func TestValidateTargetFlagsMutuallyExclusive(t *testing.T) {
94115
assert.Error(t, ValidateTargetFlags(TargetFlags{Cluster: "a", Serverless: "v4"}))
95116
assert.NoError(t, ValidateTargetFlags(TargetFlags{Cluster: "a"}))

libs/localenv/uv.go

Lines changed: 92 additions & 47 deletions
Original file line numberDiff line numberDiff line change
@@ -12,27 +12,27 @@ import (
1212
"runtime"
1313
"strings"
1414

15+
"github.com/databricks/cli/libs/cmdio"
1516
"github.com/databricks/cli/libs/env"
1617
"github.com/databricks/cli/libs/log"
1718
"github.com/databricks/cli/libs/process"
1819
)
1920

21+
// EnvAutoInstallUv opts into installing uv without an interactive prompt. It
22+
// exists so non-interactive runs (CI, IDE integrations) can allow the install
23+
// that would otherwise be declined for lack of a TTY. Any truthy value enables it.
24+
const EnvAutoInstallUv = "DATABRICKS_LOCALENV_AUTO_INSTALL_UV"
25+
2026
// uvManager implements PackageManager using the uv tool.
2127
// https://docs.astral.sh/uv/
2228
type uvManager struct {
2329
bin string
2430
}
2531

26-
// newUvManager returns a uvManager whose binary path is resolved lazily via
27-
// EnsureAvailable.
28-
func newUvManager() *uvManager {
29-
return &uvManager{}
30-
}
31-
32-
// NewUvManager returns a PackageManager backed by the uv tool.
33-
// This is the exported constructor for use outside this package.
32+
// NewUvManager returns a PackageManager backed by the uv tool. The binary path
33+
// is resolved lazily via EnsureAvailable.
3434
func NewUvManager() PackageManager {
35-
return newUvManager()
35+
return &uvManager{}
3636
}
3737

3838
// Name returns "uv".
@@ -47,6 +47,10 @@ func (m *uvManager) Name() string {
4747
func (m *uvManager) EnsureAvailable(ctx context.Context) (string, error) {
4848
bin, err := discoverUv(ctx)
4949
if err != nil {
50+
if !confirmUvInstall(ctx) {
51+
return "", NewError(ErrUvMissing, nil,
52+
"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)
53+
}
5054
if installErr := installUv(ctx); installErr != nil {
5155
return "", NewError(ErrUvMissing, installErr, "uv installation failed")
5256
}
@@ -66,17 +70,24 @@ func (m *uvManager) EnsureAvailable(ctx context.Context) (string, error) {
6670
return strings.TrimSpace(version), nil
6771
}
6872

73+
// runUv runs the uv binary with args in dir, injecting UV_INDEX_URL from pip.conf
74+
// when appropriate. An empty dir runs in the current working directory
75+
// (process.WithDir("") is a no-op). The index-url is injected only when
76+
// resolveIndexURL returns non-empty; it returns "" when UV_INDEX_URL is already
77+
// set, so an explicit value in the environment is never clobbered.
78+
func (m *uvManager) runUv(ctx context.Context, args []string, dir string) error {
79+
if indexURL := m.resolveIndexURL(ctx); indexURL != "" {
80+
_, err := process.Background(ctx, args, process.WithDir(dir), process.WithEnv("UV_INDEX_URL", indexURL))
81+
return err
82+
}
83+
_, err := process.Background(ctx, args, process.WithDir(dir))
84+
return err
85+
}
86+
6987
// EnsurePython installs the requested Python minor version via uv.
7088
func (m *uvManager) EnsurePython(ctx context.Context, minor string) error {
7189
args := append([]string{m.bin}, m.pythonInstallArgs(minor)...)
72-
indexURL := m.resolveIndexURL(ctx)
73-
var err error
74-
if indexURL != "" {
75-
_, err = process.Background(ctx, args, process.WithEnv("UV_INDEX_URL", indexURL))
76-
} else {
77-
_, err = process.Background(ctx, args)
78-
}
79-
if err != nil {
90+
if err := m.runUv(ctx, args, ""); err != nil {
8091
return uvFailure(ErrPythonInstall, err, "uv python install "+minor)
8192
}
8293
return nil
@@ -85,14 +96,7 @@ func (m *uvManager) EnsurePython(ctx context.Context, minor string) error {
8596
// Provision runs `uv sync` inside projectDir to install project dependencies.
8697
func (m *uvManager) Provision(ctx context.Context, projectDir string) error {
8798
args := append([]string{m.bin}, m.syncArgs()...)
88-
indexURL := m.resolveIndexURL(ctx)
89-
var err error
90-
if indexURL != "" {
91-
_, err = process.Background(ctx, args, process.WithDir(projectDir), process.WithEnv("UV_INDEX_URL", indexURL))
92-
} else {
93-
_, err = process.Background(ctx, args, process.WithDir(projectDir))
94-
}
95-
if err != nil {
99+
if err := m.runUv(ctx, args, projectDir); err != nil {
96100
return uvFailure(ErrProvision, err, "uv sync")
97101
}
98102
return nil
@@ -117,14 +121,7 @@ func venvPython(projectDir string) string {
117121
// activated.
118122
func (m *uvManager) PostProvision(ctx context.Context, projectDir string) error {
119123
args := append([]string{m.bin}, m.pipSeedArgs(venvPython(projectDir))...)
120-
indexURL := m.resolveIndexURL(ctx)
121-
var err error
122-
if indexURL != "" {
123-
_, err = process.Background(ctx, args, process.WithDir(projectDir), process.WithEnv("UV_INDEX_URL", indexURL))
124-
} else {
125-
_, err = process.Background(ctx, args, process.WithDir(projectDir))
126-
}
127-
if err != nil {
124+
if err := m.runUv(ctx, args, projectDir); err != nil {
128125
return uvFailure(ErrProvision, err, "uv pip seed")
129126
}
130127
return nil
@@ -136,12 +133,16 @@ func (m *uvManager) PostProvision(ctx context.Context, projectDir string) error
136133
// error: PackageNotFoundError is caught so the probe never fails just because the
137134
// package is absent. The caller decides whether an empty version is acceptable.
138135
func (m *uvManager) Validate(ctx context.Context, projectDir string) (string, string, error) {
136+
// Each value is printed with a unique prefix so parsing greps for the prefix
137+
// rather than relying on line position: any stray line uv or the interpreter
138+
// writes to stdout (e.g. a warning) would otherwise shift a positional parse.
139+
// A missing databricks-connect prints an empty DBC: value, not an error.
139140
pyCode := `import sys, importlib.metadata
140-
print(f"{sys.version_info.major}.{sys.version_info.minor}")
141+
print(f"` + validatePyPrefix + `{sys.version_info.major}.{sys.version_info.minor}")
141142
try:
142-
print(importlib.metadata.version("databricks-connect"))
143+
print("` + validateDBCPrefix + `" + importlib.metadata.version("databricks-connect"))
143144
except importlib.metadata.PackageNotFoundError:
144-
print("")`
145+
print("` + validateDBCPrefix + `")`
145146
// --no-project runs the interpreter from the created .venv without re-resolving/syncing
146147
// the project's declared dependencies, so validation observes exactly what was installed.
147148
out, err := process.Background(ctx,
@@ -151,16 +152,32 @@ except importlib.metadata.PackageNotFoundError:
151152
if err != nil {
152153
return "", "", uvFailure(ErrValidate, err, "uv run python validation")
153154
}
154-
lines := strings.Split(strings.TrimSpace(out), "\n")
155-
if len(lines) < 1 || strings.TrimSpace(lines[0]) == "" {
155+
pyVer, ok := lineWithPrefix(out, validatePyPrefix)
156+
if !ok || pyVer == "" {
156157
return "", "", NewError(ErrValidate, nil, "unexpected output from uv run: %q", out)
157158
}
158-
// The databricks-connect line is empty when the package is not installed.
159-
dbcVer := ""
160-
if len(lines) >= 2 {
161-
dbcVer = strings.TrimSpace(lines[len(lines)-1])
159+
// The databricks-connect value is empty when the package is not installed.
160+
dbcVer, _ := lineWithPrefix(out, validateDBCPrefix)
161+
return pyVer, dbcVer, nil
162+
}
163+
164+
// Validation output prefixes: uv run's stdout is grepped for these rather than
165+
// parsed positionally, so extra lines from uv or the interpreter don't break it.
166+
const (
167+
validatePyPrefix = "PYVER:"
168+
validateDBCPrefix = "DBCVER:"
169+
)
170+
171+
// lineWithPrefix returns the trimmed remainder of the first line in out that
172+
// starts with prefix, and whether such a line was found.
173+
func lineWithPrefix(out, prefix string) (string, bool) {
174+
for line := range strings.SplitSeq(out, "\n") {
175+
line = strings.TrimSpace(line)
176+
if after, ok := strings.CutPrefix(line, prefix); ok {
177+
return strings.TrimSpace(after), true
178+
}
162179
}
163-
return strings.TrimSpace(lines[0]), dbcVer, nil
180+
return "", false
164181
}
165182

166183
// syncArgs returns the argument slice for `uv sync` (without the binary).
@@ -291,18 +308,46 @@ func uvFailure(code ErrorCode, err error, action string) *PipelineError {
291308
return NewError(code, err, "%s", msg)
292309
}
293310

311+
// confirmUvInstall reports whether the caller has consented to installUv running
312+
// a remote installer that mutates the machine. The EnvAutoInstallUv opt-in wins
313+
// outright (for CI / IDE integrations); otherwise an interactive session is
314+
// prompted, and a non-interactive session without the opt-in declines rather than
315+
// silently downloading and executing an installer.
316+
func confirmUvInstall(ctx context.Context) bool {
317+
if optIn, ok := env.GetBool(ctx, EnvAutoInstallUv); ok && optIn {
318+
return true
319+
}
320+
// EnsureAvailable is a library entry point reachable with a context that has
321+
// no cmdio (e.g. Pipeline built with context.Background()); IsPromptSupported
322+
// would panic there. Treat a missing cmdio as non-interactive and decline.
323+
if !cmdio.HasIO(ctx) || !cmdio.IsPromptSupported(ctx) {
324+
return false
325+
}
326+
// Name the OS-specific installer URL that installUv will actually fetch, so
327+
// the prompt is transparent about what runs (install.ps1 on Windows).
328+
ok, err := cmdio.AskYesOrNo(ctx, "uv is not installed. Download and run the official uv installer ("+uvInstallerURL()+")?")
329+
return err == nil && ok
330+
}
331+
332+
// uvInstallerURL returns the URL of the official uv installer script that
333+
// installUv fetches for the current OS.
334+
func uvInstallerURL() string {
335+
if runtime.GOOS == "windows" {
336+
return "https://astral.sh/uv/install.ps1"
337+
}
338+
return "https://astral.sh/uv/install.sh"
339+
}
340+
294341
// installUv runs the official uv installer for the current OS. Unix uses the
295342
// shell installer; Windows uses the PowerShell installer, because the Unix
296343
// `sh`/`curl` pipeline is not available in a default Windows shell.
297344
// https://docs.astral.sh/uv/getting-started/installation/
298345
func installUv(ctx context.Context) error {
299346
var cmd []string
300347
if runtime.GOOS == "windows" {
301-
// https://astral.sh/uv/install.ps1
302-
cmd = []string{"powershell", "-ExecutionPolicy", "ByPass", "-Command", "irm https://astral.sh/uv/install.ps1 | iex"}
348+
cmd = []string{"powershell", "-ExecutionPolicy", "ByPass", "-Command", "irm " + uvInstallerURL() + " | iex"}
303349
} else {
304-
// https://astral.sh/uv/install.sh
305-
cmd = []string{"sh", "-c", "curl -LsSf https://astral.sh/uv/install.sh | sh"}
350+
cmd = []string{"sh", "-c", "curl -LsSf " + uvInstallerURL() + " | sh"}
306351
}
307352
// This downloads and runs a remote installer that mutates the user's machine
308353
// (~/.local/bin), so record exactly what ran before it fires — visible under

0 commit comments

Comments
 (0)