[WIP] Add databricks dbconnect init / sync commands

NEXT_CHANGELOG.md

@@ -5,6 +5,7 @@
 ### Notable Changes
 
 ### CLI
+* Add `databricks dbconnect init` and `databricks dbconnect sync` to provision a local Python environment (Python version, `databricks-connect` pin, and dependency constraints) matched to the selected Databricks compute target.
 
 ### Bundles
 * `bundle run` now prints the modern job run URL (`/jobs/<id>/runs/<id>`) so that non-admin users permitted to view the run are taken to the run instead of the workspace homepage.

acceptance/dbconnect/cluster-unsupported/output.txt

@@ -0,0 +1,8 @@
+=== Phase 1: preflight ===
+    status=ok  uv [UV_VERSION]
+=== Phase 2: resolve ===
+    status=ok  kind=cluster envKey=dbr/15.4.x-scala2.12
+=== Phase 3: fetch ===
+    status=failed  fetch constraints for dbr/15.4.x-scala2.12: GET [DATABRICKS_URL]/dbr/15.4.x-scala2.12/pyproject.toml: unexpected status 404 Not Found
+For more detail, re-run with --debug, or --output json to share a structured report.
+Error: fetch constraints for dbr/15.4.x-scala2.12: GET [DATABRICKS_URL]/dbr/15.4.x-scala2.12/pyproject.toml: unexpected status 404 Not Found

acceptance/dbconnect/cluster-unsupported/script

@@ -0,0 +1 @@
+musterr $CLI dbconnect init --cluster test-cluster-id --check

acceptance/dbconnect/cluster-unsupported/test.toml

@@ -0,0 +1,22 @@
+EnvMatrix.DATABRICKS_BUNDLE_ENGINE = []
+
+[Env]
+DATABRICKS_DBCONNECT_CONSTRAINT_SOURCE = "$DATABRICKS_HOST"
+
+[[Server]]
+Pattern = "GET /api/2.1/clusters/get"
+Response.Body = '''
+{
+  "cluster_id": "test-cluster-id",
+  "spark_version": "15.4.x-scala2.12"
+}
+'''
+
+[[Server]]
+Pattern = "GET /dbr/15.4.x-scala2.12/pyproject.toml"
+Response.StatusCode = 404
+Response.Body = '{"message": "Not found"}'
+
+[[Repls]]
+Old = 'uv uv \S+(?: \([^)]+\))?'
+New = 'uv [UV_VERSION]'

acceptance/dbconnect/flag-conflict/output.txt

@@ -0,0 +1 @@
+Error: if any flags in the group [cluster serverless job] are set none of the others can be; [cluster serverless] were all set

acceptance/dbconnect/flag-conflict/script

@@ -0,0 +1 @@
+musterr $CLI dbconnect init --cluster abc --serverless v4

acceptance/dbconnect/flag-conflict/test.toml

@@ -0,0 +1 @@
+EnvMatrix.DATABRICKS_BUNDLE_ENGINE = []

acceptance/dbconnect/help/output.txt

@@ -0,0 +1,40 @@
+Set up a local Python environment matched to your 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.
+
+Usage:
+  databricks dbconnect [command]
+
+Available Commands:
+  init        Create a fresh pyproject.toml and provision a matched .venv
+  sync        Merge managed dependencies into an existing pyproject.toml and re-provision
+
+Flags:
+  -h, --help   help for dbconnect
+
+Global Flags:
+      --debug            enable debug logging
+  -o, --output type      output type: text or json (default text)
+  -p, --profile string   ~/.databrickscfg profile
+  -t, --target string    bundle target to use (if applicable)
+
+Use "databricks dbconnect [command] --help" for more information about a command.
+Create a fresh pyproject.toml and provision a matched .venv
+
+Usage:
+  databricks dbconnect init [flags]
+
+Flags:
+      --check               compute the plan without writing files or provisioning
+      --cluster string      cluster ID to use as the compute target
+  -h, --help                help for init
+      --job string          job ID to use as the compute target
+      --serverless string   serverless version to use as the compute target (e.g. v4)
+
+Global Flags:
+      --debug            enable debug logging
+  -o, --output type      output type: text or json (default text)
+  -p, --profile string   ~/.databrickscfg profile
+  -t, --target string    bundle target to use (if applicable)

acceptance/dbconnect/help/script

@@ -0,0 +1,2 @@
+$CLI dbconnect --help
+$CLI dbconnect init --help

acceptance/dbconnect/help/test.toml

@@ -0,0 +1 @@
+EnvMatrix.DATABRICKS_BUNDLE_ENGINE = []

