Skip to content

Commit 52c4795

Browse files
committed
local-env: address #5832 follow-up review comments
Follow-ups deferred from the #5832 review, now that the command has landed (still hidden until #5835): - uv install consent (libs/localenv/uv.go): EnsureAvailable no longer runs the remote uv installer silently. It now requires consent — a truthy DATABRICKS_LOCALENV_AUTO_INSTALL_UV opt-in for non-interactive runs, or an interactive y/N prompt (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 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 falls back to v4. - double bundle load (cmd/localenv/sync.go): skip bundleTarget entirely 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 does not 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 is not clobbered). Adds unit tests for the install consent gate and the sentinel parse. Co-authored-by: Isaac
1 parent 3b9fe15 commit 52c4795

5 files changed

Lines changed: 140 additions & 48 deletions

File tree

cmd/localenv/compute.go

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -57,7 +57,15 @@ func (c sdkCompute) GetJobSparkVersion(ctx context.Context, jobID string) (spark
5757

5858
// Serverless jobs have Environments populated; classic compute uses JobClusters.
5959
if len(job.Settings.Environments) > 0 {
60-
return "", true, "", nil
60+
// The serverless environment version (e.g. "4") is recorded on the job's
61+
// environment spec, unlike the bundle path where it is unavailable. Return
62+
// it so ResolveTarget pins the matching serverless-vN instead of defaulting
63+
// to v4. Spec/version may be empty on older jobs; that falls back to v4.
64+
version := ""
65+
if spec := job.Settings.Environments[0].Spec; spec != nil {
66+
version = spec.EnvironmentVersion
67+
}
68+
return "", true, version, nil
6169
}
6270

6371
if len(job.Settings.JobClusters) > 0 {

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

Lines changed: 76 additions & 43 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.Split(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,6 +308,22 @@ 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+
if !cmdio.IsPromptSupported(ctx) {
321+
return false
322+
}
323+
ok, err := cmdio.AskYesOrNo(ctx, "uv is not installed. Download and run the official uv installer (https://astral.sh/uv/install.sh)?")
324+
return err == nil && ok
325+
}
326+
294327
// installUv runs the official uv installer for the current OS. Unix uses the
295328
// shell installer; Windows uses the PowerShell installer, because the Unix
296329
// `sh`/`curl` pipeline is not available in a default Windows shell.

libs/localenv/uv_test.go

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import (
88
"runtime"
99
"testing"
1010

11+
"github.com/databricks/cli/libs/cmdio"
1112
"github.com/databricks/cli/libs/env"
1213
"github.com/databricks/cli/libs/process"
1314
"github.com/stretchr/testify/assert"
@@ -221,3 +222,46 @@ func TestUvFailureIncludesStderr(t *testing.T) {
221222
assert.Equal(t, "uv sync failed", pe.Msg)
222223
})
223224
}
225+
226+
func TestConfirmUvInstall(t *testing.T) {
227+
t.Run("opt_in_env_var_consents_without_prompt", func(t *testing.T) {
228+
// Non-interactive context, but the opt-in env var grants consent.
229+
ctx := env.Set(t.Context(), EnvAutoInstallUv, "1")
230+
assert.True(t, confirmUvInstall(ctx))
231+
})
232+
233+
t.Run("non_interactive_without_opt_in_declines", func(t *testing.T) {
234+
// SetupTest defaults to PromptSupported=false, i.e. non-interactive.
235+
ctx, _ := cmdio.SetupTest(t.Context(), cmdio.TestOptions{})
236+
assert.False(t, confirmUvInstall(ctx))
237+
})
238+
239+
t.Run("falsey_opt_in_does_not_consent_when_non_interactive", func(t *testing.T) {
240+
ctx, _ := cmdio.SetupTest(t.Context(), cmdio.TestOptions{})
241+
ctx = env.Set(ctx, EnvAutoInstallUv, "0")
242+
assert.False(t, confirmUvInstall(ctx))
243+
})
244+
}
245+
246+
func TestLineWithPrefix(t *testing.T) {
247+
// A stray leading line (as uv might emit) must not shift the parse: the
248+
// value is located by prefix, not position.
249+
out := "warning: something\nPYVER:3.12\nDBCVER:17.2.0\n"
250+
251+
pyVer, ok := lineWithPrefix(out, validatePyPrefix)
252+
assert.True(t, ok)
253+
assert.Equal(t, "3.12", pyVer)
254+
255+
dbcVer, ok := lineWithPrefix(out, validateDBCPrefix)
256+
assert.True(t, ok)
257+
assert.Equal(t, "17.2.0", dbcVer)
258+
259+
// An empty DBC value (databricks-connect absent) is found but blank.
260+
empty, ok := lineWithPrefix("PYVER:3.12\nDBCVER:\n", validateDBCPrefix)
261+
assert.True(t, ok)
262+
assert.Empty(t, empty)
263+
264+
// A missing prefix reports not-found rather than a wrong line.
265+
_, ok = lineWithPrefix("PYVER:3.12\n", validateDBCPrefix)
266+
assert.False(t, ok)
267+
}

0 commit comments

Comments
 (0)