diff --git a/.nextchanges/cli/setup-local-uv-workspace.md b/.nextchanges/cli/setup-local-uv-workspace.md new file mode 100644 index 00000000000..6c692161aa6 --- /dev/null +++ b/.nextchanges/cli/setup-local-uv-workspace.md @@ -0,0 +1 @@ +* `databricks environments setup-local` now supports uv workspace members: it uses the shared `.venv` at the workspace root and writes the dependency constraints to the root `pyproject.toml`. ([#7018](https://github.com/databricks/cli/pull/7018)) diff --git a/acceptance/localenv/help/output.txt b/acceptance/localenv/help/output.txt index af1e32e15df..98abac079af 100644 --- a/acceptance/localenv/help/output.txt +++ b/acceptance/localenv/help/output.txt @@ -2,6 +2,8 @@ Set up a local Python environment that matches a Databricks cluster or serverles Use this when you want to develop or debug a Databricks project locally: it installs the matching Python version and a compatible databricks-connect, and pins your dependencies to versions known to work with your chosen compute. It creates or updates a .venv (managed by uv) in the current directory and records the setup in pyproject.toml, leaving the rest of your project untouched. +In a uv workspace member, the command uses the shared .venv at the workspace root and writes the dependency constraints to the root pyproject.toml, because uv reads them only from there. uv sync installs the current member's dependencies into the shared .venv and removes packages that only other members need. All members share one environment, so they must target compatible compute. + Usage: databricks environments setup-local [flags] diff --git a/acceptance/localenv/workspace-discovery-unsupported-check/bundles/a/pyproject.toml b/acceptance/localenv/workspace-discovery-unsupported-check/bundles/a/pyproject.toml new file mode 100644 index 00000000000..962b088a2b0 --- /dev/null +++ b/acceptance/localenv/workspace-discovery-unsupported-check/bundles/a/pyproject.toml @@ -0,0 +1,4 @@ +[project] +name = "a" +version = "0.1.0" +requires-python = ">=3.10" diff --git a/acceptance/localenv/workspace-discovery-unsupported-check/fake-uv.sh b/acceptance/localenv/workspace-discovery-unsupported-check/fake-uv.sh new file mode 100644 index 00000000000..9e18385562b --- /dev/null +++ b/acceptance/localenv/workspace-discovery-unsupported-check/fake-uv.sh @@ -0,0 +1,8 @@ +#!/bin/bash +# uv older than 0.9.9 has no `workspace` subcommand. +if [ "$1" = "workspace" ]; then + echo "error: unrecognized subcommand 'workspace'" >&2 + exit 2 +fi +echo "fake uv: unexpected call: $*" >&2 +exit 1 diff --git a/acceptance/localenv/workspace-discovery-unsupported-check/out.test.toml b/acceptance/localenv/workspace-discovery-unsupported-check/out.test.toml new file mode 100644 index 00000000000..2e1543aa802 --- /dev/null +++ b/acceptance/localenv/workspace-discovery-unsupported-check/out.test.toml @@ -0,0 +1,2 @@ +Cloud = false +GOOS.windows = false diff --git a/acceptance/localenv/workspace-discovery-unsupported-check/output.txt b/acceptance/localenv/workspace-discovery-unsupported-check/output.txt new file mode 100644 index 00000000000..dcb8fe5bc95 --- /dev/null +++ b/acceptance/localenv/workspace-discovery-unsupported-check/output.txt @@ -0,0 +1,60 @@ + +>>> [CLI] environments setup-local --serverless-version 4 --dry-run --output json +{ + "schemaVersion": 1, + "command": "environments setup-local", + "ok": true, + "mode": "default", + "dryRun": true, + "compute": { + "source": "serverless", + "serverlessVersion": "v4", + "envKey": "serverless/serverless-v4" + }, + "resolved": { + "pythonVersion": "3.12", + "dbconnectVersion": "17.2.0", + "artifactSource": "network" + }, + "greenfield": false, + "plan": { + "wouldWrite": "[TEST_TMP_DIR]/bundles/a/pyproject.toml", + "wouldBackup": "[TEST_TMP_DIR]/bundles/a/pyproject.toml.bak", + "wouldInstallPython": "3.12", + "diff": "--- pyproject.toml\n+++ pyproject.toml.new\n@@ -1,4 +1,20 @@\n [project]\n name = \"a\"\n version = \"0.1.0\"\n-requires-python = \"\u003e=3.10\"\n+requires-python = \"\u003e=3.12\"\n+\n+[dependency-groups]\n+dev = [\n+ \"databricks-connect~=17.2.0\",\n+]\n+\n+[tool.databricks.environment]\n+environment_version = \"4\"\n+\n+# managed by databricks environments setup-local — do not edit\n+[tool.uv]\n+constraint-dependencies = [\n+ \"pyarrow\u003c19\",\n+ \"pandas\u003c3\",\n+]\n+# end managed by databricks environments setup-local\n" + }, + "phases": [ + { + "phase": "preflight", + "status": "ok" + }, + { + "phase": "resolve", + "status": "ok" + }, + { + "phase": "fetch", + "status": "ok" + }, + { + "phase": "merge", + "status": "ok" + }, + { + "phase": "provision", + "status": "ok" + }, + { + "phase": "validate", + "status": "ok" + } + ], + "warnings": [ + { + "code": "W_REQUIRES_PYTHON_OVERRIDDEN", + "message": "requires-python \"\u003e=3.10\" is replaced by the environment's \"\u003e=3.12\"" + } + ], + "error": null, + "durationMs": [DURATION_MS] +} diff --git a/acceptance/localenv/workspace-discovery-unsupported-check/pyproject.toml b/acceptance/localenv/workspace-discovery-unsupported-check/pyproject.toml new file mode 100644 index 00000000000..7a29f61f24a --- /dev/null +++ b/acceptance/localenv/workspace-discovery-unsupported-check/pyproject.toml @@ -0,0 +1,7 @@ +[project] +name = "root" +version = "0.1.0" +requires-python = ">=3.10" + +[tool.uv.workspace] +members = ["bundles/*"] diff --git a/acceptance/localenv/workspace-discovery-unsupported-check/script b/acceptance/localenv/workspace-discovery-unsupported-check/script new file mode 100644 index 00000000000..f2c37466149 --- /dev/null +++ b/acceptance/localenv/workspace-discovery-unsupported-check/script @@ -0,0 +1,11 @@ +# uv older than 0.9.9 cannot report the workspace root, so in a workspace member +# the plan is the same as for a standalone project: no workspace fields. +mv fake-uv.sh uv +chmod +x uv +cleanup() { + rm -f uv +} +trap cleanup EXIT +export PATH="$(pwd):$PATH" + +withdir bundles/a trace $CLI environments setup-local --serverless-version 4 --dry-run --output json diff --git a/acceptance/localenv/workspace-discovery-unsupported-check/test.toml b/acceptance/localenv/workspace-discovery-unsupported-check/test.toml new file mode 100644 index 00000000000..1f85ab26622 --- /dev/null +++ b/acceptance/localenv/workspace-discovery-unsupported-check/test.toml @@ -0,0 +1,18 @@ +# The fake uv is a bash script, so this test cannot run on Windows runners +# (mirrors acceptance/cmd/psql). +GOOS.windows = false + +Env.DATABRICKS_LOCALENV_CONSTRAINT_SOURCE_URL_TEST_OVERRIDE = "$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"] +''' diff --git a/acceptance/localenv/workspace-member-check/bundles/a/pyproject.toml b/acceptance/localenv/workspace-member-check/bundles/a/pyproject.toml new file mode 100644 index 00000000000..962b088a2b0 --- /dev/null +++ b/acceptance/localenv/workspace-member-check/bundles/a/pyproject.toml @@ -0,0 +1,4 @@ +[project] +name = "a" +version = "0.1.0" +requires-python = ">=3.10" diff --git a/acceptance/localenv/workspace-member-check/fake-uv.sh b/acceptance/localenv/workspace-member-check/fake-uv.sh new file mode 100644 index 00000000000..34452908cdd --- /dev/null +++ b/acceptance/localenv/workspace-member-check/fake-uv.sh @@ -0,0 +1,8 @@ +#!/bin/bash +# The workspace root is the test folder, where this script is installed as uv. +if [ "$1" = "workspace" ] && [ "$2" = "dir" ]; then + cd "$(dirname "$0")" && pwd + exit 0 +fi +echo "fake uv: unexpected call: $*" >&2 +exit 1 diff --git a/acceptance/localenv/workspace-member-check/out.test.toml b/acceptance/localenv/workspace-member-check/out.test.toml new file mode 100644 index 00000000000..2e1543aa802 --- /dev/null +++ b/acceptance/localenv/workspace-member-check/out.test.toml @@ -0,0 +1,2 @@ +Cloud = false +GOOS.windows = false diff --git a/acceptance/localenv/workspace-member-check/output.txt b/acceptance/localenv/workspace-member-check/output.txt new file mode 100644 index 00000000000..b9fa432530e --- /dev/null +++ b/acceptance/localenv/workspace-member-check/output.txt @@ -0,0 +1,73 @@ + +>>> [CLI] environments setup-local --serverless-version 4 --dry-run +warning: requires-python ">=3.10" is replaced by the environment's ">=3.12" +Plan: [TEST_TMP_DIR]/bundles/a/pyproject.toml +Plan: [TEST_TMP_DIR]/pyproject.toml (uv workspace root) + changed region: requires-python + changed region: databricks-connect + changed region: tool.databricks.environment + changed region: tool.uv.constraint-dependencies (uv workspace root) +Check complete. No files were modified. + +>>> [CLI] environments setup-local --serverless-version 4 --dry-run --output json +{ + "schemaVersion": 1, + "command": "environments setup-local", + "ok": true, + "mode": "default", + "dryRun": true, + "compute": { + "source": "serverless", + "serverlessVersion": "v4", + "envKey": "serverless/serverless-v4" + }, + "resolved": { + "pythonVersion": "3.12", + "dbconnectVersion": "17.2.0", + "artifactSource": "network" + }, + "greenfield": false, + "plan": { + "wouldWrite": "[TEST_TMP_DIR]/bundles/a/pyproject.toml", + "wouldBackup": "[TEST_TMP_DIR]/bundles/a/pyproject.toml.bak", + "wouldWriteUvWorkspace": "[TEST_TMP_DIR]/pyproject.toml", + "wouldBackupUvWorkspace": "[TEST_TMP_DIR]/pyproject.toml.bak", + "wouldInstallPython": "3.12", + "diff": "--- pyproject.toml\n+++ pyproject.toml.new\n@@ -1,4 +1,12 @@\n [project]\n name = \"a\"\n version = \"0.1.0\"\n-requires-python = \"\u003e=3.10\"\n+requires-python = \"\u003e=3.12\"\n+\n+[dependency-groups]\n+dev = [\n+ \"databricks-connect~=17.2.0\",\n+]\n+\n+[tool.databricks.environment]\n+environment_version = \"4\"\n--- uv workspace root pyproject.toml\n+++ uv workspace root pyproject.toml.new\n@@ -5,3 +5,11 @@\n \n [tool.uv.workspace]\n members = [\"bundles/*\"]\n+\n+# managed by databricks environments setup-local — do not edit\n+[tool.uv]\n+constraint-dependencies = [\n+ \"pyarrow\u003c19\",\n+ \"pandas\u003c3\",\n+]\n+# end managed by databricks environments setup-local\n" + }, + "uvWorkspaceRoot": "[TEST_TMP_DIR]", + "phases": [ + { + "phase": "preflight", + "status": "ok" + }, + { + "phase": "resolve", + "status": "ok" + }, + { + "phase": "fetch", + "status": "ok" + }, + { + "phase": "merge", + "status": "ok" + }, + { + "phase": "provision", + "status": "ok" + }, + { + "phase": "validate", + "status": "ok" + } + ], + "warnings": [ + { + "code": "W_REQUIRES_PYTHON_OVERRIDDEN", + "message": "requires-python \"\u003e=3.10\" is replaced by the environment's \"\u003e=3.12\"" + } + ], + "error": null, + "durationMs": [DURATION_MS] +} diff --git a/acceptance/localenv/workspace-member-check/pyproject.toml b/acceptance/localenv/workspace-member-check/pyproject.toml new file mode 100644 index 00000000000..7a29f61f24a --- /dev/null +++ b/acceptance/localenv/workspace-member-check/pyproject.toml @@ -0,0 +1,7 @@ +[project] +name = "root" +version = "0.1.0" +requires-python = ">=3.10" + +[tool.uv.workspace] +members = ["bundles/*"] diff --git a/acceptance/localenv/workspace-member-check/script b/acceptance/localenv/workspace-member-check/script new file mode 100644 index 00000000000..ec88e4f62e2 --- /dev/null +++ b/acceptance/localenv/workspace-member-check/script @@ -0,0 +1,13 @@ +# A uv workspace member: the plan writes the member pyproject.toml and the +# constraint block into the workspace root pyproject.toml. A fake uv on PATH +# answers `uv workspace dir` (see acceptance/cmd/psql for the pattern). +mv fake-uv.sh uv +chmod +x uv +cleanup() { + rm -f uv +} +trap cleanup EXIT +export PATH="$(pwd):$PATH" + +withdir bundles/a trace $CLI environments setup-local --serverless-version 4 --dry-run +withdir bundles/a trace $CLI environments setup-local --serverless-version 4 --dry-run --output json diff --git a/acceptance/localenv/workspace-member-check/test.toml b/acceptance/localenv/workspace-member-check/test.toml new file mode 100644 index 00000000000..1f85ab26622 --- /dev/null +++ b/acceptance/localenv/workspace-member-check/test.toml @@ -0,0 +1,18 @@ +# The fake uv is a bash script, so this test cannot run on Windows runners +# (mirrors acceptance/cmd/psql). +GOOS.windows = false + +Env.DATABRICKS_LOCALENV_CONSTRAINT_SOURCE_URL_TEST_OVERRIDE = "$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"] +''' diff --git a/acceptance/localenv/workspace-member/bundles/a/pyproject.toml b/acceptance/localenv/workspace-member/bundles/a/pyproject.toml new file mode 100644 index 00000000000..962b088a2b0 --- /dev/null +++ b/acceptance/localenv/workspace-member/bundles/a/pyproject.toml @@ -0,0 +1,4 @@ +[project] +name = "a" +version = "0.1.0" +requires-python = ">=3.10" diff --git a/acceptance/localenv/workspace-member/fake-uv.sh b/acceptance/localenv/workspace-member/fake-uv.sh new file mode 100644 index 00000000000..9b63430c86e --- /dev/null +++ b/acceptance/localenv/workspace-member/fake-uv.sh @@ -0,0 +1,24 @@ +#!/bin/bash +# The workspace root is the test folder, where this script is installed as uv. +# `uv sync` creates the shared venv there with a fake interpreter that answers +# the validate probe (prefixes from libs/localenv/uv.go). +root="$(cd "$(dirname "$0")" && pwd)" +case "$1 $2" in + "--version ") echo "uv 0.12.9 (fake)"; exit 0 ;; + "workspace dir") echo "$root"; exit 0 ;; + "python install") exit 0 ;; + "pip install") exit 0 ;; + "sync --python") + mkdir -p "$root/.venv/bin" + cat > "$root/.venv/bin/python" <<'PY' +#!/bin/bash +echo "PYVER:3.12" +echo "DBCVER:17.2.0" +echo "PYSPARKVER:" +echo "DBCIMPORT:" +PY + chmod +x "$root/.venv/bin/python" + exit 0 ;; +esac +echo "fake uv: unexpected call: $*" >&2 +exit 1 diff --git a/acceptance/localenv/workspace-member/out.test.toml b/acceptance/localenv/workspace-member/out.test.toml new file mode 100644 index 00000000000..2e1543aa802 --- /dev/null +++ b/acceptance/localenv/workspace-member/out.test.toml @@ -0,0 +1,2 @@ +Cloud = false +GOOS.windows = false diff --git a/acceptance/localenv/workspace-member/output.txt b/acceptance/localenv/workspace-member/output.txt new file mode 100644 index 00000000000..5a05c6f4c08 --- /dev/null +++ b/acceptance/localenv/workspace-member/output.txt @@ -0,0 +1,105 @@ + +>>> [CLI] environments setup-local --serverless-version 4 --output json +{ + "schemaVersion": 1, + "command": "environments setup-local", + "ok": true, + "mode": "default", + "dryRun": false, + "compute": { + "source": "serverless", + "serverlessVersion": "v4", + "envKey": "serverless/serverless-v4" + }, + "resolved": { + "pythonVersion": "3.12", + "dbconnectVersion": "17.2.0", + "artifactSource": "network" + }, + "greenfield": false, + "venvPath": "../../.venv", + "uvWorkspaceRoot": "[TEST_TMP_DIR]", + "pythonResolution": "uv_install_succeeded", + "phases": [ + { + "phase": "preflight", + "status": "ok" + }, + { + "phase": "resolve", + "status": "ok" + }, + { + "phase": "fetch", + "status": "ok" + }, + { + "phase": "merge", + "status": "ok" + }, + { + "phase": "provision", + "status": "ok" + }, + { + "phase": "validate", + "status": "ok" + } + ], + "warnings": [ + { + "code": "W_REQUIRES_PYTHON_OVERRIDDEN", + "message": "requires-python \"\u003e=3.10\" is replaced by the environment's \"\u003e=3.12\"" + } + ], + "error": null, + "backupPath": "[TEST_TMP_DIR]/bundles/a/pyproject.toml.bak", + "uvWorkspaceBackupPath": "[TEST_TMP_DIR]/pyproject.toml.bak", + "durationMs": [DURATION_MS] +} + +>>> [CLI] environments setup-local --serverless-version 4 +✔ Local environment ready + + Compute target serverless 4 + Python 3.12 + databricks-connect 17.2.0 + Virtual env ../../.venv + pyproject.toml updated + uv workspace root [TEST_TMP_DIR] + Root pyproject.toml unchanged + +Next steps: + • Activate it: source ../../.venv/bin/activate + • Or select ../../.venv as the Python interpreter in VS Code / Cursor + +>>> cat pyproject.toml +[project] +name = "root" +version = "0.1.0" +requires-python = ">=3.10" + +[tool.uv.workspace] +members = ["bundles/*"] + +# managed by databricks environments setup-local — do not edit +[tool.uv] +constraint-dependencies = [ + "pyarrow<19", + "pandas<3", +] +# end managed by databricks environments setup-local + +>>> cat bundles/a/pyproject.toml +[project] +name = "a" +version = "0.1.0" +requires-python = ">=3.12" + +[dependency-groups] +dev = [ + "databricks-connect~=17.2.0", +] + +[tool.databricks.environment] +environment_version = "4" diff --git a/acceptance/localenv/workspace-member/pyproject.toml b/acceptance/localenv/workspace-member/pyproject.toml new file mode 100644 index 00000000000..7a29f61f24a --- /dev/null +++ b/acceptance/localenv/workspace-member/pyproject.toml @@ -0,0 +1,7 @@ +[project] +name = "root" +version = "0.1.0" +requires-python = ">=3.10" + +[tool.uv.workspace] +members = ["bundles/*"] diff --git a/acceptance/localenv/workspace-member/script b/acceptance/localenv/workspace-member/script new file mode 100644 index 00000000000..0e2f3b7218f --- /dev/null +++ b/acceptance/localenv/workspace-member/script @@ -0,0 +1,17 @@ +# A real run in a uv workspace member provisions the shared .venv at the +# workspace root and writes the constraint block into the root pyproject.toml. +# The member keeps requires-python and the databricks-connect pin. A fake uv on +# PATH stands in for the network-dependent sync (see acceptance/cmd/psql). +mv fake-uv.sh uv +chmod +x uv +cleanup() { + rm -f uv +} +trap cleanup EXIT +export PATH="$(pwd):$PATH" + +withdir bundles/a trace $CLI environments setup-local --serverless-version 4 --output json +# A re-run changes nothing, and the text summary says so for the root. +withdir bundles/a trace $CLI environments setup-local --serverless-version 4 +trace cat pyproject.toml +trace cat bundles/a/pyproject.toml diff --git a/acceptance/localenv/workspace-member/test.toml b/acceptance/localenv/workspace-member/test.toml new file mode 100644 index 00000000000..0f62f56b8cc --- /dev/null +++ b/acceptance/localenv/workspace-member/test.toml @@ -0,0 +1,22 @@ +# The fake uv is a bash script, so this test cannot run on Windows runners +# (mirrors acceptance/cmd/psql). +GOOS.windows = false + +# The run writes both pyproject.toml files (checked with cat below), their +# backups, and the shared .venv at the workspace root. +Ignore = [".venv", "pyproject.toml.bak", "bundles/a/pyproject.toml.bak"] + +Env.DATABRICKS_LOCALENV_CONSTRAINT_SOURCE_URL_TEST_OVERRIDE = "$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"] +''' diff --git a/cmd/environments/output.go b/cmd/environments/output.go index d7307a18eb0..f58a72b56ea 100644 --- a/cmd/environments/output.go +++ b/cmd/environments/output.go @@ -3,6 +3,7 @@ package environments import ( "context" "fmt" + "path/filepath" "runtime" "strings" @@ -75,6 +76,9 @@ func renderResult(ctx context.Context, cmd *cobra.Command, res *libslocalenv.Res if res.DryRun { if res.Plan != nil { cmdio.LogString(ctx, "Plan: "+res.Plan.WouldWrite) + if res.Plan.WouldWriteUvWorkspace != "" { + cmdio.LogString(ctx, "Plan: "+res.Plan.WouldWriteUvWorkspace+" (uv workspace root)") + } for _, region := range res.Plan.ChangedRegions { cmdio.LogString(ctx, " changed region: "+region) } @@ -114,6 +118,16 @@ func renderSuccess(ctx context.Context, res *libslocalenv.Result) { pyprojectDetail = "updated (backup: " + res.BackupPath + ")" } cmdio.LogString(ctx, fmt.Sprintf(" %-20s%s", "pyproject.toml", pyprojectDetail)) + if res.UvWorkspaceRoot != "" { + // The root file changes only when its constraint block changes, and every + // change writes a backup; --no-constraints and a no-op re-run leave it alone. + rootDetail := "unchanged" + if res.UvWorkspaceBackupPath != "" { + rootDetail = "constraints updated (backup: " + res.UvWorkspaceBackupPath + ")" + } + cmdio.LogString(ctx, fmt.Sprintf(" %-20s%s", "uv workspace root", res.UvWorkspaceRoot)) + cmdio.LogString(ctx, fmt.Sprintf(" %-20s%s", "Root pyproject.toml", rootDetail)) + } renderInstalledPythonFallback(ctx, res) cmdio.LogString(ctx, "") @@ -139,7 +153,9 @@ func renderInstalledPythonFallback(ctx context.Context, res *libslocalenv.Result // that fails with "source is not recognized" and a path that does not exist. func activateHint(venvPath string) string { if runtime.GOOS == "windows" { - return venvPath + `\Scripts\activate` + // venvPath is in slash form; cmd.exe reads "/" as a switch character, so a + // relative path such as ../../.venv needs backslashes. + return filepath.FromSlash(venvPath) + `\Scripts\activate` } return "source " + venvPath + "/bin/activate" } diff --git a/cmd/environments/output_test.go b/cmd/environments/output_test.go index e9df52b4bb3..05efa5a0305 100644 --- a/cmd/environments/output_test.go +++ b/cmd/environments/output_test.go @@ -142,6 +142,16 @@ func TestRenderDryRun(t *testing.T) { assert.Contains(t, out, "No files were modified") } +func TestActivateHintWorkspaceMember(t *testing.T) { + // venvPath is in slash form; cmd.exe needs backslashes in a relative path. + hint := activateHint("../../.venv") + if runtime.GOOS == "windows" { + assert.Equal(t, `..\..\.venv\Scripts\activate`, hint) + } else { + assert.Equal(t, "source ../../.venv/bin/activate", hint) + } +} + func TestActivateHint(t *testing.T) { hint := activateHint(".venv") if runtime.GOOS == "windows" { @@ -154,3 +164,53 @@ func TestActivateHint(t *testing.T) { assert.Equal(t, "source .venv/bin/activate", hint) } } + +func TestRenderSuccessWorkspaceMember(t *testing.T) { + res := libslocalenv.NewResult() + res.OK = true + res.Resolved = &libslocalenv.ResolvedInfo{PythonVersion: "3.12"} + res.VenvPath = "../../.venv" + res.UvWorkspaceRoot = "/repo" + res.UvWorkspaceBackupPath = "/repo/pyproject.toml.bak" + + out := renderText(t, res, nil) + + assert.Contains(t, out, "uv workspace root") + assert.Contains(t, out, "/repo") + assert.Contains(t, out, "Root pyproject.toml") + assert.Contains(t, out, "constraints updated (backup: /repo/pyproject.toml.bak)") +} + +func TestRenderSuccessStandaloneOmitsWorkspaceRows(t *testing.T) { + res := libslocalenv.NewResult() + res.OK = true + res.VenvPath = ".venv" + + out := renderText(t, res, nil) + + assert.NotContains(t, out, "uv workspace root") + assert.NotContains(t, out, "Root pyproject.toml") +} + +func TestRenderDryRunWorkspaceMember(t *testing.T) { + res := libslocalenv.NewResult() + res.OK = true + res.DryRun = true + res.Plan = &libslocalenv.Plan{WouldWrite: "/repo/a/pyproject.toml", WouldWriteUvWorkspace: "/repo/pyproject.toml"} + out := renderText(t, res, nil) + assert.Contains(t, out, "Plan: /repo/pyproject.toml (uv workspace root)") +} + +func TestRenderSuccessWorkspaceMemberUnchangedRoot(t *testing.T) { + // --no-constraints or an idempotent re-run: the root file was not changed. + res := libslocalenv.NewResult() + res.OK = true + res.VenvPath = "../../.venv" + res.UvWorkspaceRoot = "/repo" + + out := renderText(t, res, nil) + + assert.Contains(t, out, "Root pyproject.toml") + assert.Contains(t, out, "unchanged") + assert.NotContains(t, out, "constraints updated") +} diff --git a/cmd/environments/sync.go b/cmd/environments/sync.go index 779fea38182..b4e40474a03 100644 --- a/cmd/environments/sync.go +++ b/cmd/environments/sync.go @@ -22,7 +22,9 @@ func newSetupLocalCommand() *cobra.Command { Short: "Set up a local Python environment that matches your Databricks compute", Long: `Set up a local Python environment that matches a Databricks cluster or serverless version, so code you run on your machine behaves the same as it does on Databricks. -Use this when you want to develop or debug a Databricks project locally: it installs the matching Python version and a compatible databricks-connect, and pins your dependencies to versions known to work with your chosen compute. It creates or updates a .venv (managed by uv) in the current directory and records the setup in pyproject.toml, leaving the rest of your project untouched.`, +Use this when you want to develop or debug a Databricks project locally: it installs the matching Python version and a compatible databricks-connect, and pins your dependencies to versions known to work with your chosen compute. It creates or updates a .venv (managed by uv) in the current directory and records the setup in pyproject.toml, leaving the rest of your project untouched. + +In a uv workspace member, the command uses the shared .venv at the workspace root and writes the dependency constraints to the root pyproject.toml, because uv reads them only from there. uv sync installs the current member's dependencies into the shared .venv and removes packages that only other members need. All members share one environment, so they must target compatible compute.`, Example: ` # Match a serverless version databricks environments setup-local --serverless-version 5 diff --git a/libs/localenv/merge.go b/libs/localenv/merge.go index 388feed0bbd..ab784564206 100644 --- a/libs/localenv/merge.go +++ b/libs/localenv/merge.go @@ -7,6 +7,8 @@ import ( "regexp" "slices" "strings" + + "github.com/BurntSushi/toml" ) // errMultilineString and errNoProjectTable are the conditions under which @@ -31,6 +33,7 @@ const ( regionRequiresPython = "requires-python" regionDatabricksConnect = "databricks-connect" regionToolUv = "tool.uv.constraint-dependencies" + regionWorkspaceToolUv = regionToolUv + " (uv workspace root)" regionDatabricksEnvironment = "tool.databricks.environment" ) @@ -113,6 +116,11 @@ type MergeOptions struct { // databricks-connect dependency unmanaged: an existing pin is preserved and none // is injected or asserted. SkipDBConnect bool + // ConstraintsAtWorkspaceRoot is set for a uv workspace member: its constraint + // block is written to the workspace root (uv reads constraint-dependencies only + // there), so this file gets none, and a managed block that an earlier run left + // here is removed. requires-python is still managed in the member. + ConstraintsAtWorkspaceRoot bool } // MergeManaged applies the managed transforms to target, preserving every other @@ -121,79 +129,194 @@ type MergeOptions struct { // back in produces identical bytes. opts selects which axes are managed; the // environment region is always reconciled. func MergeManaged(target []byte, c Constraints, opts MergeOptions) (merged []byte, regions []string, err error) { - s := string(target) + merged, err = transformLines(target, func(lines []string) ([]string, error) { + // requires-python needs a [project] table to hold it; without one the merge + // would silently drop the pin, so fail loudly instead. Only enforced when + // constraints are managed — under SkipConstraints there is no requires-python to + // write, so a [project]-less file (e.g. one with only dependency groups) is + // merged for its other axes rather than rejected. + if !opts.SkipConstraints { + if _, _, ok := tableBounds(lines, "[project]"); !ok { + return nil, errNoProjectTable + } + } - // Detect and normalize line endings. We process on "\n" and restore "\r\n" on - // exit. Line endings are treated as a whole-file property: a file that uses - // CRLF anywhere is emitted entirely as CRLF. Real pyproject.toml files use one - // consistent ending, so this is faithful in practice; a file that deliberately - // mixes LF and CRLF would be normalized to CRLF (a benign whitespace-only - // change), which we accept rather than track a terminator per line. - crlf := strings.Contains(s, "\r\n") - if crlf { - s = strings.ReplaceAll(s, "\r\n", "\n") - } + if !opts.SkipConstraints { + var rpChanged bool + lines, rpChanged = mergeRequiresPython(lines, c.RequiresPython) + if rpChanged { + regions = append(regions, regionRequiresPython) + } + } - protected, restore, err := protectMultilineStrings(s) + if !opts.SkipDBConnect { + var dbcChanged bool + lines, _, dbcChanged = mergeDatabricksConnect(lines, c.DatabricksConnect) + // In the install flow, after the managed pin lands in the dev group, remove any + // databricks-connect pin elsewhere that is disjoint from it — the pins that would + // otherwise make uv unsatisfiable. Compatible pins are left alone. An empty pin + // (an artifact without databricks-connect) is a data no-op, distinct from the + // SkipDBConnect intent handled by the gate above. + strayChanged := false + if c.DatabricksConnect != "" { + lines, _, strayChanged = removeStrayDatabricksConnect(lines, c.DatabricksConnect) + } + if dbcChanged || strayChanged { + regions = append(regions, regionDatabricksConnect) + } + } + + lines, envChanged := mergeDatabricksEnvironment(lines, c.EnvironmentVersion) + if envChanged { + regions = append(regions, regionDatabricksEnvironment) + } + + if !opts.SkipConstraints { + var uvChanged bool + if opts.ConstraintsAtWorkspaceRoot { + // A managed block left in the member by an earlier run has no effect + // (uv ignores it), so remove it rather than keep pins that disagree + // with what uv applies. + lines, uvChanged = removeManagedToolUv(lines) + } else { + lines, uvChanged = mergeToolUv(lines, c.ConstraintDeps) + } + if uvChanged { + regions = append(regions, regionToolUv) + } + } + return lines, nil + }) if err != nil { return nil, nil, err } - lines := strings.Split(protected, "\n") + return merged, regions, nil +} - // requires-python needs a [project] table to hold it; without one the merge - // would silently drop the pin, so fail loudly instead. Only enforced when - // constraints are managed — under SkipConstraints there is no requires-python to - // write, so a [project]-less file (e.g. one with only dependency groups) is - // merged for its other axes rather than rejected. - if !opts.SkipConstraints { - if _, _, ok := tableBounds(lines, "[project]"); !ok { - return nil, nil, errNoProjectTable +// MergeToolUvBlock reconciles only the managed [tool.uv] constraint-dependencies +// block of target, preserving every other byte. It is used for the root +// pyproject.toml of a uv workspace, which may be a virtual root with no [project] +// table. changed reports whether the block differs from the input; replaced +// reports whether that change replaces constraint-dependencies the file already +// had (from an earlier run or from the user). +func MergeToolUvBlock(target []byte, deps []string) (merged []byte, changed, replaced bool, err error) { + merged, err = transformLines(target, func(lines []string) ([]string, error) { + if definesToolUvWithoutHeader(lines) { + return nil, errDottedToolUv } + hadPins := hasConstraintDeps(lines) + var out []string + out, changed = mergeToolUv(lines, deps) + replaced = changed && hadPins + return out, nil + }) + if err != nil { + return nil, false, false, err + } + // The line-based merge matches only the plain [tool.uv] spelling. Other valid + // spellings ([ tool.uv ], [tool."uv"], a quoted key) can make it define a table + // or key twice. This file is shared by every workspace member, so never return + // a merge that no longer parses. + if _, err := toml.Decode(string(merged), &map[string]any{}); err != nil { + return nil, false, false, fmt.Errorf("merged workspace root pyproject.toml is not valid TOML; write [tool.uv] and constraint-dependencies without quotes or spaces: %w", err) } + return merged, changed, replaced, nil +} - if !opts.SkipConstraints { - var rpChanged bool - lines, rpChanged = mergeRequiresPython(lines, c.RequiresPython) - if rpChanged { - regions = append(regions, regionRequiresPython) +// hasConstraintDeps reports whether lines hold pins: a managed block with at +// least one entry, or a constraint-dependencies key in the [tool.uv] table. +func hasConstraintDeps(lines []string) bool { + if start, stop, found := markerBounds(lines); found { + return slices.ContainsFunc(lines[start+1:stop], func(l string) bool { + return strings.HasPrefix(strings.TrimSpace(l), `"`) + }) + } + header, end, ok := tableBounds(lines, "[tool.uv]") + if !ok { + return false + } + for i := header + 1; i < end; i++ { + if !constraintDepsRe.MatchString(lines[i]) { + continue } + last, _ := constraintDepsArrayEnd(lines, i, end) + // Drop comments, so a quote in a comment of an empty array is not a pin. + var value strings.Builder + for _, l := range lines[i : last+1] { + if c := commentStart(l); c >= 0 { + l = l[:c] + } + value.WriteString(l) + } + _, v, _ := strings.Cut(value.String(), "=") + return strings.ContainsAny(v, `"'`) } + return false +} - if !opts.SkipDBConnect { - var dbcChanged bool - lines, _, dbcChanged = mergeDatabricksConnect(lines, c.DatabricksConnect) - // In the install flow, after the managed pin lands in the dev group, remove any - // databricks-connect pin elsewhere that is disjoint from it — the pins that would - // otherwise make uv unsatisfiable. Compatible pins are left alone. An empty pin - // (an artifact without databricks-connect) is a data no-op, distinct from the - // SkipDBConnect intent handled by the gate above. - strayChanged := false - if c.DatabricksConnect != "" { - lines, _, strayChanged = removeStrayDatabricksConnect(lines, c.DatabricksConnect) +// errDottedToolUv is returned when the file defines the [tool.uv] table without +// a [tool.uv] header (dotted, quoted, or inline keys). The line-based merge +// cannot extend such a table, and appending a [tool.uv] header would define it +// twice, which uv rejects for the whole workspace. +var errDottedToolUv = errors.New("pyproject.toml defines [tool.uv] without a [tool.uv] header (for example tool.uv.workspace.members or an inline table); use [tool.uv] and [tool.uv.workspace] tables so the constraint block can be added") + +// toolUvTopRe and toolUvInToolRe match a key that defines the [tool.uv] table +// without its header: tool.uv.… or an inline tool = {…} at the top level, and +// uv.… or uv = {…} inside [tool]. Keys may be quoted. +var ( + toolUvTopRe = regexp.MustCompile(`^\s*["']?tool["']?\s*(\.\s*["']?uv["']?\s*[.=]|=\s*\{)`) + toolUvInToolRe = regexp.MustCompile(`^\s*["']?uv["']?\s*[.=]`) +) + +// definesToolUvWithoutHeader reports whether lines define the [tool.uv] table +// only through keys, with no [tool.uv] header. +func definesToolUvWithoutHeader(lines []string) bool { + if _, _, ok := tableBounds(lines, "[tool.uv]"); ok { + return false + } + table := "" + for _, line := range lines { + if name := headerName(line); name != "" { + table = name + continue } - if dbcChanged || strayChanged { - regions = append(regions, regionDatabricksConnect) + if (table == "" && toolUvTopRe.MatchString(line)) || (table == "[tool]" && toolUvInToolRe.MatchString(line)) { + return true } } + return false +} + +// transformLines runs fn over target split into lines, with CRLF normalized to +// LF and multiline strings protected, and restores both on the way out. +func transformLines(target []byte, fn func([]string) ([]string, error)) ([]byte, error) { + s := string(target) - lines, envChanged := mergeDatabricksEnvironment(lines, c.EnvironmentVersion) - if envChanged { - regions = append(regions, regionDatabricksEnvironment) + // Detect and normalize line endings. We process on "\n" and restore "\r\n" on + // exit. Line endings are treated as a whole-file property: a file that uses + // CRLF anywhere is emitted entirely as CRLF. Real pyproject.toml files use one + // consistent ending, so this is faithful in practice; a file that deliberately + // mixes LF and CRLF would be normalized to CRLF (a benign whitespace-only + // change), which we accept rather than track a terminator per line. + crlf := strings.Contains(s, "\r\n") + if crlf { + s = strings.ReplaceAll(s, "\r\n", "\n") } - if !opts.SkipConstraints { - var uvChanged bool - lines, uvChanged = mergeToolUv(lines, c.ConstraintDeps) - if uvChanged { - regions = append(regions, regionToolUv) - } + protected, restore, err := protectMultilineStrings(s) + if err != nil { + return nil, err + } + lines, err := fn(strings.Split(protected, "\n")) + if err != nil { + return nil, err } out := restore(strings.Join(lines, "\n")) if crlf { out = strings.ReplaceAll(out, "\n", "\r\n") } - return []byte(out), regions, nil + return []byte(out), nil } // headerName returns the bracketed table name of a TOML table-header line (e.g. @@ -948,6 +1071,62 @@ func mergeToolUv(lines, deps []string) ([]string, bool) { return appendManagedBlock(lines, renderToolUvBlock(deps, true)), true } +// removeManagedToolUv removes the marker-bracketed managed [tool.uv] block, and +// the one blank line appendManagedBlock puts before a standalone block. A +// header-less block inside a user [tool.uv] table leaves the table's own keys +// and blank lines. +// A standalone block keeps its [tool.uv] header when user keys follow its end +// marker, because those keys belong to that table. +func removeManagedToolUv(lines []string) ([]string, bool) { + start, stop, found := markerBounds(lines) + if !found { + return removeUnmarkedConstraintDeps(lines) + } + var keep []string + if slices.ContainsFunc(lines[start:stop+1], func(l string) bool { return headerName(l) == "[tool.uv]" }) { + if tableBodyHasKeys(lines, stop+1) { + keep = []string{"[tool.uv]"} + } else if start > 0 && strings.TrimSpace(lines[start-1]) == "" { + start-- + } + } + out := make([]string, 0, len(lines)-(stop-start+1)+len(keep)) + out = append(out, lines[:start]...) + out = append(out, keep...) + out = append(out, lines[stop+1:]...) + return out, true +} + +// removeUnmarkedConstraintDeps removes constraint-dependencies from a plain +// [tool.uv] table, the form a run before the managed markers wrote. Like +// mergeToolUv, it drops the whole table when that key is all it holds. +func removeUnmarkedConstraintDeps(lines []string) ([]string, bool) { + header, end, ok := tableBounds(lines, "[tool.uv]") + if !ok || !slices.ContainsFunc(lines[header+1:end], constraintDepsRe.MatchString) { + return lines, false + } + if toolUvHasOnlyConstraintDeps(lines, header, end) { + out := make([]string, 0, len(lines)-(end-header)) + out = append(out, lines[:header]...) + return append(out, lines[end:]...), true + } + return removeConstraintDeps(lines, header, end), true +} + +// tableBodyHasKeys reports whether the lines from index from up to the next +// table header hold anything other than blank lines and comments. +func tableBodyHasKeys(lines []string, from int) bool { + for _, line := range lines[from:] { + if tableHeaderRe.MatchString(line) { + return false + } + if t := strings.TrimSpace(line); t != "" && !strings.HasPrefix(t, "#") { + return true + } + } + return false +} + // markerAttachedToToolUv reports whether the managed marker region beginning at // index start sits inside an existing [tool.uv] table — i.e. the nearest table // header above it is [tool.uv]. When true, the managed block must omit its own diff --git a/libs/localenv/merge_test.go b/libs/localenv/merge_test.go index 3e175ae0e2a..2bf2b741b32 100644 --- a/libs/localenv/merge_test.go +++ b/libs/localenv/merge_test.go @@ -1271,3 +1271,158 @@ func TestRenderFreshPyprojectSkipDBConnectOmitsDBConnect(t *testing.T) { assert.Contains(t, s, "constraint-dependencies") requireValidTOML(t, out) } + +func TestMergeManagedConstraintsAtWorkspaceRootAddsNoBlock(t *testing.T) { + in := []byte("[project]\nname = \"a\"\nrequires-python = \">=3.10\"\n") + c := Constraints{RequiresPython: ">=3.12", ConstraintDeps: []string{"numpy==1.26.4"}} + out, regions, err := MergeManaged(in, c, MergeOptions{ConstraintsAtWorkspaceRoot: true}) + require.NoError(t, err) + assert.Contains(t, string(out), `requires-python = ">=3.12"`) + assert.NotContains(t, string(out), "constraint-dependencies") + assert.NotContains(t, regions, regionToolUv) +} + +func TestMergeToolUvBlockVirtualRoot(t *testing.T) { + in := []byte("[tool.uv.workspace]\nmembers = [\"bundles/*\"]\n") + out, changed, _, err := MergeToolUvBlock(in, []string{"numpy==1.26.4"}) + require.NoError(t, err) + assert.True(t, changed) + assert.Contains(t, string(out), managedMarkerStart) + assert.Contains(t, string(out), `"numpy==1.26.4"`) + assert.Contains(t, string(out), "[tool.uv.workspace]") + + again, changed2, _, err := MergeToolUvBlock(out, []string{"numpy==1.26.4"}) + require.NoError(t, err) + assert.False(t, changed2) + assert.Equal(t, out, again) +} + +func TestMergeToolUvBlockKeepsCRLF(t *testing.T) { + in := []byte("[tool.uv.workspace]\r\nmembers = [\"a\"]\r\n") + out, _, _, err := MergeToolUvBlock(in, []string{"x==1"}) + require.NoError(t, err) + assert.NotContains(t, strings.ReplaceAll(string(out), "\r\n", ""), "\n") +} + +func TestMergeManagedConstraintsAtWorkspaceRootRemovesStaleStandaloneBlock(t *testing.T) { + // A run before workspace support wrote the managed block into the member, + // where uv ignores it. The member merge removes it; the root holds the pins. + in := "[project]\nname = \"a\"\nrequires-python = \">=3.12\"\n\n" + + managedMarkerStart + "\n[tool.uv]\nconstraint-dependencies = [\n \"x==1\",\n]\n" + managedMarkerEnd + "\n" + c := Constraints{RequiresPython: ">=3.12"} + out, regions, err := MergeManaged([]byte(in), c, MergeOptions{ConstraintsAtWorkspaceRoot: true, SkipDBConnect: true}) + require.NoError(t, err) + assert.Equal(t, "[project]\nname = \"a\"\nrequires-python = \">=3.12\"\n", string(out)) + assert.Contains(t, regions, regionToolUv) + + again, regions2, err := MergeManaged(out, c, MergeOptions{ConstraintsAtWorkspaceRoot: true, SkipDBConnect: true}) + require.NoError(t, err) + assert.Equal(t, out, again) + assert.NotContains(t, regions2, regionToolUv) +} + +func TestMergeManagedConstraintsAtWorkspaceRootRemovesStaleBlockKeepsUserTable(t *testing.T) { + in := "[project]\nname = \"a\"\nrequires-python = \">=3.12\"\n\n[tool.uv]\npackage = false\n" + + managedMarkerStart + "\nconstraint-dependencies = [\n \"x==1\",\n]\n" + managedMarkerEnd + "\n" + c := Constraints{RequiresPython: ">=3.12"} + out, _, err := MergeManaged([]byte(in), c, MergeOptions{ConstraintsAtWorkspaceRoot: true, SkipDBConnect: true}) + require.NoError(t, err) + assert.Equal(t, "[project]\nname = \"a\"\nrequires-python = \">=3.12\"\n\n[tool.uv]\npackage = false\n", string(out)) +} + +func TestMergeManagedConstraintsAtWorkspaceRootKeepsHeaderForTrailingUserKeys(t *testing.T) { + // The user added a key below an earlier standalone block. It belongs to that + // block's [tool.uv] table, so the header must survive the block's removal. + in := "[project]\nname = \"a\"\nrequires-python = \">=3.12\"\n\n" + + managedMarkerStart + "\n[tool.uv]\nconstraint-dependencies = [\n \"x==1\",\n]\n" + managedMarkerEnd + "\npackage = false\n" + c := Constraints{RequiresPython: ">=3.12"} + out, _, err := MergeManaged([]byte(in), c, MergeOptions{ConstraintsAtWorkspaceRoot: true, SkipDBConnect: true}) + require.NoError(t, err) + assert.Equal(t, "[project]\nname = \"a\"\nrequires-python = \">=3.12\"\n\n[tool.uv]\npackage = false\n", string(out)) +} + +func TestMergeToolUvBlockDottedKeyGuard(t *testing.T) { + // A [tool.uv] table defined without its own header cannot be extended by the + // line-based merge; appending a [tool.uv] header would define it twice, which + // uv rejects for the whole workspace. + tests := []struct { + name string + in string + wantErr bool + }{ + {"workspace subtable header", "[tool.uv.workspace]\nmembers = [\"bundles/*\"]\n", false}, + {"dotted key at top level", "tool.uv.workspace.members = [\"bundles/*\"]\n", true}, + {"quoted dotted key", "\"tool\".\"uv\".workspace.members = [\"bundles/*\"]\n", true}, + {"dotted key in [tool]", "[tool]\nuv.workspace.members = [\"bundles/*\"]\n", true}, + {"inline tool table", "tool = { uv = { workspace = { members = [\"bundles/*\"] } } }\n", true}, + {"inline uv table in [tool]", "[tool]\nuv = { workspace = { members = [\"bundles/*\"] } }\n", true}, + // Header spellings the line-based merge does not match; the merged file + // would define [tool.uv] twice. + {"header with spaces", "[ tool.uv ]\ndev-dependencies = []\n", true}, + {"quoted header", "[tool.\"uv\"]\ndev-dependencies = []\n", true}, + {"quoted constraint key", "[tool.uv]\n\"constraint-dependencies\" = [\"y==1\"]\n", true}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + _, _, _, err := MergeToolUvBlock([]byte(tc.in), []string{"x==1"}) + if tc.wantErr { + assert.Error(t, err) + } else { + assert.NoError(t, err) + } + }) + } +} + +func TestMergeToolUvBlockReportsReplacedPins(t *testing.T) { + tests := []struct { + name string + in string + wantReplaced bool + }{ + {"no pins", "[tool.uv.workspace]\nmembers = [\"a\"]\n", false}, + {"only a comment names the key", "# constraint-dependencies are added by setup-local\n[tool.uv.workspace]\nmembers = [\"a\"]\n", false}, + {"user pins", "[tool.uv]\nconstraint-dependencies = [\"urllib3<2\"]\n", true}, + {"managed pins", managedMarkerStart + "\n[tool.uv]\nconstraint-dependencies = [\n \"old==1\",\n]\n" + managedMarkerEnd + "\n", true}, + {"empty unmarked key", "[tool.uv]\nconstraint-dependencies = []\n", false}, + {"quote only in a comment", "[tool.uv]\nconstraint-dependencies = [\n # \"placeholder\"\n] # 'none yet'\n", false}, + {"pin with a trailing comment", "[tool.uv]\nconstraint-dependencies = [\n \"urllib3<2\", # why\n]\n", true}, + {"empty managed block", managedMarkerStart + "\n[tool.uv]\nconstraint-dependencies = [\n]\n" + managedMarkerEnd + "\n", false}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + _, _, replaced, err := MergeToolUvBlock([]byte(tc.in), []string{"x==1"}) + require.NoError(t, err) + assert.Equal(t, tc.wantReplaced, replaced) + }) + } +} + +func TestMergeManagedConstraintsAtWorkspaceRootRemovesUnmarkedPins(t *testing.T) { + // A run before the managed markers wrote a plain [tool.uv] table with only + // constraint-dependencies; the standalone merge treats it as its own, so the + // member merge removes it too. Other [tool.uv] keys stay. + c := Constraints{RequiresPython: ">=3.12"} + opts := MergeOptions{ConstraintsAtWorkspaceRoot: true, SkipDBConnect: true} + head := "[project]\nname = \"a\"\nrequires-python = \">=3.12\"\n" + + out, regions, err := MergeManaged([]byte(head+"\n[tool.uv]\nconstraint-dependencies = [\"x==1\"]\n"), c, opts) + require.NoError(t, err) + assert.Equal(t, head, string(out)) + assert.Contains(t, regions, regionToolUv) + + out, _, err = MergeManaged([]byte(head+"\n[tool.uv]\npackage = false\nconstraint-dependencies = [\"x==1\"]\n"), c, opts) + require.NoError(t, err) + assert.Equal(t, head+"\n[tool.uv]\npackage = false\n", string(out)) +} + +func TestMergeManagedConstraintsAtWorkspaceRootKeepsUserBlankLine(t *testing.T) { + // A header-less managed block sits inside the user's [tool.uv] table; the + // blank line before it is the user's, not the merge's. + in := "[project]\nname = \"a\"\nrequires-python = \">=3.12\"\n\n[tool.uv]\npackage = false\n\n" + + managedMarkerStart + "\nconstraint-dependencies = [\n \"x==1\",\n]\n" + managedMarkerEnd + "\n" + c := Constraints{RequiresPython: ">=3.12"} + out, _, err := MergeManaged([]byte(in), c, MergeOptions{ConstraintsAtWorkspaceRoot: true, SkipDBConnect: true}) + require.NoError(t, err) + assert.Equal(t, "[project]\nname = \"a\"\nrequires-python = \">=3.12\"\n\n[tool.uv]\npackage = false\n\n", string(out)) +} diff --git a/libs/localenv/pipeline.go b/libs/localenv/pipeline.go index 81c55837e2d..aa64acacb07 100644 --- a/libs/localenv/pipeline.go +++ b/libs/localenv/pipeline.go @@ -81,6 +81,10 @@ type Pipeline struct { // Result.VenvPath. Set in preflight. venvRel string + // workspaceRoot is the uv workspace root when ProjectDir is a workspace + // member, else "". Set in preflight. + workspaceRoot string + // nowFn supplies the time for backup filenames; nil means time.Now. Injected // in tests so backup names are deterministic. nowFn func() time.Time @@ -181,8 +185,6 @@ func (p *Pipeline) run(ctx context.Context) error { if m := detectManager(p.ProjectDir); m != managerUv { return p.fail(PhasePreflight, false, NewError(ErrManagerUnsupported, nil, "%s", managerGuidance(m))) } - p.venvRel = venvDir - p.venv = filepath.Join(p.ProjectDir, p.venvRel) // Under --dry-run the pipeline only reads and reports a plan, so it must not // mutate anything at preflight. Two preflight steps can write: // - ensureWritable creates and removes a temp file (and would fail a @@ -190,9 +192,8 @@ func (p *Pipeline) run(ctx context.Context) error { // - PackageManager.EnsureAvailable may install the manager (uv) if missing. // Both exist to fail fast before real writes, which --dry-run never performs, so // they are skipped in a dry run. Neither result is needed to compute the plan. - if p.Check { - p.markOK(PhasePreflight, "check") - } else { + detail := "check" + if !p.Check { if err := ensureWritable(p.ProjectDir); err != nil { return p.fail(PhasePreflight, false, NewError(ErrNotWritable, err, "project directory %s is not writable", filepath.ToSlash(p.ProjectDir))) } @@ -200,8 +201,35 @@ func (p *Pipeline) run(ctx context.Context) error { if err != nil { return p.fail(PhasePreflight, false, asPipelineError(err, ErrUvMissing, "%s unavailable", p.PM.Name())) } - p.markOK(PhasePreflight, p.PM.Name()+" "+version) + detail = p.PM.Name() + " " + version + } + // A uv workspace member shares the venv at the workspace root, where uv sync + // creates it, and keeps its constraint block in the root pyproject.toml, the + // only file uv reads constraint-dependencies from. + // https://docs.astral.sh/uv/concepts/projects/workspaces/ + // Every run asks uv, also for a standalone project: only uv applies its member + // globs and exclude rules, and one call costs little next to uv sync. + var rootRel string + p.workspaceRoot, rootRel = p.discoverWorkspaceRoot(ctx) + // A discovery failure keeps the standalone behavior, but an interrupt must + // not: the run would go on to merge as a standalone project and write files. + if err := ctx.Err(); err != nil { + return p.fail(PhasePreflight, false, NewError(ErrCanceled, err, "interrupted")) + } + venvBase := p.ProjectDir + p.venvRel = venvDir + if p.workspaceRoot != "" { + venvBase = p.workspaceRoot + p.venvRel = filepath.Join(rootRel, venvDir) + p.res.UvWorkspaceRoot = filepath.ToSlash(p.workspaceRoot) + if !p.Check { + if err := ensureWritable(p.workspaceRoot); err != nil { + return p.fail(PhasePreflight, false, NewError(ErrNotWritable, err, "uv workspace root %s is not writable", filepath.ToSlash(p.workspaceRoot))) + } + } } + p.venv = filepath.Join(venvBase, venvDir) + p.markOK(PhasePreflight, detail) // Phase: resolve — compute target → environment key. p.report(ctx, PhaseResolve) @@ -228,9 +256,13 @@ func (p *Pipeline) run(ctx context.Context) error { // values, so c keeps the real artifact values throughout. --constraints-only maps // to --no-dbconnect: it stops managing the pin (not reported here, not asserted in // validate) but keeps the resolved value on c. + // A workspace member's constraint block goes to the workspace root (see + // planRootMerge). Discovery never reports a greenfield project as a member + // (see discoverWorkspaceRoot), so RenderFreshPyproject needs no such option. opts := MergeOptions{ - SkipConstraints: p.SkipConstraints, - SkipDBConnect: p.Mode == ModeConstraintsOnly, + SkipConstraints: p.SkipConstraints, + SkipDBConnect: p.Mode == ModeConstraintsOnly, + ConstraintsAtWorkspaceRoot: p.workspaceRoot != "", } dbcVersion := dbcVersionFromPin(c.DatabricksConnect) @@ -247,7 +279,7 @@ func (p *Pipeline) run(ctx context.Context) error { // The serverless environment version (empty for cluster targets) is written // into [tool.databricks.environment] so the project also runs in serverless Jobs. p.report(ctx, PhaseMerge) - mergedBytes, greenfield, err := p.mergePlan(ctx, pyMinor, c, compute.ServerlessEnvironmentVersion(), opts) + mergedBytes, root, greenfield, err := p.mergePlan(ctx, pyMinor, c, compute.ServerlessEnvironmentVersion(), opts) if err != nil { return err } @@ -262,7 +294,7 @@ func (p *Pipeline) run(ctx context.Context) error { } // Apply the merged content to disk (backup first for existing projects). - if err := p.applyMerge(ctx, mergedBytes, greenfield); err != nil { + if err := p.applyMerge(ctx, mergedBytes, root, greenfield); err != nil { return err } p.markOK(PhaseMerge, "") @@ -385,7 +417,7 @@ func (p *Pipeline) plannedBackupName(dir string) (string, error) { // --dry-run). opts selects which managed axes are written. envVersion is the // serverless environment version to write into [tool.databricks.environment], or // "" for a cluster target. -func (p *Pipeline) mergePlan(_ context.Context, pyMinor string, c *Constraints, envVersion string, opts MergeOptions) (merged []byte, greenfield bool, err error) { +func (p *Pipeline) mergePlan(_ context.Context, pyMinor string, c *Constraints, envVersion string, opts MergeOptions) (merged []byte, root *rootMerge, greenfield bool, err error) { pyproject := p.pyprojectPath() // The merge base is the live pyproject.toml, not the backup. MergeManaged @@ -406,7 +438,7 @@ func (p *Pipeline) mergePlan(_ context.Context, pyMinor string, c *Constraints, // after detectManager saw it) must not be misread as greenfield — that // would render a fresh file and overwrite the user's project with no // backup. Fail instead of destroying unrecoverable state (invariant 2). - return nil, false, p.fail(PhaseMerge, false, NewError(ErrMerge, rerr, "read pyproject.toml %s failed", filepath.ToSlash(pyproject))) + return nil, nil, false, p.fail(PhaseMerge, false, NewError(ErrMerge, rerr, "read pyproject.toml %s failed", filepath.ToSlash(pyproject))) } greenfield = baseBytes == nil @@ -440,7 +472,7 @@ func (p *Pipeline) mergePlan(_ context.Context, pyMinor string, c *Constraints, } else { merged, changedRegions, err = MergeManaged(baseBytes, effective, opts) if err != nil { - return nil, greenfield, p.fail(PhaseMerge, false, NewError(ErrMerge, err, "merge managed regions failed")) + return nil, nil, greenfield, p.fail(PhaseMerge, false, NewError(ErrMerge, err, "merge managed regions failed")) } // Surface merge-quality warnings (overridden, consolidated, or duplicated pins, // conflicting user constraints) from the pre-merge file. Greenfield has nothing of @@ -452,12 +484,24 @@ func (p *Pipeline) mergePlan(_ context.Context, pyMinor string, c *Constraints, detectMergeWarnings(baseBytes, effective, planDBConnect(baseBytes, effective, opts), opts)...) } + root, err = p.planRootMerge(c.ConstraintDeps, opts) + if err != nil { + return nil, nil, greenfield, err + } + if root != nil && root.changed { + changedRegions = append(changedRegions, regionWorkspaceToolUv) + } + // Under --dry-run, build the plan (with a diff) for reporting. A real run does // not need the diff. if p.Check { + diff := pyprojectDiff(pyprojectFile, baseBytes, merged, greenfield) + if root != nil { + diff += pyprojectDiff("uv workspace root "+pyprojectFile, root.base, root.merged, false) + } plan := &Plan{ WouldWrite: filepath.ToSlash(pyproject), - Diff: pyprojectDiff(pyprojectFile, baseBytes, merged, greenfield), + Diff: diff, ChangedRegions: changedRegions, WouldInstallPython: pyMinor, } @@ -468,9 +512,52 @@ func (p *Pipeline) mergePlan(_ context.Context, pyMinor string, c *Constraints, plan.WouldBackup = filepath.ToSlash(backupName) } } + // Unlike the member file, the root is reported only when the run changes + // it; Result.UvWorkspaceRoot already names the root. + if root != nil && root.changed { + plan.WouldWriteUvWorkspace = filepath.ToSlash(root.path) + if backupName, statErr := p.plannedBackupName(p.workspaceRoot); statErr == nil { + plan.WouldBackupUvWorkspace = filepath.ToSlash(backupName) + } + } p.res.Plan = plan } - return merged, greenfield, nil + return merged, root, greenfield, nil +} + +// rootMerge is the planned change to the workspace root pyproject.toml. +type rootMerge struct { + path string + base []byte + merged []byte + changed bool +} + +// planRootMerge merges the constraint block into the workspace root +// pyproject.toml. It returns nil when there is no workspace root or when +// constraints are not managed (--no-constraints). +func (p *Pipeline) planRootMerge(deps []string, opts MergeOptions) (*rootMerge, error) { + if p.workspaceRoot == "" || opts.SkipConstraints { + return nil, nil + } + path := filepath.Join(p.workspaceRoot, pyprojectFile) + base, err := os.ReadFile(path) + if err != nil { + return nil, p.fail(PhaseMerge, false, NewError(ErrMerge, err, "read workspace root pyproject.toml %s failed", filepath.ToSlash(path))) + } + merged, changed, replaced, err := MergeToolUvBlock(base, deps) + if err != nil { + return nil, p.fail(PhaseMerge, false, NewError(ErrMerge, err, "merge workspace root constraints failed")) + } + // Earlier pins are replaced, whether an earlier run (maybe for another + // bundle) wrote them or the user did, and the new ones apply to every member. + if replaced { + p.res.Warnings = append(p.res.Warnings, Warning{ + Code: WarnUvWorkspaceConstraintsReplaced, + Message: fmt.Sprintf("the constraint pins in %s are replaced; all members of this uv workspace share them", filepath.ToSlash(path)), + }) + } + return &rootMerge{path: path, base: base, merged: merged, changed: changed}, nil } // pyprojectDiff renders the unified diff the --dry-run plan reports for one @@ -492,7 +579,7 @@ func pyprojectDiff(label string, base, merged []byte, greenfield bool) string { // applyMerge writes the merged bytes to disk, backing up the current // pyproject.toml first. From the backup copy onward, disk has been mutated. -func (p *Pipeline) applyMerge(_ context.Context, mergedBytes []byte, greenfield bool) error { +func (p *Pipeline) applyMerge(_ context.Context, mergedBytes []byte, root *rootMerge, greenfield bool) error { backup, mutated, perr := p.writePyproject(p.pyprojectPath(), mergedBytes, greenfield) // A backup written before a failed write is still reported, so the consumer // can restore from it. @@ -502,6 +589,18 @@ func (p *Pipeline) applyMerge(_ context.Context, mergedBytes []byte, greenfield if perr != nil { return p.fail(PhaseMerge, mutated, perr) } + if root != nil && root.changed { + // The member is written first, so a root failure reports disk as mutated + // when either write changed it. + rootBackup, rootMutated, perr := p.writePyproject(root.path, root.merged, false) + if rootBackup != "" { + p.res.UvWorkspaceBackupPath = filepath.ToSlash(rootBackup) + } + if perr != nil { + perr.Msg = "uv workspace root: " + perr.Msg + return p.fail(PhaseMerge, mutated || rootMutated, perr) + } + } return nil } @@ -655,7 +754,8 @@ func (p *Pipeline) validate(ctx context.Context, expectedPyMinor, dbcPin string, // venvPath is reported relative to the project root (spec §6.1), not as an // absolute path: the VS Code consumer already knows the project root (it sets // the working directory when it shells out). For a standalone project this is - // ".venv". + // ".venv"; for a workspace member it is computed between the symlink-resolved + // folders (see discoverWorkspaceRoot), for example "../../.venv". p.res.VenvPath = filepath.ToSlash(p.venvRel) if p.res.Resolved != nil { if managed { diff --git a/libs/localenv/pipeline_test.go b/libs/localenv/pipeline_test.go index fda5f63e0ea..e1c2e735e28 100644 --- a/libs/localenv/pipeline_test.go +++ b/libs/localenv/pipeline_test.go @@ -27,13 +27,30 @@ const cancelPMStderr = "error: no solution found: databricks-connect==17.2 confl // below — a plain ">= 0" assertion would also hold for the unset field. const fetchDelay = 25 * time.Millisecond -type fakePM struct{ py, dbc, pyspark, dbcImportErr string } +type fakePM struct { + py, dbc, pyspark, dbcImportErr string + // root and rootErr are what WorkspaceRoot reports; an empty root means the + // project is not a workspace member. + root string + rootErr error +} func (fakePM) Name() string { return "fake" } func (fakePM) EnsureAvailable(context.Context) (string, error) { return "fake 1.0", nil } func (fakePM) EnsurePython(_ context.Context, minor string) (PythonSelection, error) { return PythonSelection{Executable: minor, Resolution: PythonResolutionUVInstallSucceeded}, nil } + +func (f fakePM) WorkspaceRoot(_ context.Context, projectDir string) (string, error) { + if f.rootErr != nil { + return "", f.rootErr + } + if f.root != "" { + return f.root, nil + } + return projectDir, nil +} + func (fakePM) Provision(context.Context, string, string) error { return nil } func (fakePM) PostProvision(context.Context, string, string) error { return nil } func (f fakePM) Validate(context.Context, string, string) (VenvInfo, error) { @@ -54,6 +71,10 @@ func (noProvisionPM) EnsurePython(context.Context, string) (PythonSelection, err return PythonSelection{}, errors.New("EnsurePython must not be called under --dry-run") } +func (noProvisionPM) WorkspaceRoot(_ context.Context, projectDir string) (string, error) { + return projectDir, nil +} + func (noProvisionPM) Provision(context.Context, string, string) error { return errors.New("Provision must not be called under --dry-run") } @@ -840,7 +861,7 @@ func TestApplyMergeFailsOnUnreadableDirWithoutOverwritingBackup(t *testing.T) { t.Cleanup(func() { _ = os.Chmod(dir, 0o755) }) p := &Pipeline{ProjectDir: dir, res: &Result{Phases: initialPhases()}} - err := p.applyMerge(t.Context(), []byte("merged"), false) + err := p.applyMerge(t.Context(), []byte("merged"), nil, false) require.Error(t, err) var pe *PipelineError require.ErrorAs(t, err, &pe) @@ -1068,13 +1089,13 @@ func TestApplyMergeSameInstantBackupsGetUniqueNames(t *testing.T) { } // First modifying apply: backs up "current-1" under a timestamped name. - require.NoError(t, p.applyMerge(t.Context(), []byte("merged-1\n"), false)) + require.NoError(t, p.applyMerge(t.Context(), []byte("merged-1\n"), nil, false)) first := p.res.BackupPath require.NotEmpty(t, first) // Second modifying apply at the SAME instant: backs up "merged-1" but must not // collide with or overwrite the first timestamped backup. - require.NoError(t, p.applyMerge(t.Context(), []byte("merged-2\n"), false)) + require.NoError(t, p.applyMerge(t.Context(), []byte("merged-2\n"), nil, false)) second := p.res.BackupPath require.NotEmpty(t, second) diff --git a/libs/localenv/pkgmanager.go b/libs/localenv/pkgmanager.go index dba3b084fa3..18c3767a2e4 100644 --- a/libs/localenv/pkgmanager.go +++ b/libs/localenv/pkgmanager.go @@ -16,6 +16,11 @@ type PackageManager interface { // minor. The returned executable is passed unchanged to Provision. EnsurePython(ctx context.Context, minor string) (PythonSelection, error) + // WorkspaceRoot returns the absolute root of the workspace that the package + // manager uses for projectDir. For a project that is not a workspace member, + // this is projectDir itself. + WorkspaceRoot(ctx context.Context, projectDir string) (string, error) + // Provision installs the project dependencies inside projectDir, pinning the // environment to python, which is either a minor request (e.g. "3.12") or an // exact interpreter path selected by EnsurePython. diff --git a/libs/localenv/provision_integration_test.go b/libs/localenv/provision_integration_test.go index 71420c20a49..c3e648bffc0 100644 --- a/libs/localenv/provision_integration_test.go +++ b/libs/localenv/provision_integration_test.go @@ -6,6 +6,8 @@ import ( "net/http/httptest" "os" "path/filepath" + "regexp" + "strconv" "testing" "github.com/databricks/cli/libs/env" @@ -105,6 +107,53 @@ func TestProvisionCreatesAndValidatesVenv(t *testing.T) { "expected databricks-connect major 17, got %q", res.Resolved.DBConnectVersion) } +// TestProvisionWorkspaceMember runs the real pipeline from a member of a uv +// workspace with two members. uv sync creates the shared venv at the workspace +// root, the constraint block lands in the root pyproject.toml, and uv applies it +// (the lock honors the artifact's pyarrow<19 pin, a databricks-connect +// dependency). +func TestProvisionWorkspaceMember(t *testing.T) { + requireRealProvision(t) + if _, err := process.Background(t.Context(), []string{mustUv(t, t.Context()), "workspace", "dir", "--help"}); err != nil { + t.Skipf("uv has no `workspace dir` command (added in uv 0.9.9): %v", err) + } + root := t.TempDir() + for _, m := range []string{"a", "b"} { + dir := filepath.Join(root, "bundles", m) + require.NoError(t, os.MkdirAll(dir, 0o755)) + require.NoError(t, os.WriteFile(filepath.Join(dir, pyprojectFile), + []byte("[project]\nname = \""+m+"\"\nversion = \"0.1.0\"\nrequires-python = \">=3.10\"\n"), 0o644)) + } + require.NoError(t, os.WriteFile(filepath.Join(root, pyprojectFile), + []byte("[project]\nname = \"root\"\nversion = \"0.1.0\"\nrequires-python = \">=3.10\"\n\n[tool.uv.workspace]\nmembers = [\"bundles/*\"]\n"), 0o644)) + member := filepath.Join(root, "bundles", "a") + srv := serverlessConstraintServer(t) + + p := &Pipeline{ + Mode: ModeDefault, Check: false, ProjectDir: member, + ConstraintBaseURL: srv.URL, CacheDir: t.TempDir(), + Flags: ComputeFlags{Serverless: "5"}, + Compute: stubCompute{}, PM: NewUvManager(), + } + res, err := p.Run(realProvisionCtx(t)) + require.NoError(t, err) + assert.Equal(t, "../../.venv", res.VenvPath) + assert.FileExists(t, venvPython(filepath.Join(root, venvDir))) + assert.NoDirExists(t, filepath.Join(member, venvDir)) + + rootPyproject, err := os.ReadFile(filepath.Join(root, pyprojectFile)) + require.NoError(t, err) + assert.Contains(t, string(rootPyproject), `"pyarrow<19"`) + + lock, err := os.ReadFile(filepath.Join(root, "uv.lock")) + require.NoError(t, err) + m := regexp.MustCompile(`name = "pyarrow"\nversion = "(\d+)\.`).FindSubmatch(lock) + require.NotNil(t, m, "uv.lock has no pyarrow entry") + major, err := strconv.Atoi(string(m[1])) + require.NoError(t, err) + assert.Less(t, major, 19, "the root constraint pyarrow<19 must apply") +} + // TestProvisionValidateIgnoresActiveVirtualEnv is the regression guard for the // validate fix: an active VIRTUAL_ENV pointing at a *real* interpreter of a // different version must not make validation read that interpreter instead of diff --git a/libs/localenv/result.go b/libs/localenv/result.go index ea256661c9f..2ceb7dc6bbc 100644 --- a/libs/localenv/result.go +++ b/libs/localenv/result.go @@ -213,10 +213,14 @@ type ResolvedInfo struct { // Plan describes the changes a --dry-run run would apply (spec §6.3). // ChangedRegions is retained for text output only and is not serialized. type Plan struct { - WouldWrite string `json:"wouldWrite"` - WouldBackup string `json:"wouldBackup,omitempty"` - WouldInstallPython string `json:"wouldInstallPython,omitempty"` - Diff string `json:"diff"` + WouldWrite string `json:"wouldWrite"` + WouldBackup string `json:"wouldBackup,omitempty"` + // WouldWriteUvWorkspace and WouldBackupUvWorkspace name the workspace root + // pyproject.toml and its planned backup, for a uv workspace member only. + WouldWriteUvWorkspace string `json:"wouldWriteUvWorkspace,omitempty"` + WouldBackupUvWorkspace string `json:"wouldBackupUvWorkspace,omitempty"` + WouldInstallPython string `json:"wouldInstallPython,omitempty"` + Diff string `json:"diff"` ChangedRegions []string `json:"-"` } @@ -293,6 +297,11 @@ const ( // the user's own pyproject pins it (constraints-only mode) — so it agrees with the // validate hard-fail, which keys on the installed venv rather than the mode. WarnStandalonePysparkConflict = "W_STANDALONE_PYSPARK_CONFLICT" + // WarnUvWorkspaceConstraintsReplaced: the project is a uv workspace member and + // the run replaced constraint pins the workspace root already had, from an + // earlier run (for example for another bundle) or written by the user. Every + // member of the workspace now resolves against the new pins. + WarnUvWorkspaceConstraintsReplaced = "W_UV_WORKSPACE_CONSTRAINTS_REPLACED" ) // Result is the full outcome of a sync run and the root of the --json object @@ -303,16 +312,20 @@ const ( // a distinction that trips JSON consumers and golden diffs. Construct a Result // with NewResult (or otherwise seed both) rather than a bare Result{} literal. type Result struct { - SchemaVersion int `json:"schemaVersion"` - Command string `json:"command"` - OK bool `json:"ok"` - Mode string `json:"mode"` - DryRun bool `json:"dryRun"` - Compute *ComputeInfo `json:"compute,omitempty"` - Resolved *ResolvedInfo `json:"resolved,omitempty"` - Greenfield bool `json:"greenfield"` - Plan *Plan `json:"plan,omitempty"` - VenvPath string `json:"venvPath,omitempty"` + SchemaVersion int `json:"schemaVersion"` + Command string `json:"command"` + OK bool `json:"ok"` + Mode string `json:"mode"` + DryRun bool `json:"dryRun"` + Compute *ComputeInfo `json:"compute,omitempty"` + Resolved *ResolvedInfo `json:"resolved,omitempty"` + Greenfield bool `json:"greenfield"` + Plan *Plan `json:"plan,omitempty"` + VenvPath string `json:"venvPath,omitempty"` + // UvWorkspaceRoot is the uv workspace root (absolute, slash form) when the + // project is a workspace member. The [tool.uv] constraint block and the venv + // live there. Empty for a standalone project. + UvWorkspaceRoot string `json:"uvWorkspaceRoot,omitempty"` PythonResolution PythonResolution `json:"pythonResolution,omitempty"` // PythonInterpreter is the exact interpreter chosen by the installed-Python // fallback, named in the text summary so the user can tell which interpreter @@ -328,6 +341,9 @@ type Result struct { Warnings []Warning `json:"warnings"` Error *PipelineError `json:"error"` BackupPath string `json:"backupPath,omitempty"` + // UvWorkspaceBackupPath is the backup of the workspace root pyproject.toml, set + // only when the run changed that file. + UvWorkspaceBackupPath string `json:"uvWorkspaceBackupPath,omitempty"` // DurationMs is the pipeline's wall time in milliseconds (spec §6). It covers the // CLI pipeline only; the extension measures its own end-to-end latency (process // spawn, interpreter adoption) separately. diff --git a/libs/localenv/uv.go b/libs/localenv/uv.go index c62b84429ac..40ef1deff3e 100644 --- a/libs/localenv/uv.go +++ b/libs/localenv/uv.go @@ -300,6 +300,31 @@ func lineWithPrefix(out, prefix string) (string, bool) { return "", false } +// WorkspaceRoot asks uv for the workspace root of projectDir (`uv workspace +// dir`, added in uv 0.9.9). Under --dry-run EnsureAvailable has not run, so the +// binary is discovered here without installing it. +// https://docs.astral.sh/uv/reference/cli/#uv-workspace-dir +func (m *uvManager) WorkspaceRoot(ctx context.Context, projectDir string) (string, error) { + bin := m.bin + if bin == "" { + found, err := discoverUv(ctx) + if err != nil { + return "", err + } + bin = found + } + out, err := m.runUvOutput(ctx, append([]string{bin}, m.workspaceDirArgs()...), projectDir) + if err != nil { + return "", err + } + return strings.TrimSpace(out), nil +} + +// workspaceDirArgs returns the argument slice for `uv workspace dir`. +func (m *uvManager) workspaceDirArgs() []string { + return []string{"workspace", "dir"} +} + // syncArgs returns the argument slice for `uv sync` (without the binary), // pinning the interpreter to python via --python. func (m *uvManager) syncArgs(python string) []string { diff --git a/libs/localenv/uv_test.go b/libs/localenv/uv_test.go index 3add58cb4ac..26f6786b421 100644 --- a/libs/localenv/uv_test.go +++ b/libs/localenv/uv_test.go @@ -56,6 +56,34 @@ func TestUvArgs(t *testing.T) { "python", "find", "--system", "--no-python-downloads", "cpython@==3.12.*", }, m.pythonFindArgs("3.12")) assert.Equal(t, []string{"pip", "install", "pip", "--python", "/p/.venv/bin/python"}, m.pipSeedArgs("/p/.venv/bin/python")) + assert.Equal(t, []string{"workspace", "dir"}, m.workspaceDirArgs()) +} + +func TestWorkspaceRootRunsInProjectDirAndTrimsOutput(t *testing.T) { + project := filepath.Join("p", "proj") + ctx, stub := process.WithStub(t.Context()) + var dir string + stub.WithCallback(func(cmd *exec.Cmd) error { + dir = cmd.Dir + _, err := cmd.Stdout.Write([]byte("/repo\n")) + return err + }) + m := &uvManager{bin: "uv"} + + root, err := m.WorkspaceRoot(ctx, project) + require.NoError(t, err) + assert.Equal(t, "/repo", root) + assert.Equal(t, project, dir) + assert.Equal(t, []string{uvStubCommand("workspace dir")}, stub.Commands()) +} + +func TestWorkspaceRootReturnsUvError(t *testing.T) { + ctx, stub := process.WithStub(t.Context()) + stub.WithFailureFor(`workspace dir`, errors.New("exit status 2")) + m := &uvManager{bin: "uv"} + + _, err := m.WorkspaceRoot(ctx, "proj") + assert.Error(t, err) } func TestEnsurePythonStopsAfterSuccessfulInstall(t *testing.T) { diff --git a/libs/localenv/workspace.go b/libs/localenv/workspace.go new file mode 100644 index 00000000000..59ae59ac35c --- /dev/null +++ b/libs/localenv/workspace.go @@ -0,0 +1,86 @@ +package localenv + +import ( + "context" + "os" + "path/filepath" + + "github.com/databricks/cli/libs/log" +) + +// discoverWorkspaceRoot returns the uv workspace root when ProjectDir is a +// workspace member, and the root relative to ProjectDir (for example "../.."). +// It returns "" for both when ProjectDir is not a member: a standalone project, +// the root itself, or a project uv cannot place. +// +// A folder with no pyproject.toml yet (greenfield) is not treated as a member. +// When a members glob of a parent workspace matches it, the file this command +// creates makes it a member: the first run then fails at provision (uv syncs +// the root venv), and a re-run finds the member. +// +// A discovery failure (uv older than 0.9.9, uv not found under --dry-run, a +// malformed workspace) keeps the standalone behavior. A malformed workspace +// still fails later in uv sync, with uv's own message. +func (p *Pipeline) discoverWorkspaceRoot(ctx context.Context) (root, rel string) { + if !fileExists(p.pyprojectPath()) { + return "", "" + } + root, err := p.PM.WorkspaceRoot(ctx, p.ProjectDir) + if err != nil { + log.Debugf(ctx, CommandName+": workspace discovery failed, using the project folder: %v: %s", err, uvStderr(err)) + return "", "" + } + // uv prints the root with symlinks resolved and in the on-disk letter case, + // while ProjectDir keeps the form the shell gave (on macOS, /var vs + // /private/var, and any letter case). So compare folders by identity. + // + // The relative path is counted from the symlink-resolved ProjectDir, because + // that is where the OS resolves "..". The root is reported in ProjectDir's + // form when one of its parents is the root, so the paths the command reports + // match the folder the user ran it from; otherwise uv's form is used. + resolved, err := filepath.EvalSymlinks(p.ProjectDir) + if err != nil { + log.Debugf(ctx, CommandName+": cannot resolve the project folder: %v", err) + return "", "" + } + _, levels, ok := sameFolderAncestor(resolved, root) + if !ok { + log.Debugf(ctx, CommandName+": workspace root %s is not a parent of the project folder", filepath.ToSlash(root)) + return "", "" + } + if levels == 0 { + return "", "" + } + if ancestor, _, ok := sameFolderAncestor(p.ProjectDir, root); ok { + return ancestor, upLevels(levels) + } + return root, upLevels(levels) +} + +// sameFolderAncestor returns the first of dir and its parents that is the same +// folder as target (by file identity, not by name), and how many levels above +// dir it is. +func sameFolderAncestor(dir, target string) (ancestor string, levels int, ok bool) { + targetInfo, err := os.Stat(target) + if err != nil { + return "", 0, false + } + for d := dir; ; d = filepath.Dir(d) { + if info, err := os.Stat(d); err == nil && os.SameFile(info, targetInfo) { + return d, levels, true + } + if filepath.Dir(d) == d { + return "", 0, false + } + levels++ + } +} + +// upLevels returns the relative path n folders up, for example "../.." for 2. +func upLevels(n int) string { + parts := make([]string, n) + for i := range parts { + parts[i] = ".." + } + return filepath.Join(parts...) +} diff --git a/libs/localenv/workspace_test.go b/libs/localenv/workspace_test.go new file mode 100644 index 00000000000..84a96608960 --- /dev/null +++ b/libs/localenv/workspace_test.go @@ -0,0 +1,394 @@ +package localenv + +import ( + "context" + "errors" + "os" + "path/filepath" + "runtime" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// memberLayout creates a uv workspace root with one member, bundles/a. +func memberLayout(t *testing.T) (root, member string) { + t.Helper() + root = t.TempDir() + member = filepath.Join(root, "bundles", "a") + require.NoError(t, os.MkdirAll(member, 0o755)) + require.NoError(t, os.WriteFile(filepath.Join(root, pyprojectFile), []byte("[tool.uv.workspace]\nmembers = [\"bundles/*\"]\n"), 0o644)) + require.NoError(t, os.WriteFile(filepath.Join(member, pyprojectFile), []byte("[project]\nname = \"a\"\n"), 0o644)) + return root, member +} + +func TestDiscoverWorkspaceRoot(t *testing.T) { + root, member := memberLayout(t) + greenfield := filepath.Join(root, "new") + require.NoError(t, os.MkdirAll(greenfield, 0o755)) + + tests := []struct { + name string + projectDir string + pm fakePM + wantRoot string + wantRel string + }{ + {"member", member, fakePM{root: root}, root, filepath.Join("..", "..")}, + {"standalone", member, fakePM{root: member}, "", ""}, + // The created pyproject.toml is not a member yet; see discoverWorkspaceRoot. + {"greenfield", greenfield, fakePM{root: root}, "", ""}, + {"uv failure", member, fakePM{rootErr: errors.New("unrecognized subcommand 'workspace'")}, "", ""}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + p := &Pipeline{ProjectDir: tc.projectDir, PM: tc.pm} + gotRoot, gotRel := p.discoverWorkspaceRoot(t.Context()) + assert.Equal(t, tc.wantRoot, gotRoot) + assert.Equal(t, tc.wantRel, gotRel) + }) + } +} + +func TestDiscoverWorkspaceRootSymlinkedSameDir(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("symlinks need extra privileges on Windows") + } + _, member := memberLayout(t) + link := filepath.Join(t.TempDir(), "link") + require.NoError(t, os.Symlink(member, link)) + // uv prints the resolved path; ProjectDir is the symlinked one. + p := &Pipeline{ProjectDir: link, PM: fakePM{root: member}} + got, _ := p.discoverWorkspaceRoot(t.Context()) + assert.Empty(t, got) +} + +// caseFoldedPath returns path with its last element in upper case, and skips +// the test on a case-sensitive file system, where that path does not exist. +func caseFoldedPath(t *testing.T, path string) string { + t.Helper() + folded := filepath.Join(filepath.Dir(path), strings.ToUpper(filepath.Base(path))) + if _, err := os.Stat(folded); err != nil { + t.Skip("the file system is case-sensitive") + } + return folded +} + +func TestDiscoverWorkspaceRootCaseMismatchStandalone(t *testing.T) { + // The shell keeps the letter case the user typed; uv prints the case on disk. + _, member := memberLayout(t) + p := &Pipeline{ProjectDir: caseFoldedPath(t, member), PM: fakePM{root: member}} + got, _ := p.discoverWorkspaceRoot(t.Context()) + assert.Empty(t, got) +} + +func TestDiscoverWorkspaceRootCaseMismatchMember(t *testing.T) { + root, _ := memberLayout(t) + foldedRoot := caseFoldedPath(t, root) + p := &Pipeline{ProjectDir: filepath.Join(foldedRoot, "bundles", "a"), PM: fakePM{root: root}} + gotRoot, gotRel := p.discoverWorkspaceRoot(t.Context()) + assert.Equal(t, foldedRoot, gotRoot) + assert.Equal(t, filepath.Join("..", ".."), gotRel) +} + +// newMemberPipeline returns a pipeline for the member of a memberLayout. +func newMemberPipeline(t *testing.T) (p *Pipeline, root, member string) { + t.Helper() + root, member = memberLayout(t) + // The member needs requires-python so the existing merge accepts it. + require.NoError(t, os.WriteFile(filepath.Join(member, pyprojectFile), + []byte("[project]\nname = \"a\"\nversion = \"0.1.0\"\nrequires-python = \">=3.10\"\n"), 0o644)) + p = newServerlessTestPipeline(t, member) + p.PM = fakePM{py: "3.12", dbc: "17.2.0", root: root} + return p, root, member +} + +func TestPipelineWorkspaceMember(t *testing.T) { + p, root, member := newMemberPipeline(t) + res, err := p.Run(t.Context()) + require.NoError(t, err) + + rootBytes, err := os.ReadFile(filepath.Join(root, pyprojectFile)) + require.NoError(t, err) + memberBytes, err := os.ReadFile(filepath.Join(member, pyprojectFile)) + require.NoError(t, err) + assert.Contains(t, string(rootBytes), "constraint-dependencies") + assert.Contains(t, string(rootBytes), "[tool.uv.workspace]") + assert.NotContains(t, string(memberBytes), "constraint-dependencies") + assert.Contains(t, string(memberBytes), "requires-python") + assert.Contains(t, string(memberBytes), `"databricks-connect~=17.2.0"`) + + assert.Equal(t, "../../.venv", res.VenvPath) + assert.Equal(t, filepath.ToSlash(root), res.UvWorkspaceRoot) + assert.Equal(t, filepath.ToSlash(filepath.Join(root, backupFile)), res.UvWorkspaceBackupPath) + assert.Equal(t, filepath.ToSlash(filepath.Join(member, backupFile)), res.BackupPath) + assert.NotContains(t, codes(res.Warnings), WarnUvWorkspaceConstraintsReplaced, "first run must not warn") +} + +func TestPipelineWorkspaceMemberIdempotent(t *testing.T) { + p, root, member := newMemberPipeline(t) + _, err := p.Run(t.Context()) + require.NoError(t, err) + res, err := p.Run(t.Context()) + require.NoError(t, err) + assert.Empty(t, res.BackupPath) + assert.Empty(t, res.UvWorkspaceBackupPath) + assert.NotContains(t, codes(res.Warnings), WarnUvWorkspaceConstraintsReplaced) + for _, dir := range []string{root, member} { + matches, err := filepath.Glob(filepath.Join(dir, "pyproject.toml.*.bak")) + require.NoError(t, err) + assert.Empty(t, matches, "no timestamped backup in %s", dir) + } +} + +func TestPipelineWorkspaceConstraintsReplacedWarning(t *testing.T) { + p, root, _ := newMemberPipeline(t) + _, err := p.Run(t.Context()) + require.NoError(t, err) + // Simulate another bundle's run: change the pins inside the managed block. + path := filepath.Join(root, pyprojectFile) + b, err := os.ReadFile(path) + require.NoError(t, err) + changed := strings.Replace(string(b), "constraint-dependencies = [", "constraint-dependencies = [\n \"other==1.0\",", 1) + require.NoError(t, os.WriteFile(path, []byte(changed), 0o644)) + + res, err := p.Run(t.Context()) + require.NoError(t, err) + assert.Contains(t, codes(res.Warnings), WarnUvWorkspaceConstraintsReplaced) +} + +func TestPipelineWorkspaceMemberNoConstraints(t *testing.T) { + p, root, _ := newMemberPipeline(t) + p.SkipConstraints = true + before, err := os.ReadFile(filepath.Join(root, pyprojectFile)) + require.NoError(t, err) + res, err := p.Run(t.Context()) + require.NoError(t, err) + after, err := os.ReadFile(filepath.Join(root, pyprojectFile)) + require.NoError(t, err) + assert.Equal(t, before, after) + assert.Empty(t, res.UvWorkspaceBackupPath) + assert.NoFileExists(t, filepath.Join(root, backupFile)) +} + +func TestPipelineWorkspaceMemberDryRun(t *testing.T) { + p, root, member := newMemberPipeline(t) + p.Check = true + before, err := os.ReadFile(filepath.Join(root, pyprojectFile)) + require.NoError(t, err) + res, err := p.Run(t.Context()) + require.NoError(t, err) + after, err := os.ReadFile(filepath.Join(root, pyprojectFile)) + require.NoError(t, err) + assert.Equal(t, before, after) + require.NotNil(t, res.Plan) + assert.Equal(t, filepath.ToSlash(filepath.Join(member, pyprojectFile)), res.Plan.WouldWrite) + assert.Equal(t, filepath.ToSlash(filepath.Join(root, pyprojectFile)), res.Plan.WouldWriteUvWorkspace) + assert.Equal(t, filepath.ToSlash(filepath.Join(root, backupFile)), res.Plan.WouldBackupUvWorkspace) + assert.Contains(t, res.Plan.Diff, "workspace root") +} + +func TestPipelineWorkspaceRootNotWritable(t *testing.T) { + if runtime.GOOS == "windows" || os.Getuid() == 0 { + t.Skip("directory permissions do not block writes on Windows or as root") + } + p, root, member := newMemberPipeline(t) + require.NoError(t, os.Chmod(root, 0o555)) + t.Cleanup(func() { _ = os.Chmod(root, 0o755) }) + memberBefore, err := os.ReadFile(filepath.Join(member, pyprojectFile)) + require.NoError(t, err) + + res, err := p.Run(t.Context()) + require.Error(t, err) + require.NotNil(t, res.Error) + assert.Equal(t, ErrNotWritable, res.Error.Code) + assert.Contains(t, res.Error.Error(), filepath.ToSlash(root)) + memberAfter, err := os.ReadFile(filepath.Join(member, pyprojectFile)) + require.NoError(t, err) + assert.Equal(t, memberBefore, memberAfter) +} + +func TestDiscoverWorkspaceRootKeepsProjectPathForm(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("symlinks need extra privileges on Windows") + } + root, _ := memberLayout(t) + link := filepath.Join(t.TempDir(), "link") + require.NoError(t, os.Symlink(root, link)) + resolved, err := filepath.EvalSymlinks(root) + require.NoError(t, err) + // uv prints the symlink-resolved root; ProjectDir reaches the member through + // the link. The root must come back in ProjectDir's form so the venv path + // relative to the project stays "../../.venv". + p := &Pipeline{ProjectDir: filepath.Join(link, "bundles", "a"), PM: fakePM{root: resolved}} + got, _ := p.discoverWorkspaceRoot(t.Context()) + assert.Equal(t, link, got) +} + +// symlinkBelowRootLayout reaches the member through a link to the workspace's +// bundles folder (link -> /bundles), so the link sits below the root. +func symlinkBelowRootLayout(t *testing.T) (resolvedRoot, projectDir string) { + t.Helper() + if runtime.GOOS == "windows" { + t.Skip("symlinks need extra privileges on Windows") + } + root, _ := memberLayout(t) + link := filepath.Join(t.TempDir(), "link") + require.NoError(t, os.Symlink(filepath.Join(root, "bundles"), link)) + resolvedRoot, err := filepath.EvalSymlinks(root) + require.NoError(t, err) + return resolvedRoot, filepath.Join(link, "a") +} + +func TestDiscoverWorkspaceRootSymlinkBelowRoot(t *testing.T) { + resolvedRoot, projectDir := symlinkBelowRootLayout(t) + // Joining "../.." onto the link path would give the link's parent, not the + // root; the root must still name the real workspace root. + p := &Pipeline{ProjectDir: projectDir, PM: fakePM{root: resolvedRoot}} + gotRoot, gotRel := p.discoverWorkspaceRoot(t.Context()) + assert.Equal(t, resolvedRoot, gotRoot) + // Relative to the member's real folder (/bundles/a). + assert.Equal(t, filepath.Join("..", ".."), gotRel) +} + +func TestPipelineWorkspaceMemberVenvPathWithResolvedUvRoot(t *testing.T) { + // Real uv prints the symlink-resolved root (macOS /private/var). + p, root, _ := newMemberPipeline(t) + resolved, err := filepath.EvalSymlinks(root) + require.NoError(t, err) + p.PM = fakePM{py: "3.12", dbc: "17.2.0", root: resolved} + res, err := p.Run(t.Context()) + require.NoError(t, err) + assert.Equal(t, "../../.venv", res.VenvPath) +} + +func TestPipelineWorkspaceMemberSymlinkBelowRoot(t *testing.T) { + resolvedRoot, projectDir := symlinkBelowRootLayout(t) + require.NoError(t, os.WriteFile(filepath.Join(projectDir, pyprojectFile), + []byte("[project]\nname = \"a\"\nversion = \"0.1.0\"\nrequires-python = \">=3.10\"\n"), 0o644)) + p := newServerlessTestPipeline(t, projectDir) + p.PM = fakePM{py: "3.12", dbc: "17.2.0", root: resolvedRoot} + res, err := p.Run(t.Context()) + require.NoError(t, err) + // Relative to the member's real folder (/bundles/a), which is what the + // OS resolves a relative path against. + assert.Equal(t, "../../.venv", res.VenvPath) + rootBytes, err := os.ReadFile(filepath.Join(resolvedRoot, pyprojectFile)) + require.NoError(t, err) + assert.Contains(t, string(rootBytes), "constraint-dependencies") +} + +func TestPipelineWorkspaceReplacesUserRootPinsWithWarning(t *testing.T) { + p, root, _ := newMemberPipeline(t) + require.NoError(t, os.WriteFile(filepath.Join(root, pyprojectFile), + []byte("[tool.uv.workspace]\nmembers = [\"bundles/*\"]\n\n[tool.uv]\nconstraint-dependencies = [\"urllib3<2\"]\n"), 0o644)) + + res, err := p.Run(t.Context()) + require.NoError(t, err) + assert.Contains(t, codes(res.Warnings), WarnUvWorkspaceConstraintsReplaced) +} + +func TestApplyMergeRootFailureBeforeAnyWriteIsNotDiskMutated(t *testing.T) { + // The member merge is a no-op and the root file cannot be read, so nothing on + // disk changed and the failure must say so. + _, member := memberLayout(t) + memberPath := filepath.Join(member, pyprojectFile) + current, err := os.ReadFile(memberPath) + require.NoError(t, err) + p := &Pipeline{ProjectDir: member, res: NewResult()} + p.res.Phases = initialPhases() + missing := filepath.Join(t.TempDir(), pyprojectFile) + + err = p.applyMerge(t.Context(), current, &rootMerge{path: missing, merged: []byte("x"), changed: true}, false) + require.Error(t, err) + require.NotNil(t, p.res.Error) + assert.False(t, p.res.Error.DiskMutated) + assert.Contains(t, p.res.Error.Error(), "workspace root") +} + +func TestApplyMergeSkipsUnchangedRoot(t *testing.T) { + // The plan already read the root and found no change, so applyMerge must not + // touch it again: a missing root file is not an error here. + _, member := memberLayout(t) + current, err := os.ReadFile(filepath.Join(member, pyprojectFile)) + require.NoError(t, err) + p := &Pipeline{ProjectDir: member, res: NewResult()} + p.res.Phases = initialPhases() + missing := filepath.Join(t.TempDir(), pyprojectFile) + + err = p.applyMerge(t.Context(), current, &rootMerge{path: missing, merged: []byte("x")}, false) + require.NoError(t, err) + assert.Empty(t, p.res.UvWorkspaceBackupPath) +} + +func TestDiscoverWorkspaceRootSymlinkChangesDepth(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("symlinks need extra privileges on Windows") + } + // /link -> /deep/bundles: the logical path reaches the root in + // two levels, but the OS resolves a relative path from the real member folder + // (/deep/bundles/a), three levels below the root. + root := t.TempDir() + real := filepath.Join(root, "deep", "bundles", "a") + require.NoError(t, os.MkdirAll(real, 0o755)) + require.NoError(t, os.WriteFile(filepath.Join(real, pyprojectFile), []byte("[project]\nname = \"a\"\n"), 0o644)) + require.NoError(t, os.Symlink(filepath.Join(root, "deep", "bundles"), filepath.Join(root, "link"))) + + p := &Pipeline{ProjectDir: filepath.Join(root, "link", "a"), PM: fakePM{root: root}} + gotRoot, gotRel := p.discoverWorkspaceRoot(t.Context()) + assert.Equal(t, root, gotRoot) + assert.Equal(t, filepath.Join("..", "..", ".."), gotRel) +} + +func TestPipelineWorkspaceMemberDryRunReRunPlansNoRootWrite(t *testing.T) { + p, _, _ := newMemberPipeline(t) + _, err := p.Run(t.Context()) + require.NoError(t, err) + + p.Check = true + res, err := p.Run(t.Context()) + require.NoError(t, err) + require.NotNil(t, res.Plan) + assert.Empty(t, res.Plan.WouldWriteUvWorkspace) + assert.Empty(t, res.Plan.WouldBackupUvWorkspace) + assert.NotEmpty(t, res.UvWorkspaceRoot) +} + +// cancelDiscoveryPM cancels the run while workspace discovery runs, like a +// SIGINT that interrupts `uv workspace dir`. +type cancelDiscoveryPM struct { + fakePM + cancel context.CancelFunc +} + +func (f cancelDiscoveryPM) WorkspaceRoot(ctx context.Context, _ string) (string, error) { + f.cancel() + return "", ctx.Err() +} + +func TestPipelineCancelDuringDiscoveryWritesNothing(t *testing.T) { + p, root, member := newMemberPipeline(t) + // A first run fills the constraint cache, so the canceled run below can + // still fetch; then the member file is reset so a merge would change it. + _, err := p.Run(t.Context()) + require.NoError(t, err) + require.NoError(t, os.WriteFile(filepath.Join(member, pyprojectFile), + []byte("[project]\nname = \"a\"\nversion = \"0.1.0\"\nrequires-python = \">=3.10\"\n"), 0o644)) + + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + p.PM = cancelDiscoveryPM{fakePM: fakePM{py: "3.12", dbc: "17.2.0", root: root}, cancel: cancel} + before, err := os.ReadFile(filepath.Join(member, pyprojectFile)) + require.NoError(t, err) + + res, err := p.Run(ctx) + require.Error(t, err) + require.NotNil(t, res.Error) + assert.Equal(t, ErrCanceled, res.Error.Code) + assert.False(t, res.Error.DiskMutated) + after, err := os.ReadFile(filepath.Join(member, pyprojectFile)) + require.NoError(t, err) + assert.Equal(t, before, after) +}