[VPEX][6/8] Add local-env uv backend and hidden CLI command

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))

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)
+}

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
+}

cmd/localenv/output.go

@@ -0,0 +1,81 @@
+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.")
+		// The failing phase's message was already printed by the phase loop above
+		// (Pipeline.fail sets the errored phase's Detail to the error text).
+		// Returning pipelineErr would make root print "Error: ..." with the same
+		// message again, since PipelineError.Unwrap yields the cause, not the
+		// ErrAlreadyPrinted sentinel. Signal already-printed to exit non-zero once.
+		return root.ErrAlreadyPrinted
+	}
+
+	// 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
+}

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,
+	}
+}