diff --git a/cmd/localenv/compute.go b/cmd/localenv/compute.go index c573ca351b..8ce8c91314 100644 --- a/cmd/localenv/compute.go +++ b/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 +} diff --git a/cmd/localenv/compute_test.go b/cmd/localenv/compute_test.go new file mode 100644 index 0000000000..8620821a8f --- /dev/null +++ b/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)) + }) + } +} diff --git a/cmd/localenv/sync.go b/cmd/localenv/sync.go index a0b07f62e4..2814627801 100644 --- a/cmd/localenv/sync.go +++ b/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{ diff --git a/libs/localenv/target.go b/libs/localenv/target.go index b3ccb3171e..cba5666d3f 100644 --- a/libs/localenv/target.go +++ b/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" diff --git a/libs/localenv/target_test.go b/libs/localenv/target_test.go index 22beaee338..7e80dce241 100644 --- a/libs/localenv/target_test.go +++ b/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"})) diff --git a/libs/localenv/uv.go b/libs/localenv/uv.go index e9b4c911ce..f2996ba29b 100644 --- a/libs/localenv/uv.go +++ b/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,36 @@ 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. @@ -298,11 +345,9 @@ func uvFailure(code ErrorCode, err error, action string) *PipelineError { 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 diff --git a/libs/localenv/uv_test.go b/libs/localenv/uv_test.go index 106a97df93..8678c45481 100644 --- a/libs/localenv/uv_test.go +++ b/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) +}