acceptance/dbconnect/json-error/output.txt

@@ -0,0 +1,20 @@
+{
+  "mode": "init",
+  "check": false,
+  "phases": [
+    {
+      "name": "preflight",
+      "status": "ok",
+      "detail": "uv [UV_VERSION]"
+    },
+    {
+      "name": "resolve",
+      "status": "failed",
+      "detail": "No compute target is selected. Select a cluster or serverless target, or pass --cluster/--serverless/--job"
+    }
+  ],
+  "error": {
+    "code": "no_target_selected",
+    "message": "No compute target is selected. Select a cluster or serverless target, or pass --cluster/--serverless/--job"
+  }
+}

acceptance/dbconnect/json-error/script

@@ -0,0 +1 @@
+musterr $CLI dbconnect init --output json

acceptance/dbconnect/json-error/test.toml

@@ -0,0 +1,5 @@
+EnvMatrix.DATABRICKS_BUNDLE_ENGINE = []
+
+[[Repls]]
+Old = 'uv uv \S+(?: \([^)]+\))?'
+New = 'uv [UV_VERSION]'

acceptance/dbconnect/no-target/output.txt

@@ -0,0 +1,6 @@
+=== Phase 1: preflight ===
+    status=ok  uv [UV_VERSION]
+=== Phase 2: resolve ===
+    status=failed  No compute target is selected. Select a cluster or serverless target, or pass --cluster/--serverless/--job
+For more detail, re-run with --debug, or --output json to share a structured report.
+Error: No compute target is selected. Select a cluster or serverless target, or pass --cluster/--serverless/--job

acceptance/dbconnect/no-target/script

@@ -0,0 +1 @@
+musterr $CLI dbconnect init

acceptance/dbconnect/no-target/test.toml

@@ -0,0 +1,5 @@
+EnvMatrix.DATABRICKS_BUNDLE_ENGINE = []
+
+[[Repls]]
+Old = 'uv uv \S+(?: \([^)]+\))?'
+New = 'uv [UV_VERSION]'

acceptance/dbconnect/serverless-check/output.txt

@@ -0,0 +1,17 @@
+
+>>> [CLI] dbconnect init --serverless v4 --check
+=== Phase 1: preflight ===
+    status=ok  uv [UV_VERSION]
+=== Phase 2: resolve ===
+    status=ok  kind=serverless envKey=serverless/serverless-v4
+=== Phase 3: fetch ===
+    status=ok  source=[DATABRICKS_URL]/serverless/serverless-v4/pyproject.toml fromCache=false
+=== Phase 4: parse-python-version ===
+    status=ok  3.12
+=== Phase 5: plan ===
+    status=ok  changed=requires-python,databricks-connect,tool.uv.constraint-dependencies
+Plan: [TEST_TMP_DIR]/pyproject.toml
+  changed region: requires-python
+  changed region: databricks-connect
+  changed region: tool.uv.constraint-dependencies
+Check complete. No files were modified.

acceptance/dbconnect/serverless-check/script

@@ -0,0 +1 @@
+trace $CLI dbconnect init --serverless v4 --check

acceptance/dbconnect/serverless-check/test.toml

@@ -0,0 +1,21 @@
+EnvMatrix.DATABRICKS_BUNDLE_ENGINE = []
+
+[Env]
+DATABRICKS_DBCONNECT_CONSTRAINT_SOURCE = "$DATABRICKS_HOST"
+
+[[Server]]
+Pattern = "GET /serverless/serverless-v4/pyproject.toml"
+Response.Body = '''
+[project]
+requires-python = ">=3.12"
+
+[dependency-groups]
+dev = ["databricks-connect~=17.2.0"]
+
+[tool.uv]
+constraint-dependencies = ["pyarrow<19", "pandas<3"]
+'''
+
+[[Repls]]
+Old = 'uv uv \S+(?: \([^)]+\))?'
+New = 'uv [UV_VERSION]'

acceptance/dbconnect/serverless-json/output.txt

