diff --git a/cmd/cmd.go b/cmd/cmd.go index a2c3280b94..bcd81619fc 100644 --- a/cmd/cmd.go +++ b/cmd/cmd.go @@ -18,6 +18,7 @@ import ( "github.com/databricks/cli/cmd/experimental" "github.com/databricks/cli/cmd/fs" "github.com/databricks/cli/cmd/labs" + "github.com/databricks/cli/cmd/localenv" "github.com/databricks/cli/cmd/pipelines" "github.com/databricks/cli/cmd/quickstart" "github.com/databricks/cli/cmd/root" @@ -113,6 +114,7 @@ func New(ctx context.Context) *cobra.Command { cli.AddCommand(cache.New()) cli.AddCommand(experimental.New()) cli.AddCommand(psql.New()) + cli.AddCommand(localenv.New()) cli.AddCommand(configure.New()) cli.AddCommand(fs.New()) cli.AddCommand(labs.New(ctx)) diff --git a/cmd/localenv/compute.go b/cmd/localenv/compute.go new file mode 100644 index 0000000000..c573ca351b --- /dev/null +++ b/cmd/localenv/compute.go @@ -0,0 +1,80 @@ +package localenv + +import ( + "context" + "fmt" + "strconv" + + databricks "github.com/databricks/databricks-sdk-go" +) + +// sdkCompute adapts the Databricks SDK to the localenv.ComputeClient interface. +type sdkCompute struct { + w *databricks.WorkspaceClient +} + +// GetClusterSparkVersion returns the Spark version string for a running cluster. +func (c sdkCompute) GetClusterSparkVersion(ctx context.Context, clusterID string) (string, error) { + d, err := c.w.Clusters.GetByClusterId(ctx, clusterID) + if err != nil { + return "", fmt.Errorf("get cluster %s: %w", clusterID, err) + } + return d.SparkVersion, nil +} + +// GetJobSparkVersion inspects the job's configuration to determine compute type. +// +// A job is considered serverless when it has non-empty Environments (JobEnvironment +// entries), which signals the Databricks serverless runtime. A job with classic compute +// uses JobClusters; we read SparkVersion from the first job cluster's NewCluster spec. +// +// Task-level compute (tasks[].new_cluster / tasks[].existing_cluster_id with no +// job-level job_clusters) is not resolved here: it may vary per task and an +// existing_cluster_id would need a second lookup, which is out of scope for the +// initial job support. Such a job returns an actionable error rather than a wrong +// guess; use --cluster or --serverless explicitly instead. +func (c sdkCompute) GetJobSparkVersion(ctx context.Context, jobID string) (sparkVersion string, isServerless bool, version string, err error) { + id, err := strconv.ParseInt(jobID, 10, 64) + if err != nil { + return "", false, "", fmt.Errorf("invalid job ID %q: must be an integer: %w", jobID, err) + } + + job, err := c.w.Jobs.GetByJobId(ctx, id) + if err != nil { + return "", false, "", fmt.Errorf("get job %d: %w", id, err) + } + + if job.Settings == nil { + return "", false, "", fmt.Errorf("job %d has no settings", id) + } + + // A job that declares both serverless environments and classic job clusters is + // ambiguous: its tasks can run on different compute, so there is no single + // correct local environment to provision. Refuse rather than guess serverless. + if len(job.Settings.Environments) > 0 && len(job.Settings.JobClusters) > 0 { + return "", false, "", fmt.Errorf("job %d has both serverless environments and job clusters; pass --cluster or --serverless explicitly to disambiguate", id) + } + + // Serverless jobs have Environments populated; classic compute uses JobClusters. + if len(job.Settings.Environments) > 0 { + return "", true, "", nil + } + + if len(job.Settings.JobClusters) > 0 { + sv := job.Settings.JobClusters[0].NewCluster.SparkVersion + if sv == "" { + return "", false, "", fmt.Errorf("could not determine compute for job %d: first job cluster has no spark_version", id) + } + // Tasks can reference any job_cluster_key, so if the job's clusters do not + // all share one Spark version there is no single correct local environment. + // Refuse rather than silently provisioning for the first cluster. + for _, jc := range job.Settings.JobClusters[1:] { + if jc.NewCluster.SparkVersion != sv { + return "", false, "", fmt.Errorf("job %d has job clusters with differing spark_version; pass --cluster or --serverless explicitly to disambiguate", id) + } + } + return sv, false, sv, nil + } + + 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) +} diff --git a/cmd/localenv/localenv.go b/cmd/localenv/localenv.go new file mode 100644 index 0000000000..b93c36f710 --- /dev/null +++ b/cmd/localenv/localenv.go @@ -0,0 +1,45 @@ +package localenv + +import ( + "github.com/databricks/cli/cmd/root" + libslocalenv "github.com/databricks/cli/libs/localenv" + "github.com/spf13/cobra" +) + +// New returns the local-env command group. The group, subgroup, and verb names +// come from the single command-name constants in libs/localenv so a rename is a +// one-location change (spec ยง0 / invariant 8). +// +// The command is Hidden while the feature lands across the stacked PRs: it is +// wired and runnable for dogfooding, but stays out of help and completion until +// the final PR unveils it (removes this flag, adds the help line and changelog). +func New() *cobra.Command { + cmd := &cobra.Command{ + Use: libslocalenv.CommandGroup, + Short: "Manage local development environments matched to Databricks compute", + GroupID: "development", + Hidden: true, + Long: `Manage local development environments matched to a Databricks compute target. + +Derives the Python version, databricks-connect version, and dependency +constraints from the selected compute (cluster, serverless, or job) so that +local resolution matches the Databricks runtime.`, + RunE: root.ReportUnknownSubcommand, + } + cmd.AddCommand(newPythonCommand()) + return cmd +} + +// newPythonCommand returns the "python" subgroup. It is a parent-only node: with +// no verb it reports an unknown-subcommand error (mirroring the generated command +// groups) rather than doing nothing. +func newPythonCommand() *cobra.Command { + cmd := &cobra.Command{ + Use: libslocalenv.CommandSubgroup, + Short: "Manage the local Python environment", + Long: `Manage the local Python environment matched to a Databricks compute target.`, + RunE: root.ReportUnknownSubcommand, + } + cmd.AddCommand(newSyncCommand()) + return cmd +} diff --git a/cmd/localenv/output.go b/cmd/localenv/output.go new file mode 100644 index 0000000000..96c0322bad --- /dev/null +++ b/cmd/localenv/output.go @@ -0,0 +1,76 @@ +package localenv + +import ( + "context" + "fmt" + + "github.com/databricks/cli/cmd/root" + "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/cli/libs/flags" + libslocalenv "github.com/databricks/cli/libs/localenv" + "github.com/spf13/cobra" +) + +// renderResult renders the pipeline result to the command's output. +// In JSON mode it renders the full structured result (even on error). +// In text mode it prints phase headers and a summary, then returns the error. +// +// res is always non-nil: Pipeline.Run constructs and returns a fully-populated +// Result (with the canonical phase list and error object) on every path, +// including failures, so no nil guard is needed here. +func renderResult(ctx context.Context, cmd *cobra.Command, res *libslocalenv.Result, pipelineErr error) error { + if root.OutputType(cmd) == flags.OutputJSON { + if err := cmdio.Render(ctx, res); err != nil { + return err + } + // The JSON object is the only thing written to stdout. On failure we still + // need a non-zero exit, but returning pipelineErr would make the root print + // "Error: ..." to stderr. ErrAlreadyPrinted exits non-zero without that. + if pipelineErr != nil { + return root.ErrAlreadyPrinted + } + return nil + } + + // Text mode: print each phase in execution order. + for _, phase := range res.Phases { + if phase.Detail != "" { + cmdio.LogString(ctx, fmt.Sprintf("%-10s %s %s", phase.Phase, phase.Status, phase.Detail)) + } else { + cmdio.LogString(ctx, fmt.Sprintf("%-10s %s", phase.Phase, phase.Status)) + } + } + + for _, w := range res.Warnings { + cmdio.LogString(ctx, "warning: "+w.Message) + } + + if pipelineErr != nil { + cmdio.LogString(ctx, "For more detail, re-run with --debug, or --output json to share a structured report.") + return pipelineErr + } + + // Print a final success / check summary. + if res.DryRun { + if res.Plan != nil { + cmdio.LogString(ctx, "Plan: "+res.Plan.WouldWrite) + for _, region := range res.Plan.ChangedRegions { + cmdio.LogString(ctx, " changed region: "+region) + } + } + cmdio.LogString(ctx, "Check complete. No files were modified.") + return nil + } + + if res.Resolved != nil { + summary := "Success: python=" + res.Resolved.PythonVersion + if res.Resolved.DBConnectVersion != "" { + summary += " databricks-connect=" + res.Resolved.DBConnectVersion + } + if res.VenvPath != "" { + summary += " venv=" + res.VenvPath + } + cmdio.LogString(ctx, summary) + } + return nil +} diff --git a/cmd/localenv/sync.go b/cmd/localenv/sync.go new file mode 100644 index 0000000000..a0b07f62e4 --- /dev/null +++ b/cmd/localenv/sync.go @@ -0,0 +1,154 @@ +package localenv + +import ( + "context" + "os" + "path/filepath" + + "github.com/databricks/cli/cmd/root" + "github.com/databricks/cli/libs/cmdctx" + "github.com/databricks/cli/libs/env" + libslocalenv "github.com/databricks/cli/libs/localenv" + "github.com/spf13/cobra" +) + +// envConstraintSource is the environment variable that overrides the constraint +// source with a full base URL (used e.g. by tests pointing at a local server). +// When unset, the base URL is derived from the hosting repo via +// libslocalenv.RepoConstraintBaseURL (which reads its own repo env var). +const envConstraintSource = "DATABRICKS_LOCALENV_CONSTRAINT_SOURCE" + +func newSyncCommand() *cobra.Command { + cmd := &cobra.Command{ + Use: libslocalenv.CommandVerb, + Short: "Provision a local Python environment matched to a Databricks compute target", + Long: `Provision (or update) a local Python environment matched to a Databricks compute target. + +Resolves the target to an environment key, fetches the pinned Python version, +databricks-connect version, and dependency constraints published for that key, +then provisions a matched .venv with uv. A project with no pyproject.toml is +initialized from scratch; an existing pyproject.toml is merged in place (its +env-owned sections are refreshed, user-owned content is preserved).`, + } + // The target is selected via flags; reject stray positional args rather than + // silently ignoring them. + cmd.Args = cobra.NoArgs + cmd.PreRunE = root.MustWorkspaceClient + addTargetFlags(cmd) + cmd.RunE = func(cmd *cobra.Command, args []string) error { + return runPipeline(cmd) + } + return cmd +} + +// addTargetFlags adds the shared target and mode flags to a command. +func addTargetFlags(cmd *cobra.Command) { + cmd.Flags().String("cluster", "", "cluster ID to use as the compute target") + cmd.Flags().String("serverless", "", "serverless version to use as the compute target (e.g. v4)") + cmd.Flags().String("job", "", "job ID to use as the compute target") + cmd.Flags().Bool("constraints-only", false, "apply the Python version and constraints without adding the databricks-connect dependency") + cmd.Flags().Bool("check", false, "compute the plan without writing files or provisioning") + cmd.Flags().String("constraint-source", "", "URL for the constraint source (overrides "+envConstraintSource+")") + // Hide constraint-source from casual --help output; it is a power-user escape hatch. + _ = cmd.Flags().MarkHidden("constraint-source") + cmd.MarkFlagsMutuallyExclusive("cluster", "serverless", "job") +} + +// runPipeline builds and runs the local-env Pipeline. +func runPipeline(cmd *cobra.Command) error { + ctx := cmd.Context() + + cluster, _ := cmd.Flags().GetString("cluster") + serverless, _ := cmd.Flags().GetString("serverless") + job, _ := cmd.Flags().GetString("job") + constraintsOnly, _ := cmd.Flags().GetBool("constraints-only") + check, _ := cmd.Flags().GetBool("check") + constraintSource, _ := cmd.Flags().GetString("constraint-source") + + targetFlags := libslocalenv.TargetFlags{ + Cluster: cluster, + Serverless: serverless, + Job: job, + } + // ValidateTargetFlags is kept despite MarkFlagsMutuallyExclusive above: + // it also validates the library path (no Cobra equivalent) and guards + // non-Cobra call paths such as tests that invoke runPipeline directly. + if err := libslocalenv.ValidateTargetFlags(targetFlags); err != nil { + return err + } + + mode := libslocalenv.ModeDefault + if constraintsOnly { + mode = libslocalenv.ModeConstraintsOnly + } + + constraintBaseURL := resolveConstraintBaseURL(ctx, constraintSource) + + projectDir, err := os.Getwd() + if err != nil { + return err + } + + cacheDir, err := os.UserCacheDir() + if err != nil { + return err + } + cacheDir = filepath.Join(cacheDir, "databricks", "localenv") + + bt := bundleTarget(cmd) + + w := cmdctx.WorkspaceClient(ctx) + p := &libslocalenv.Pipeline{ + Mode: mode, + Check: check, + ProjectDir: projectDir, + ConstraintBaseURL: constraintBaseURL, + CacheDir: cacheDir, + Flags: targetFlags, + Compute: sdkCompute{w: w}, + Bundle: bt, + PM: libslocalenv.NewUvManager(), + } + + res, pipelineErr := p.Run(ctx) + return renderResult(ctx, cmd, res, pipelineErr) +} + +// resolveConstraintBaseURL returns the constraint base URL using ordered precedence: +// an explicit --constraint-source flag, then a full-URL override from +// DATABRICKS_LOCALENV_CONSTRAINT_SOURCE, then the URL derived from the hosting repo +// (libslocalenv.RepoConstraintBaseURL). All three may be unset, in which case it +// returns "" and the pipeline reports the missing source at the fetch phase. +func resolveConstraintBaseURL(ctx context.Context, flagValue string) string { + if flagValue != "" { + return flagValue + } + if v, ok := env.Lookup(ctx, envConstraintSource); ok && v != "" { + return v + } + return libslocalenv.RepoConstraintBaseURL(ctx) +} + +// bundleTarget reads the active bundle (if any) and maps its compute configuration +// to a libslocalenv.BundleTarget. +// +// Only the top-level bundle.cluster_id field is consulted here; serverless is not +// recorded in the bundle config, so Selected=true is set only when a cluster ID is +// present. If the bundle is absent or has no cluster_id, Selected=false is returned +// so the pipeline falls through to requiring an explicit flag. +// +// TODO: extend once bundle config exposes a serverless field at the bundle level. +func bundleTarget(cmd *cobra.Command) libslocalenv.BundleTarget { + b := root.TryConfigureBundle(cmd) + if b == nil { + return libslocalenv.BundleTarget{Selected: false} + } + clusterID := b.Config.Bundle.ClusterId + if clusterID == "" { + return libslocalenv.BundleTarget{Selected: false} + } + return libslocalenv.BundleTarget{ + ClusterID: clusterID, + Selected: true, + } +} diff --git a/libs/localenv/uv.go b/libs/localenv/uv.go new file mode 100644 index 0000000000..5440d0cf66 --- /dev/null +++ b/libs/localenv/uv.go @@ -0,0 +1,356 @@ +package localenv + +import ( + "bufio" + "context" + "errors" + "net/url" + "os" + "os/exec" + "path/filepath" + "regexp" + "runtime" + "strings" + + "github.com/databricks/cli/libs/env" + "github.com/databricks/cli/libs/log" + "github.com/databricks/cli/libs/process" +) + +// 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. +func NewUvManager() PackageManager { + return newUvManager() +} + +// Name returns "uv". +func (m *uvManager) Name() string { + return "uv" +} + +// EnsureAvailable discovers or installs uv and records the binary path. +// It runs the official installer when uv is not found on the PATH or in the +// standard candidate locations. +// https://docs.astral.sh/uv/getting-started/installation/ +func (m *uvManager) EnsureAvailable(ctx context.Context) (string, error) { + bin, err := discoverUv(ctx) + if err != nil { + if installErr := installUv(ctx); installErr != nil { + return "", NewError(ErrUvMissing, installErr, "uv installation failed") + } + bin, err = discoverUv(ctx) + if err != nil { + return "", err + } + } + log.Debugf(ctx, "uv: discovered binary at %s", bin) + m.bin = bin + + // Use --version (not "version") to avoid project-scoped sub-command that requires pyproject.toml. + version, err := process.Background(ctx, []string{m.bin, "--version"}) + if err != nil { + return "", uvFailure(ErrUvMissing, err, "uv version check") + } + return strings.TrimSpace(version), nil +} + +// 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 { + return uvFailure(ErrPythonInstall, err, "uv python install "+minor) + } + return nil +} + +// 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 { + return uvFailure(ErrProvision, err, "uv sync") + } + return nil +} + +// venvPython returns the path to the virtualenv's Python interpreter, +// accounting for the Windows (Scripts/python.exe) vs Unix (bin/python) layout. +func venvPython(projectDir string) string { + if runtime.GOOS == "windows" { + return filepath.Join(projectDir, venvDir, "Scripts", "python.exe") + } + return filepath.Join(projectDir, venvDir, "bin", "python") +} + +// PostProvision seeds pip into the project's virtual environment. +// +// VS Code's ms-python.vscode-python-envs extension falls back to +// `python -m pip list` when its `uv --version` probe fails on the GUI PATH. +// uv virtual environments do not include pip by default, and `uv sync` strips +// pip if it was previously present. Seeding pip after every sync ensures the +// VS Code integration works correctly regardless of how the environment was +// 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 { + return uvFailure(ErrProvision, err, "uv pip seed") + } + return nil +} + +// Validate reads the Python minor version and databricks-connect package +// version from the project's virtual environment. When databricks-connect is not +// installed (constraints-only mode), the second line is empty rather than an +// 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) { + pyCode := `import sys, importlib.metadata +print(f"{sys.version_info.major}.{sys.version_info.minor}") +try: + print(importlib.metadata.version("databricks-connect")) +except importlib.metadata.PackageNotFoundError: + print("")` + // --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, + []string{m.bin, "run", "--no-project", "python", "-c", pyCode}, + process.WithDir(projectDir), + ) + 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]) == "" { + 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]) + } + return strings.TrimSpace(lines[0]), dbcVer, nil +} + +// syncArgs returns the argument slice for `uv sync` (without the binary). +func (m *uvManager) syncArgs() []string { + return []string{"sync"} +} + +// pythonInstallArgs returns the argument slice for `uv python install `. +func (m *uvManager) pythonInstallArgs(minor string) []string { + return []string{"python", "install", minor} +} + +// pipSeedArgs returns the argument slice for seeding pip into the venv. +func (m *uvManager) pipSeedArgs(venvPython string) []string { + return []string{"pip", "install", "pip", "--python", venvPython} +} + +// pipIndexURLRe matches `index-url = ` lines in pip.conf. +var pipIndexURLRe = regexp.MustCompile(`(?i)^\s*index-url\s*=\s*(\S+)`) + +// pipConfIndexURL reads the user's pip config and returns the index-url value. +// uv ignores pip.conf; on Databricks-managed machines pypi.org is blocked and +// the corporate PyPI proxy is declared via pip.conf. Bridging the value through +// UV_INDEX_URL lets uv reach the proxy. +// +// pip stores its per-user config in OS-specific locations, so we probe them in +// order and return the first index-url found. +// https://pip.pypa.io/en/stable/topics/configuration/#location +func pipConfIndexURL(ctx context.Context) string { + for _, confPath := range pipConfPaths(ctx) { + if url := indexURLFromFile(confPath); url != "" { + return url + } + } + return "" +} + +// pipConfPaths returns the per-user pip config file locations to probe, in +// precedence order, for the current OS. An empty home yields no paths. +// https://pip.pypa.io/en/stable/topics/configuration/#location +func pipConfPaths(ctx context.Context) []string { + home, err := env.UserHomeDir(ctx) + if err != nil || home == "" { + return nil + } + switch runtime.GOOS { + case "windows": + // pip reads %APPDATA%\pip\pip.ini; fall back to the home-relative path + // when APPDATA is unset. + if appData, ok := env.Lookup(ctx, "APPDATA"); ok && appData != "" { + return []string{filepath.Join(appData, "pip", "pip.ini")} + } + return []string{filepath.Join(home, "pip", "pip.ini")} + case "darwin": + return []string{ + filepath.Join(home, "Library", "Application Support", "pip", "pip.conf"), + filepath.Join(home, ".config", "pip", "pip.conf"), + // Legacy per-user location pip still reads. + filepath.Join(home, ".pip", "pip.conf"), + } + default: + return []string{ + filepath.Join(home, ".config", "pip", "pip.conf"), + // Legacy per-user location pip still reads. + filepath.Join(home, ".pip", "pip.conf"), + } + } +} + +// indexURLFromFile returns the first index-url value in a pip config file, or "" +// if the file is absent or has no index-url entry. +func indexURLFromFile(confPath string) string { + f, err := os.Open(confPath) + if err != nil { + return "" + } + defer f.Close() + + scanner := bufio.NewScanner(f) + for scanner.Scan() { + if m := pipIndexURLRe.FindStringSubmatch(scanner.Text()); m != nil { + return strings.TrimSpace(m[1]) + } + } + return "" +} + +// resolveIndexURL returns a UV_INDEX_URL value to inject, or "" when none is +// needed. It returns "" when UV_INDEX_URL is already set in the context env +// (so the caller's explicit value is never overridden) and also when pip.conf +// has no index-url entry. +func (m *uvManager) resolveIndexURL(ctx context.Context) string { + if _, ok := env.Lookup(ctx, "UV_INDEX_URL"); ok { + log.Debugf(ctx, "uv: UV_INDEX_URL already set in environment, not overriding") + return "" + } + url := pipConfIndexURL(ctx) + if url != "" { + // Redact any embedded credentials: private PyPI proxies often carry + // userinfo (https://user:pass@host/simple) that must not reach debug logs. + log.Debugf(ctx, "uv: using package index %s from pip.conf", redactURLCredentials(url)) + } else { + log.Debugf(ctx, "uv: no UV_INDEX_URL and no index-url in pip.conf; uv will use its default index (pypi.org)") + } + return url +} + +// redactURLCredentials strips userinfo from a URL so it is safe to log. A value +// that does not parse as a URL is returned unchanged (it carries no parseable +// userinfo to leak). +func redactURLCredentials(raw string) string { + u, err := url.Parse(raw) + if err != nil || u.User == nil { + return raw + } + u.User = nil + return u.String() +} + +// uvFailure builds a PipelineError from a failed uv invocation, appending uv's +// stderr to the message so callers can see the actual failure reason (e.g. +// "Connection refused") rather than just the exit code. +func uvFailure(code ErrorCode, err error, action string) *PipelineError { + msg := action + " failed" + if perr, ok := errors.AsType[*process.ProcessError](err); ok && strings.TrimSpace(perr.Stderr) != "" { + msg = msg + ": " + strings.TrimSpace(perr.Stderr) + } + return NewError(code, err, "%s", msg) +} + +// 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"} + } else { + // https://astral.sh/uv/install.sh + cmd = []string{"sh", "-c", "curl -LsSf https://astral.sh/uv/install.sh | sh"} + } + _, err := process.Background(ctx, cmd) + return err +} + +// discoverUv searches for the uv binary on PATH and in well-known install +// locations. It returns NewError(ErrUvMissing, ...) if uv is not found. +// +// Candidate locations follow the uv installer defaults: +// https://docs.astral.sh/uv/getting-started/installation/ +// XDG_BIN_HOME is specified by the XDG Base Directory Specification: +// https://specifications.freedesktop.org/basedir-spec/latest/ +func discoverUv(ctx context.Context) (string, error) { + // Prefer PATH lookup first; it respects user customisation. exec.LookPath + // applies PATHEXT on Windows, so "uv" resolves to uv.exe there. + if p, err := exec.LookPath("uv"); err == nil { + return p, nil + } + + home, _ := env.UserHomeDir(ctx) + + // XDG_BIN_HOME defaults to $HOME/.local/bin when unset. + xdgBinHome, _ := env.Lookup(ctx, "XDG_BIN_HOME") + + // The installer writes uv.exe on Windows; os.Stat needs the exact name. + exe := "uv" + if runtime.GOOS == "windows" { + exe = "uv.exe" + } + candidates := []string{ + filepath.Join(home, ".local", "bin", exe), + filepath.Join(xdgBinHome, exe), + "/opt/homebrew/bin/uv", + "/usr/local/bin/uv", + } + + for _, c := range candidates { + // Skip relative paths produced when home or xdgBinHome is empty (e.g. + // filepath.Join("", "uv") == "uv"): os.Stat would resolve them against the + // CWD and could pick up a stray ./uv that is not the real binary. + if !filepath.IsAbs(c) { + continue + } + if _, err := os.Stat(c); err == nil { + return c, nil + } + } + + return "", NewError(ErrUvMissing, nil, + "uv not found on PATH or in well-known locations (%s)", strings.Join(candidates, ", ")) +} diff --git a/libs/localenv/uv_test.go b/libs/localenv/uv_test.go new file mode 100644 index 0000000000..106a97df93 --- /dev/null +++ b/libs/localenv/uv_test.go @@ -0,0 +1,223 @@ +package localenv + +import ( + "context" + "errors" + "os" + "path/filepath" + "runtime" + "testing" + + "github.com/databricks/cli/libs/env" + "github.com/databricks/cli/libs/process" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// writePipConf writes conf to the primary OS-specific pip config path under a +// fresh temp home and returns a context rooted at that home. On Windows the +// primary path is %APPDATA%\pip\pip.ini, so APPDATA is pointed inside the temp +// home too; without this the file would land in a location pipConfPaths never +// probes on that OS and the test would read an empty index-url. +func writePipConf(t *testing.T, conf string) context.Context { + t.Helper() + tmp := t.TempDir() + ctx := env.WithUserHomeDir(t.Context(), tmp) + if runtime.GOOS == "windows" { + ctx = env.Set(ctx, "APPDATA", filepath.Join(tmp, "AppData", "Roaming")) + } + paths := pipConfPaths(ctx) + require.NotEmpty(t, paths) + primary := paths[0] + require.NoError(t, os.MkdirAll(filepath.Dir(primary), 0o755)) + require.NoError(t, os.WriteFile(primary, []byte(conf), 0o644)) + return ctx +} + +func TestUvArgs(t *testing.T) { + m := &uvManager{bin: "uv"} + assert.Equal(t, []string{"sync"}, m.syncArgs()) + assert.Equal(t, []string{"python", "install", "3.12"}, m.pythonInstallArgs("3.12")) + assert.Equal(t, []string{"pip", "install", "pip", "--python", "/p/.venv/bin/python"}, m.pipSeedArgs("/p/.venv/bin/python")) +} + +func TestDiscoverUvFindsBinOnPath(t *testing.T) { + dir := t.TempDir() + // exec.LookPath only resolves a bare "uv" to a file with a PATHEXT extension + // on Windows, so the on-PATH binary must be named uv.exe there. + exe := "uv" + if runtime.GOOS == "windows" { + exe = "uv.exe" + } + bin := filepath.Join(dir, exe) + require.NoError(t, os.WriteFile(bin, []byte("#!/bin/sh\n"), 0o755)) + t.Setenv("PATH", dir) + got, err := discoverUv(t.Context()) + require.NoError(t, err) + assert.Equal(t, bin, got) +} + +func TestDiscoverUvFindsBinInLocalBin(t *testing.T) { + // The installer drops uv (uv.exe on Windows) under ~/.local/bin; discoverUv + // must find it there when it is not on PATH, using the OS-appropriate name. + home := t.TempDir() + binDir := filepath.Join(home, ".local", "bin") + require.NoError(t, os.MkdirAll(binDir, 0o755)) + exe := "uv" + if runtime.GOOS == "windows" { + exe = "uv.exe" + } + bin := filepath.Join(binDir, exe) + require.NoError(t, os.WriteFile(bin, []byte("#!/bin/sh\n"), 0o755)) + + t.Setenv("PATH", t.TempDir()) // no uv on PATH + ctx := env.WithUserHomeDir(t.Context(), home) + ctx = env.Set(ctx, "XDG_BIN_HOME", "") + got, err := discoverUv(ctx) + require.NoError(t, err) + assert.Equal(t, bin, got) +} + +func TestDiscoverUvSkipsRelativeCandidatesWhenHomeUnset(t *testing.T) { + // Regression: when HOME and XDG_BIN_HOME are unset the candidate paths + // collapse to relative "uv" / ".local/bin/uv". discoverUv must not os.Stat + // them against the CWD and return a stray ./uv as if it were the real binary. + dir := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(dir, "uv"), []byte("#!/bin/sh\n"), 0o755)) + t.Chdir(dir) + t.Setenv("PATH", t.TempDir()) // a PATH with no uv, so LookPath falls through + + ctx := env.WithUserHomeDir(t.Context(), "") + ctx = env.Set(ctx, "XDG_BIN_HOME", "") + got, err := discoverUv(ctx) + if err == nil { + assert.True(t, filepath.IsAbs(got), "discoverUv must not return a relative path; got %q", got) + } +} + +func TestPipConfIndexURLReadsOSSpecificPath(t *testing.T) { + // The primary pip config path for the current OS must be honored, not just + // the Linux/XDG ~/.config/pip/pip.conf location. + tmp := t.TempDir() + ctx := env.WithUserHomeDir(t.Context(), tmp) + if runtime.GOOS == "windows" { + ctx = env.Set(ctx, "APPDATA", filepath.Join(tmp, "AppData", "Roaming")) + } + + paths := pipConfPaths(ctx) + require.NotEmpty(t, paths) + primary := paths[0] + require.NoError(t, os.MkdirAll(filepath.Dir(primary), 0o755)) + require.NoError(t, os.WriteFile(primary, []byte("[global]\nindex-url = https://proxy.example/simple\n"), 0o644)) + + assert.Equal(t, "https://proxy.example/simple", pipConfIndexURL(ctx)) +} + +func TestRedactURLCredentials(t *testing.T) { + cases := []struct { + in string + want string + }{ + {"https://user:pass@proxy.example/simple", "https://proxy.example/simple"}, + {"https://token@proxy.example/simple", "https://proxy.example/simple"}, + {"https://proxy.example/simple", "https://proxy.example/simple"}, + {"not a url", "not a url"}, + } + for _, tc := range cases { + assert.Equal(t, tc.want, redactURLCredentials(tc.in)) + } +} + +func TestPipConfIndexURLReadsLegacyPath(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("legacy ~/.pip path is Unix/macOS only") + } + // The legacy ~/.pip/pip.conf location pip still reads must be honored when the + // newer per-user locations are absent. + tmp := t.TempDir() + legacy := filepath.Join(tmp, ".pip") + require.NoError(t, os.MkdirAll(legacy, 0o755)) + require.NoError(t, os.WriteFile(filepath.Join(legacy, "pip.conf"), []byte("[global]\nindex-url = https://legacy.example/simple\n"), 0o644)) + + ctx := env.WithUserHomeDir(t.Context(), tmp) + assert.Equal(t, "https://legacy.example/simple", pipConfIndexURL(ctx)) +} + +func TestPipConfIndexURL(t *testing.T) { + t.Run("returns_url_from_pip_conf", func(t *testing.T) { + ctx := writePipConf(t, "[global]\nindex-url = https://proxy.example/simple\n") + got := pipConfIndexURL(ctx) + assert.Equal(t, "https://proxy.example/simple", got) + }) + + t.Run("returns_empty_when_no_pip_conf", func(t *testing.T) { + tmp := t.TempDir() + ctx := env.WithUserHomeDir(t.Context(), tmp) + got := pipConfIndexURL(ctx) + assert.Empty(t, got) + }) + + t.Run("returns_empty_when_no_index_url_in_conf", func(t *testing.T) { + ctx := writePipConf(t, "[global]\nextra-index-url = https://other.example/simple\n") + got := pipConfIndexURL(ctx) + assert.Empty(t, got) + }) +} + +func TestResolveIndexURLRespectsExistingEnv(t *testing.T) { + m := &uvManager{} + + t.Run("returns_empty_when_UV_INDEX_URL_already_set", func(t *testing.T) { + // Set up a pip.conf that would otherwise be used. + ctx := writePipConf(t, "[global]\nindex-url = https://proxy.example/simple\n") + // When UV_INDEX_URL is in ctx, resolveIndexURL must not override it. + ctx = env.Set(ctx, "UV_INDEX_URL", "https://explicit.example/simple") + + got := m.resolveIndexURL(ctx) + assert.Empty(t, got) + }) + + t.Run("returns_pip_conf_url_when_UV_INDEX_URL_unset", func(t *testing.T) { + // CI sets UV_INDEX_URL to a package mirror in the process environment, and + // env.Lookup falls back to os.LookupEnv, so it must actually be removed + // (setting it to "" still reports ok=true) for resolveIndexURL to fall + // through to pip.conf. t.Setenv registers the cleanup that restores it. + t.Setenv("UV_INDEX_URL", "") + os.Unsetenv("UV_INDEX_URL") + + ctx := writePipConf(t, "[global]\nindex-url = https://proxy.example/simple\n") + got := m.resolveIndexURL(ctx) + assert.Equal(t, "https://proxy.example/simple", got) + }) +} + +func TestUvFailureIncludesStderr(t *testing.T) { + t.Run("includes_stderr_when_present", func(t *testing.T) { + underlying := &process.ProcessError{ + Command: "uv sync", + Err: errors.New("exit status 2"), + Stderr: "error: Connection refused\n", + } + pe := uvFailure(ErrProvision, underlying, "uv sync") + assert.Equal(t, ErrProvision, pe.Code) + assert.Contains(t, pe.Msg, "Connection refused") + assert.NotEqual(t, '\n', pe.Msg[len(pe.Msg)-1], "Msg must not end with a newline") + }) + + t.Run("omits_stderr_suffix_when_empty", func(t *testing.T) { + underlying := &process.ProcessError{ + Command: "uv sync", + Err: errors.New("exit status 2"), + Stderr: "", + } + pe := uvFailure(ErrProvision, underlying, "uv sync") + assert.Equal(t, ErrProvision, pe.Code) + assert.Equal(t, "uv sync failed", pe.Msg) + }) + + t.Run("non_process_error_uses_action_only", func(t *testing.T) { + pe := uvFailure(ErrProvision, errors.New("some other error"), "uv sync") + assert.Equal(t, ErrProvision, pe.Code) + assert.Equal(t, "uv sync failed", pe.Msg) + }) +}