Skip to content

Commit 2fd09fe

Browse files
committed
Add local-env uv backend and hidden CLI command
Sixth in the stacked series. Wires the engine to a runnable CLI command, registered Hidden so it does not appear in help or completion until the final unveil PR (it is invocable for dogfooding in the meantime). - libs/localenv/uv.go: the uv implementation of the PackageManager interface (discover/install uv, install Python, uv sync, seed pip, validate the venv), plus the pip.conf -> UV_INDEX_URL bridge for Databricks-managed machines. - cmd/localenv: the command tree matching the target path "local-env python sync" — a top group (local-env), a python subgroup, and the sync verb. Parent nodes use root.ReportUnknownSubcommand so an unknown subcommand errors while a bare group shows help. sync resolves flags/bundle target, builds the Pipeline with the uv manager, and renders text or --json. - cmd/cmd.go: register the group (Hidden:true). The sync verb rejects stray positional args (cobra.NoArgs) rather than silently ignoring them, since the target is chosen via flags. Job compute resolution reads job-level environments/job_clusters; task-level compute is out of scope and returns an actionable error pointing at --cluster/--serverless. This is the first layer reachable from main, so it makes the whole libs/localenv package live for the deadcode checker without any pragmas. Build, unit tests, lint, and deadcode are clean; the hidden three-level command was smoke-tested end to end (help at each level, unknown-subcommand errors, flags). Co-authored-by: Isaac
1 parent d8f0de9 commit 2fd09fe

7 files changed

Lines changed: 746 additions & 0 deletions

File tree

