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