Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions cmd/cmd.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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))
Expand Down
65 changes: 65 additions & 0 deletions cmd/localenv/compute.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,65 @@
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)
}

// 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)
}
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)
}
45 changes: 45 additions & 0 deletions cmd/localenv/localenv.go
Original file line number Diff line number Diff line change
@@ -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
}
76 changes: 76 additions & 0 deletions cmd/localenv/output.go
Original file line number Diff line number Diff line change
@@ -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
}
154 changes: 154 additions & 0 deletions cmd/localenv/sync.go
Original file line number Diff line number Diff line change
@@ -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"
)

const (
// defaultConstraintBaseURL is the default URL for the constraint source.
defaultConstraintBaseURL = "https://raw.githubusercontent.com/rugpanov/databricks-environments/main"

// envConstraintSource is the environment variable for overriding the constraint source URL.
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
}

// Resolve constraint base URL: flag → env var → default constant.
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:
// flag → env var → default constant.
func resolveConstraintBaseURL(ctx context.Context, flagValue string) string {
if flagValue != "" {
return flagValue
}
if v, ok := env.Lookup(ctx, envConstraintSource); ok {
return v
}
return defaultConstraintBaseURL
}

// 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,
}
}
Loading
Loading