cmd/cmd.go

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ import (
1818
"github.com/databricks/cli/cmd/experimental"
1919
"github.com/databricks/cli/cmd/fs"
2020
"github.com/databricks/cli/cmd/labs"
21+
"github.com/databricks/cli/cmd/localenv"
2122
"github.com/databricks/cli/cmd/pipelines"
2223
"github.com/databricks/cli/cmd/quickstart"
2324
"github.com/databricks/cli/cmd/root"
@@ -113,6 +114,7 @@ func New(ctx context.Context) *cobra.Command {
113114
cli.AddCommand(cache.New())
114115
cli.AddCommand(experimental.New())
115116
cli.AddCommand(psql.New())
117+
cli.AddCommand(localenv.New())
116118
cli.AddCommand(configure.New())
117119
cli.AddCommand(fs.New())
118120
cli.AddCommand(labs.New(ctx))

cmd/localenv/compute.go

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,65 @@
1+
package localenv
2+
3+
import (
4+
"context"
5+
"fmt"
6+
"strconv"
7+
8+
databricks "github.com/databricks/databricks-sdk-go"
9+
)
10+
11+
// sdkCompute adapts the Databricks SDK to the localenv.ComputeClient interface.
12+
type sdkCompute struct {
13+
w *databricks.WorkspaceClient
14+
}
15+
16+
// GetClusterSparkVersion returns the Spark version string for a running cluster.
17+
func (c sdkCompute) GetClusterSparkVersion(ctx context.Context, clusterID string) (string, error) {
18+
d, err := c.w.Clusters.GetByClusterId(ctx, clusterID)
19+
if err != nil {
20+
return "", fmt.Errorf("get cluster %s: %w", clusterID, err)
21+
}
22+
return d.SparkVersion, nil
23+
}
24+
25+
// GetJobSparkVersion inspects the job's configuration to determine compute type.
26+
//
27+
// A job is considered serverless when it has non-empty Environments (JobEnvironment
28+
// entries), which signals the Databricks serverless runtime. A job with classic compute
29+
// uses JobClusters; we read SparkVersion from the first job cluster's NewCluster spec.
30+
//
31+
// Task-level compute (tasks[].new_cluster / tasks[].existing_cluster_id with no
32+
// job-level job_clusters) is not resolved here: it may vary per task and an
33+
// existing_cluster_id would need a second lookup, which is out of scope for the
34+
// initial job support. Such a job returns an actionable error rather than a wrong
35+
// guess; use --cluster or --serverless explicitly instead.
36+
func (c sdkCompute) GetJobSparkVersion(ctx context.Context, jobID string) (sparkVersion string, isServerless bool, version string, err error) {
37+
id, err := strconv.ParseInt(jobID, 10, 64)
38+
if err != nil {
39+
return "", false, "", fmt.Errorf("invalid job ID %q: must be an integer: %w", jobID, err)
40+
}
41+
42+
job, err := c.w.Jobs.GetByJobId(ctx, id)
43+
if err != nil {
44+
return "", false, "", fmt.Errorf("get job %d: %w", id, err)
45+
}
46+
47+
if job.Settings == nil {
48+
return "", false, "", fmt.Errorf("job %d has no settings", id)
49+
}
50+
51+
// Serverless jobs have Environments populated; classic compute uses JobClusters.
52+
if len(job.Settings.Environments) > 0 {
53+
return "", true, "", nil
54+
}
55+
56+
if len(job.Settings.JobClusters) > 0 {
57+
sv := job.Settings.JobClusters[0].NewCluster.SparkVersion
58+
if sv == "" {
59+
return "", false, "", fmt.Errorf("could not determine compute for job %d: first job cluster has no spark_version", id)
60+
}
61+
return sv, false, sv, nil
62+
}
63+
64+
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)
65+
}

cmd/localenv/localenv.go

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
package localenv
2+
3+
import (
4+
"github.com/databricks/cli/cmd/root"
5+
libslocalenv "github.com/databricks/cli/libs/localenv"
6+
"github.com/spf13/cobra"
7+
)
8+
9+
// New returns the local-env command group. The group, subgroup, and verb names
10+
// come from the single command-name constants in libs/localenv so a rename is a
11+
// one-location change (spec §0 / invariant 8).
12+
//
13+
// The command is Hidden while the feature lands across the stacked PRs: it is
14+
// wired and runnable for dogfooding, but stays out of help and completion until
15+
// the final PR unveils it (removes this flag, adds the help line and changelog).
16+
func New() *cobra.Command {
17+
cmd := &cobra.Command{
18+
Use: libslocalenv.CommandGroup,
19+
Short: "Manage local development environments matched to Databricks compute",
20+
GroupID: "development",
21+
Hidden: true,
22+
Long: `Manage local development environments matched to a Databricks compute target.
23+
24+
Derives the Python version, databricks-connect version, and dependency
25+
constraints from the selected compute (cluster, serverless, or job) so that
26+
local resolution matches the Databricks runtime.`,
27+
RunE: root.ReportUnknownSubcommand,
28+
}
29+
cmd.AddCommand(newPythonCommand())
30+
return cmd
31+
}
32+
33+
// newPythonCommand returns the "python" subgroup. It is a parent-only node: with
34+
// no verb it reports an unknown-subcommand error (mirroring the generated command
35+
// groups) rather than doing nothing.
36+
func newPythonCommand() *cobra.Command {
37+
cmd := &cobra.Command{
38+
Use: libslocalenv.CommandSubgroup,
39+
Short: "Manage the local Python environment",
40+
Long: `Manage the local Python environment matched to a Databricks compute target.`,
41+
RunE: root.ReportUnknownSubcommand,
42+
}
43+
cmd.AddCommand(newSyncCommand())
44+
return cmd
45+
}

cmd/localenv/output.go

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
package localenv
2+
3+
import (
4+
"context"
5+
"fmt"
6+
7+
"github.com/databricks/cli/cmd/root"
8+
"github.com/databricks/cli/libs/cmdio"
9+
"github.com/databricks/cli/libs/flags"
10+
libslocalenv "github.com/databricks/cli/libs/localenv"
11+
"github.com/spf13/cobra"
12+
)
13+
14+
// renderResult renders the pipeline result to the command's output.
15+
// In JSON mode it renders the full structured result (even on error).
16+
// In text mode it prints phase headers and a summary, then returns the error.
17+
//
18+
// res is always non-nil: Pipeline.Run constructs and returns a fully-populated
19+
// Result (with the canonical phase list and error object) on every path,
20+
// including failures, so no nil guard is needed here.
21+
func renderResult(ctx context.Context, cmd *cobra.Command, res *libslocalenv.Result, pipelineErr error) error {
22+
if root.OutputType(cmd) == flags.OutputJSON {
23+
if err := cmdio.Render(ctx, res); err != nil {
24+
return err
25+
}
26+
// The JSON object is the only thing written to stdout. On failure we still
27+
// need a non-zero exit, but returning pipelineErr would make the root print
28+
// "Error: ..." to stderr. ErrAlreadyPrinted exits non-zero without that.
29+
if pipelineErr != nil {
30+
return root.ErrAlreadyPrinted
31+
}
32+
return nil
33+
}
34+
35+
// Text mode: print each phase in execution order.
36+
for _, phase := range res.Phases {
37+
if phase.Detail != "" {
38+
cmdio.LogString(ctx, fmt.Sprintf("%-10s %s %s", phase.Phase, phase.Status, phase.Detail))
39+
} else {
40+
cmdio.LogString(ctx, fmt.Sprintf("%-10s %s", phase.Phase, phase.Status))
41+
}
42+
}
43+
44+
for _, w := range res.Warnings {
45+
cmdio.LogString(ctx, "warning: "+w.Message)
46+
}
47+
48+
if pipelineErr != nil {
49+
cmdio.LogString(ctx, "For more detail, re-run with --debug, or --output json to share a structured report.")
50+
return pipelineErr
51+
}
52+
53+
// Print a final success / check summary.
54+
if res.DryRun {
55+
if res.Plan != nil {
56+
cmdio.LogString(ctx, "Plan: "+res.Plan.WouldWrite)
57+
for _, region := range res.Plan.ChangedRegions {
58+
cmdio.LogString(ctx, " changed region: "+region)
59+
}
60+
}
61+
cmdio.LogString(ctx, "Check complete. No files were modified.")
62+
return nil
63+
}
64+
65+
if res.Resolved != nil {
66+
summary := "Success: python=" + res.Resolved.PythonVersion
67+
if res.Resolved.DBConnectVersion != "" {
68+
summary += " databricks-connect=" + res.Resolved.DBConnectVersion
69+
}
70+
if res.VenvPath != "" {
71+
summary += " venv=" + res.VenvPath
72+
}
73+
cmdio.LogString(ctx, summary)
74+
}
75+
return nil
76+
}

cmd/localenv/sync.go

Lines changed: 154 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,154 @@
1+
package localenv
2+
3+
import (
4+
"context"
5+
"os"
6+
"path/filepath"
7+
8+
"github.com/databricks/cli/cmd/root"
9+
"github.com/databricks/cli/libs/cmdctx"
10+
"github.com/databricks/cli/libs/env"
11+
libslocalenv "github.com/databricks/cli/libs/localenv"
12+
"github.com/spf13/cobra"
13+
)
14+
15+
const (
16+
// defaultConstraintBaseURL is the default URL for the constraint source.
17+
defaultConstraintBaseURL = "https://raw.githubusercontent.com/rugpanov/databricks-environments/main"
18+
19+
// envConstraintSource is the environment variable for overriding the constraint source URL.
20+
envConstraintSource = "DATABRICKS_LOCALENV_CONSTRAINT_SOURCE"
21+
)
22+
23+
func newSyncCommand() *cobra.Command {
24+
cmd := &cobra.Command{
25+
Use: libslocalenv.CommandVerb,
26+
Short: "Provision a local Python environment matched to a Databricks compute target",
27+
Long: `Provision (or update) a local Python environment matched to a Databricks compute target.
28+
29+
Resolves the target to an environment key, fetches the pinned Python version,
30+
databricks-connect version, and dependency constraints published for that key,
31+
then provisions a matched .venv with uv. A project with no pyproject.toml is
32+
initialized from scratch; an existing pyproject.toml is merged in place (its
33+
env-owned sections are refreshed, user-owned content is preserved).`,
34+
}
35+
// The target is selected via flags; reject stray positional args rather than
36+
// silently ignoring them.
37+
cmd.Args = cobra.NoArgs
38+
cmd.PreRunE = root.MustWorkspaceClient
39+
addTargetFlags(cmd)
40+
cmd.RunE = func(cmd *cobra.Command, args []string) error {
41+
return runPipeline(cmd)
42+
}
43+
return cmd
44+
}
45+
46+
// addTargetFlags adds the shared target and mode flags to a command.
47+
func addTargetFlags(cmd *cobra.Command) {
48+
cmd.Flags().String("cluster", "", "cluster ID to use as the compute target")
49+
cmd.Flags().String("serverless", "", "serverless version to use as the compute target (e.g. v4)")
50+
cmd.Flags().String("job", "", "job ID to use as the compute target")
51+
cmd.Flags().Bool("constraints-only", false, "apply the Python version and constraints without adding the databricks-connect dependency")
52+
cmd.Flags().Bool("check", false, "compute the plan without writing files or provisioning")
53+
cmd.Flags().String("constraint-source", "", "URL for the constraint source (overrides "+envConstraintSource+")")
54+
// Hide constraint-source from casual --help output; it is a power-user escape hatch.
55+
_ = cmd.Flags().MarkHidden("constraint-source")
56+
cmd.MarkFlagsMutuallyExclusive("cluster", "serverless", "job")
57+
}
58+
59+
// runPipeline builds and runs the local-env Pipeline.
60+
func runPipeline(cmd *cobra.Command) error {
61+
ctx := cmd.Context()
62+
63+
cluster, _ := cmd.Flags().GetString("cluster")
64+
serverless, _ := cmd.Flags().GetString("serverless")
65+
job, _ := cmd.Flags().GetString("job")
66+
constraintsOnly, _ := cmd.Flags().GetBool("constraints-only")
67+
check, _ := cmd.Flags().GetBool("check")
68+
constraintSource, _ := cmd.Flags().GetString("constraint-source")
69+
70+
targetFlags := libslocalenv.TargetFlags{
71+
Cluster: cluster,
72+
Serverless: serverless,
73+
Job: job,
74+
}
75+
// ValidateTargetFlags is kept despite MarkFlagsMutuallyExclusive above:
76+
// it also validates the library path (no Cobra equivalent) and guards
77+
// non-Cobra call paths such as tests that invoke runPipeline directly.
78+
if err := libslocalenv.ValidateTargetFlags(targetFlags); err != nil {
79+
return err
80+
}
81+
82+
mode := libslocalenv.ModeDefault
83+
if constraintsOnly {
84+
mode = libslocalenv.ModeConstraintsOnly
85+
}
86+
87+
// Resolve constraint base URL: flag → env var → default constant.
88+
constraintBaseURL := resolveConstraintBaseURL(ctx, constraintSource)
89+
90+
projectDir, err := os.Getwd()
91+
if err != nil {
92+
return err
93+
}
94+
95+
cacheDir, err := os.UserCacheDir()
96+
if err != nil {
97+
return err
98+
}
99+
cacheDir = filepath.Join(cacheDir, "databricks", "localenv")
100+
101+
bt := bundleTarget(cmd)
102+
103+
w := cmdctx.WorkspaceClient(ctx)
104+
p := &libslocalenv.Pipeline{
105+
Mode: mode,
106+
Check: check,
107+
ProjectDir: projectDir,
108+
ConstraintBaseURL: constraintBaseURL,
109+
CacheDir: cacheDir,
110+
Flags: targetFlags,
111+
Compute: sdkCompute{w: w},
112+
Bundle: bt,
113+
PM: libslocalenv.NewUvManager(),
114+
}
115+
116+
res, pipelineErr := p.Run(ctx)
117+
return renderResult(ctx, cmd, res, pipelineErr)
118+
}
119+
120+
// resolveConstraintBaseURL returns the constraint base URL using ordered precedence:
121+
// flag → env var → default constant.
122+
func resolveConstraintBaseURL(ctx context.Context, flagValue string) string {
123+
if flagValue != "" {
124+
return flagValue
125+
}
126+
if v, ok := env.Lookup(ctx, envConstraintSource); ok {
127+
return v
128+
}
129+
return defaultConstraintBaseURL
130+
}
131+
132+
// bundleTarget reads the active bundle (if any) and maps its compute configuration
133+
// to a libslocalenv.BundleTarget.
134+
//
135+
// Only the top-level bundle.cluster_id field is consulted here; serverless is not
136+
// recorded in the bundle config, so Selected=true is set only when a cluster ID is
137+
// present. If the bundle is absent or has no cluster_id, Selected=false is returned
138+
// so the pipeline falls through to requiring an explicit flag.
139+
//
140+
// TODO: extend once bundle config exposes a serverless field at the bundle level.
141+
func bundleTarget(cmd *cobra.Command) libslocalenv.BundleTarget {
142+
b := root.TryConfigureBundle(cmd)
143+
if b == nil {
144+
return libslocalenv.BundleTarget{Selected: false}
145+
}
146+
clusterID := b.Config.Bundle.ClusterId
147+
if clusterID == "" {
148+
return libslocalenv.BundleTarget{Selected: false}
149+
}
150+
return libslocalenv.BundleTarget{
151+
ClusterID: clusterID,
152+
Selected: true,
153+
}
154+
}

0 commit comments

Comments
 (0)