@@ -0,0 +1,57 @@
+
+>>> [CLI] dbconnect init --serverless v4 --check --output json
+{
+  "mode": "init",
+  "check": true,
+  "target": {
+    "kind": "serverless",
+    "cluster_id": "",
+    "spark_version": "",
+    "env_key": "serverless/serverless-v4",
+    "python_version": "3.12"
+  },
+  "constraints": {
+    "source_url": "[DATABRICKS_URL]/serverless/serverless-v4/pyproject.toml",
+    "from_cache": false,
+    "requires_python": "\u003e=3.12",
+    "databricks_connect": "databricks-connect~=17.2.0",
+    "constraint_count": 2
+  },
+  "plan": {
+    "pyproject_path": "[TEST_TMP_DIR]/pyproject.toml",
+    "backup_path": "[TEST_TMP_DIR]/pyproject.toml.bak",
+    "diff": "--- pyproject.toml\n+++ pyproject.toml\n@@ -1 +1,16 @@\n+[project]\n+name = \"001\"\n+requires-python = \"\u003e=3.12\"\n+\n+[dependency-groups]\n+dev = [\n+    \"databricks-connect~=17.2.0\",\n+]\n+\n+# managed by databricks dbconnect — do not edit\n+[tool.uv]\n+constraint-dependencies = [\n+    \"pyarrow\u003c19\",\n+    \"pandas\u003c3\",\n+]\n+# end managed by databricks dbconnect\n",
+    "changed_regions": [
+      "requires-python",
+      "databricks-connect",
+      "tool.uv.constraint-dependencies"
+    ]
+  },
+  "phases": [
+    {
+      "name": "preflight",
+      "status": "ok",
+      "detail": "uv [UV_VERSION]"
+    },
+    {
+      "name": "resolve",
+      "status": "ok",
+      "detail": "kind=serverless envKey=serverless/serverless-v4"
+    },
+    {
+      "name": "fetch",
+      "status": "ok",
+      "detail": "source=[DATABRICKS_URL]/serverless/serverless-v4/pyproject.toml fromCache=false"
+    },
+    {
+      "name": "parse-python-version",
+      "status": "ok",
+      "detail": "3.12"
+    },
+    {
+      "name": "plan",
+      "status": "ok",
+      "detail": "changed=requires-python,databricks-connect,tool.uv.constraint-dependencies"
+    }
+  ]
+}

acceptance/dbconnect/serverless-json/script

@@ -0,0 +1 @@
+trace $CLI dbconnect init --serverless v4 --check --output json

acceptance/dbconnect/serverless-json/test.toml

@@ -0,0 +1,21 @@
+EnvMatrix.DATABRICKS_BUNDLE_ENGINE = []
+
+[Env]
+DATABRICKS_DBCONNECT_CONSTRAINT_SOURCE = "$DATABRICKS_HOST"
+
+[[Server]]
+Pattern = "GET /serverless/serverless-v4/pyproject.toml"
+Response.Body = '''
+[project]
+requires-python = ">=3.12"
+
+[dependency-groups]
+dev = ["databricks-connect~=17.2.0"]
+
+[tool.uv]
+constraint-dependencies = ["pyarrow<19", "pandas<3"]
+'''
+
+[[Repls]]
+Old = 'uv uv \S+(?: \([^)]+\))?'
+New = 'uv [UV_VERSION]'

cmd/cmd.go

@@ -15,6 +15,7 @@ import (
 	"github.com/databricks/cli/cmd/cache"
 	"github.com/databricks/cli/cmd/completion"
 	"github.com/databricks/cli/cmd/configure"
+	"github.com/databricks/cli/cmd/dbconnect"
 	"github.com/databricks/cli/cmd/experimental"
 	"github.com/databricks/cli/cmd/fs"
 	"github.com/databricks/cli/cmd/labs"
@@ -120,6 +121,7 @@ func New(ctx context.Context) *cobra.Command {
 	cli.AddCommand(cache.New())
 	cli.AddCommand(experimental.New())
 	cli.AddCommand(psql.New())
+	cli.AddCommand(dbconnect.New())
 	cli.AddCommand(configure.New())
 	cli.AddCommand(fs.New())
 	cli.AddCommand(labs.New(ctx))

cmd/dbconnect/compute.go

@@ -0,0 +1,60 @@
+package dbconnect
+
+import (
+	"context"
+	"fmt"
+	"strconv"
+
+	databricks "github.com/databricks/databricks-sdk-go"
+)
+
+// sdkCompute adapts the Databricks SDK to the dbconnect.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.
+// If neither indicator is present the job's compute cannot be determined.
+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: no environments or job clusters found", id)
+}

cmd/dbconnect/dbconnect.go

@@ -0,0 +1,20 @@
+package dbconnect
+
+import "github.com/spf13/cobra"
+
+// New returns the `dbconnect` command group.
+func New() *cobra.Command {
+	cmd := &cobra.Command{
+		Use:     "dbconnect",
+		Short:   "Set up a local Python environment matched to your Databricks compute",
+		GroupID: "development",
+		Long: `Set up a local Python environment matched to your 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.`,
+	}
+	cmd.AddCommand(newInitCommand())
+	cmd.AddCommand(newSyncCommand())
+	return cmd
+}