diff --git a/CHANGELOG.md b/CHANGELOG.md index fad9ce620..3347da6df 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,29 @@ # Version changelog +## [Release] Release v0.243.0 + +CLI: + * Upgrade Go SDK to 0.59.0 ([#2425](https://github.com/databricks/cli/pull/2425)). + +Bundles: + * Added a warning when `config` section is used in apps ([#2416](https://github.com/databricks/cli/pull/2416)). + * Switch to use GET workspaces-files/{name} instead of workspace/export for state files to avoid 10MB limit ([#2423](https://github.com/databricks/cli/pull/2423)). + * Use schema field for pipeline in builtin template ([#2347](https://github.com/databricks/cli/pull/2347)). + * Add warning when variable interpolation is used for auth fields ([#2399](https://github.com/databricks/cli/pull/2399)). + * Add warning when include is used in config files other than databricks.yml ([#2389](https://github.com/databricks/cli/pull/2389)). + * Add support for schemas in deployment bind/unbind commands ([#2406](https://github.com/databricks/cli/pull/2406)). + * Do not modify/create .gitignore in bundle root ([#2429](https://github.com/databricks/cli/pull/2429)). + * Raise an error when there are multiple local libraries with the same basename used ([#2382](https://github.com/databricks/cli/pull/2382)). + * Upgrade TF provider to 1.68.0 ([#2426](https://github.com/databricks/cli/pull/2426)). + +API Changes: + * Changed `databricks experiments log-inputs` command with new required argument order. + * Added `databricks genie get-space` command. + * Added `databricks providers list-provider-share-assets` command. + * Changed `databricks shares update-permissions` command return type to become non-empty. + +OpenAPI commit e5c870006a536121442cfd2441bdc8a5fb76ae1e (2025-03-03) + ## [Release] Release v0.242.0 Notable changes: diff --git a/acceptance/bin/setmtime.py b/acceptance/bin/setmtime.py new file mode 100755 index 000000000..dd546d65c --- /dev/null +++ b/acceptance/bin/setmtime.py @@ -0,0 +1,17 @@ +#!/usr/bin/env python3 +""" +Cross-platform set mtime with nanosecond precision. +Usage: setmtime.py +""" + +import sys +import os +import datetime + +timestamp = sys.argv[1] +ts, ns = timestamp.split(".") +dt = datetime.datetime.strptime(ts, "%Y-%m-%d %H:%M:%S").replace(tzinfo=datetime.timezone.utc) +ns = int(ns.ljust(9, "0")) +ts = int(dt.timestamp()) * 10**9 + ns +for filename in sys.argv[2:]: + os.utime(filename, ns=(ts, ts)) diff --git a/acceptance/cmd/patchwhl/my_test_code-0.0.1-py3-none-any.whl b/acceptance/cmd/patchwhl/my_test_code-0.0.1-py3-none-any.whl new file mode 100644 index 000000000..4bb80477c Binary files /dev/null and b/acceptance/cmd/patchwhl/my_test_code-0.0.1-py3-none-any.whl differ diff --git a/acceptance/cmd/patchwhl/output.txt b/acceptance/cmd/patchwhl/output.txt new file mode 100644 index 000000000..13f1f1dce --- /dev/null +++ b/acceptance/cmd/patchwhl/output.txt @@ -0,0 +1,32 @@ + +=== Test prebuilt wheel: +>>> setmtime.py 2025-03-05 15:07:33.123456700 my_test_code-0.0.1-py3-none-any.whl + +>>> [CLI] selftest patchwhl my_test_code-0.0.1-py3-none-any.whl +Warn: Patched whl: my_test_code-0.0.1-py3-none-any.whl -> my_test_code-0.0.1+1741187253123456700-py3-none-any.whl + +>>> diff.py original output +--- original/my_test_code-0.0.1+1741187253123456700.dist-info/METADATA ++++ output/my_test_code-0.0.1+1741187253123456700.dist-info/METADATA +@@ -1,5 +1,5 @@ + Metadata-Version: 2.1 + Name: my-test-code +-Version: 0.0.1 ++Version: 0.0.1+1741187253123456700 + Summary: my test wheel + Home-page: https://databricks.com +--- original/my_test_code-0.0.1+1741187253123456700.dist-info/RECORD ++++ output/my_test_code-0.0.1+1741187253123456700.dist-info/RECORD +@@ -1,7 +1,7 @@ + src/__init__.py,sha256=BRmKeYehopKv4NG_SFa7t6wn248RrPHJivu7DM1R-Rw,48 + src/__main__.py,sha256=8TtsnLsaJEM35Y4L8ocrv-qfxusgYpRL2HPyYiabHng,242 +-my_test_code-0.0.1.dist-info/METADATA,sha256=6fyVq4hexXGUP_J2mB1CI-ijZ6CenvKNIlHx0bKPRJ0,197 +-my_test_code-0.0.1.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92 +-my_test_code-0.0.1.dist-info/entry_points.txt,sha256=oDWOW9SsBlk4Uejj1ftYPBxfhJ5ZJctb4JOUIG1rc-4,34 +-my_test_code-0.0.1.dist-info/top_level.txt,sha256=74rtVfumQlgAPzR5_2CgYN24MB0XARCg0t-gzk6gTrM,4 +-my_test_code-0.0.1.dist-info/RECORD,, ++my_test_code-0.0.1+1741187253123456700.dist-info/METADATA,sha256=H99P2vEwB_hBVPNtPwsXZotaDQzmWEGeSlOtMzWe62U,217 ++my_test_code-0.0.1+1741187253123456700.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92 ++my_test_code-0.0.1+1741187253123456700.dist-info/entry_points.txt,sha256=oDWOW9SsBlk4Uejj1ftYPBxfhJ5ZJctb4JOUIG1rc-4,34 ++my_test_code-0.0.1+1741187253123456700.dist-info/top_level.txt,sha256=74rtVfumQlgAPzR5_2CgYN24MB0XARCg0t-gzk6gTrM,4 ++my_test_code-0.0.1+1741187253123456700.dist-info/RECORD,, diff --git a/acceptance/cmd/patchwhl/output/my_test_code-0.0.1+1741187253123456700.dist-info/METADATA b/acceptance/cmd/patchwhl/output/my_test_code-0.0.1+1741187253123456700.dist-info/METADATA new file mode 100644 index 000000000..2308a523d --- /dev/null +++ b/acceptance/cmd/patchwhl/output/my_test_code-0.0.1+1741187253123456700.dist-info/METADATA @@ -0,0 +1,9 @@ +Metadata-Version: 2.1 +Name: my-test-code +Version: 0.0.1+1741187253123456700 +Summary: my test wheel +Home-page: https://databricks.com +Author: Databricks +Author-email: john.doe@databricks.com +Requires-Dist: setuptools + diff --git a/acceptance/cmd/patchwhl/output/my_test_code-0.0.1+1741187253123456700.dist-info/RECORD b/acceptance/cmd/patchwhl/output/my_test_code-0.0.1+1741187253123456700.dist-info/RECORD new file mode 100644 index 000000000..d7eb55dc1 --- /dev/null +++ b/acceptance/cmd/patchwhl/output/my_test_code-0.0.1+1741187253123456700.dist-info/RECORD @@ -0,0 +1,7 @@ +src/__init__.py,sha256=BRmKeYehopKv4NG_SFa7t6wn248RrPHJivu7DM1R-Rw,48 +src/__main__.py,sha256=8TtsnLsaJEM35Y4L8ocrv-qfxusgYpRL2HPyYiabHng,242 +my_test_code-0.0.1+1741187253123456700.dist-info/METADATA,sha256=H99P2vEwB_hBVPNtPwsXZotaDQzmWEGeSlOtMzWe62U,217 +my_test_code-0.0.1+1741187253123456700.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92 +my_test_code-0.0.1+1741187253123456700.dist-info/entry_points.txt,sha256=oDWOW9SsBlk4Uejj1ftYPBxfhJ5ZJctb4JOUIG1rc-4,34 +my_test_code-0.0.1+1741187253123456700.dist-info/top_level.txt,sha256=74rtVfumQlgAPzR5_2CgYN24MB0XARCg0t-gzk6gTrM,4 +my_test_code-0.0.1+1741187253123456700.dist-info/RECORD,, diff --git a/acceptance/cmd/patchwhl/output/my_test_code-0.0.1+1741187253123456700.dist-info/WHEEL b/acceptance/cmd/patchwhl/output/my_test_code-0.0.1+1741187253123456700.dist-info/WHEEL new file mode 100644 index 000000000..98c0d20b7 --- /dev/null +++ b/acceptance/cmd/patchwhl/output/my_test_code-0.0.1+1741187253123456700.dist-info/WHEEL @@ -0,0 +1,5 @@ +Wheel-Version: 1.0 +Generator: bdist_wheel (0.42.0) +Root-Is-Purelib: true +Tag: py3-none-any + diff --git a/acceptance/cmd/patchwhl/output/my_test_code-0.0.1+1741187253123456700.dist-info/entry_points.txt b/acceptance/cmd/patchwhl/output/my_test_code-0.0.1+1741187253123456700.dist-info/entry_points.txt new file mode 100644 index 000000000..85a57f7d3 --- /dev/null +++ b/acceptance/cmd/patchwhl/output/my_test_code-0.0.1+1741187253123456700.dist-info/entry_points.txt @@ -0,0 +1,2 @@ +[group_1] +run = src.__main__:main diff --git a/acceptance/cmd/patchwhl/output/my_test_code-0.0.1+1741187253123456700.dist-info/top_level.txt b/acceptance/cmd/patchwhl/output/my_test_code-0.0.1+1741187253123456700.dist-info/top_level.txt new file mode 100644 index 000000000..85de9cf93 --- /dev/null +++ b/acceptance/cmd/patchwhl/output/my_test_code-0.0.1+1741187253123456700.dist-info/top_level.txt @@ -0,0 +1 @@ +src diff --git a/acceptance/cmd/patchwhl/output/src/__init__.py b/acceptance/cmd/patchwhl/output/src/__init__.py new file mode 100644 index 000000000..909f1f322 --- /dev/null +++ b/acceptance/cmd/patchwhl/output/src/__init__.py @@ -0,0 +1,2 @@ +__version__ = "0.0.1" +__author__ = "Databricks" diff --git a/acceptance/cmd/patchwhl/output/src/__main__.py b/acceptance/cmd/patchwhl/output/src/__main__.py new file mode 100644 index 000000000..73d045afb --- /dev/null +++ b/acceptance/cmd/patchwhl/output/src/__main__.py @@ -0,0 +1,16 @@ +""" +The entry point of the Python Wheel +""" + +import sys + + +def main(): + # This method will print the provided arguments + print('Hello from my func') + print('Got arguments:') + print(sys.argv) + + +if __name__ == '__main__': + main() diff --git a/acceptance/cmd/patchwhl/ruff.toml b/acceptance/cmd/patchwhl/ruff.toml new file mode 100644 index 000000000..9ae5435d0 --- /dev/null +++ b/acceptance/cmd/patchwhl/ruff.toml @@ -0,0 +1 @@ +exclude = ["output"] diff --git a/acceptance/cmd/patchwhl/script b/acceptance/cmd/patchwhl/script new file mode 100644 index 000000000..26a0688eb --- /dev/null +++ b/acceptance/cmd/patchwhl/script @@ -0,0 +1,14 @@ +title "Test prebuilt wheel:" + +trace setmtime.py "2025-03-05 15:07:33.123456700" my_test_code-0.0.1-py3-none-any.whl +trace $CLI selftest patchwhl my_test_code-0.0.1-py3-none-any.whl + +mkdir output original +unzip -q my_test_code-0.0.1+1*-py3-none-any.whl -d output +unzip -q my_test_code-0.0.1-py3-none-any.whl -d original +rm my_test_code-0.0.1+1*-py3-none-any.whl + +# rename directory to match so that we can compare contents +mv original/my_test_code-0.0.1.dist-info original/my_test_code-0.0.1+1741187253123456700.dist-info +trace diff.py original output +rm -fr original diff --git a/bundle/config/mutator/apply_source_linked_deployment_preset_test.go b/bundle/config/mutator/apply_source_linked_deployment_preset_test.go index 42fda8ea7..0dbe3a441 100644 --- a/bundle/config/mutator/apply_source_linked_deployment_preset_test.go +++ b/bundle/config/mutator/apply_source_linked_deployment_preset_test.go @@ -36,13 +36,13 @@ func TestApplyPresetsSourceLinkedDeployment(t *testing.T) { }{ { name: "preset enabled, bundle in Workspace, databricks runtime", - ctx: dbr.MockRuntime(testContext, true), + ctx: dbr.MockRuntime(testContext, dbr.Environment{IsDbr: true, Version: "15.4"}), initialValue: &enabled, expectedValue: &enabled, }, { name: "preset enabled, bundle not in Workspace, databricks runtime", - ctx: dbr.MockRuntime(testContext, true), + ctx: dbr.MockRuntime(testContext, dbr.Environment{IsDbr: true, Version: "15.4"}), mutateBundle: func(b *bundle.Bundle) { b.SyncRootPath = "/Users/user.name@company.com" }, @@ -52,26 +52,26 @@ func TestApplyPresetsSourceLinkedDeployment(t *testing.T) { }, { name: "preset enabled, bundle in Workspace, not databricks runtime", - ctx: dbr.MockRuntime(testContext, false), + ctx: dbr.MockRuntime(testContext, dbr.Environment{}), initialValue: &enabled, expectedValue: &disabled, expectedWarning: "source-linked deployment is available only in the Databricks Workspace", }, { name: "preset disabled, bundle in Workspace, databricks runtime", - ctx: dbr.MockRuntime(testContext, true), + ctx: dbr.MockRuntime(testContext, dbr.Environment{IsDbr: true, Version: "15.4"}), initialValue: &disabled, expectedValue: &disabled, }, { name: "preset nil, bundle in Workspace, databricks runtime", - ctx: dbr.MockRuntime(testContext, true), + ctx: dbr.MockRuntime(testContext, dbr.Environment{IsDbr: true, Version: "15.4"}), initialValue: nil, expectedValue: nil, }, { name: "preset nil, dev mode true, bundle in Workspace, databricks runtime", - ctx: dbr.MockRuntime(testContext, true), + ctx: dbr.MockRuntime(testContext, dbr.Environment{IsDbr: true, Version: "15.4"}), mutateBundle: func(b *bundle.Bundle) { b.Config.Bundle.Mode = config.Development }, @@ -80,7 +80,7 @@ func TestApplyPresetsSourceLinkedDeployment(t *testing.T) { }, { name: "preset enabled, workspace.file_path is defined by user", - ctx: dbr.MockRuntime(testContext, true), + ctx: dbr.MockRuntime(testContext, dbr.Environment{IsDbr: true, Version: "15.4"}), mutateBundle: func(b *bundle.Bundle) { b.Config.Workspace.FilePath = "file_path" }, @@ -90,7 +90,7 @@ func TestApplyPresetsSourceLinkedDeployment(t *testing.T) { }, { name: "preset enabled, apps is defined by user", - ctx: dbr.MockRuntime(testContext, true), + ctx: dbr.MockRuntime(testContext, dbr.Environment{IsDbr: true, Version: "15.4"}), mutateBundle: func(b *bundle.Bundle) { b.Config.Resources.Apps = map[string]*resources.App{ "app": {}, diff --git a/bundle/config/mutator/configure_wsfs_test.go b/bundle/config/mutator/configure_wsfs_test.go index 6f76293e0..6762a446b 100644 --- a/bundle/config/mutator/configure_wsfs_test.go +++ b/bundle/config/mutator/configure_wsfs_test.go @@ -47,7 +47,7 @@ func TestConfigureWSFS_SkipsIfNotRunningOnRuntime(t *testing.T) { originalSyncRoot := b.SyncRoot ctx := context.Background() - ctx = dbr.MockRuntime(ctx, false) + ctx = dbr.MockRuntime(ctx, dbr.Environment{}) diags := bundle.Apply(ctx, b, mutator.ConfigureWSFS()) assert.Empty(t, diags) assert.Equal(t, originalSyncRoot, b.SyncRoot) @@ -58,7 +58,7 @@ func TestConfigureWSFS_SwapSyncRoot(t *testing.T) { originalSyncRoot := b.SyncRoot ctx := context.Background() - ctx = dbr.MockRuntime(ctx, true) + ctx = dbr.MockRuntime(ctx, dbr.Environment{IsDbr: true, Version: "15.4"}) diags := bundle.Apply(ctx, b, mutator.ConfigureWSFS()) assert.Empty(t, diags) assert.NotEqual(t, originalSyncRoot, b.SyncRoot) diff --git a/bundle/tests/loader.go b/bundle/tests/loader.go index 6748e6409..bd5ba4cba 100644 --- a/bundle/tests/loader.go +++ b/bundle/tests/loader.go @@ -69,7 +69,7 @@ func initializeTarget(t *testing.T, path, env string) (*bundle.Bundle, diag.Diag b := load(t, path) configureMock(t, b) - ctx := dbr.MockRuntime(context.Background(), false) + ctx := dbr.MockRuntime(context.Background(), dbr.Environment{}) diags := bundle.Apply(ctx, b, mutator.SelectTarget(env)) diags = diags.Extend(phases.Initialize(ctx, b)) diff --git a/cmd/selftest/patchwhl.go b/cmd/selftest/patchwhl.go new file mode 100644 index 000000000..25a792952 --- /dev/null +++ b/cmd/selftest/patchwhl.go @@ -0,0 +1,24 @@ +package selftest + +import ( + "github.com/databricks/cli/libs/log" + "github.com/databricks/cli/libs/patchwheel" + "github.com/spf13/cobra" +) + +func newPatchWhl() *cobra.Command { + return &cobra.Command{ + Use: "patchwhl", + Run: func(cmd *cobra.Command, args []string) { + ctx := cmd.Context() + for _, arg := range args { + out, err := patchwheel.PatchWheel(ctx, arg, ".") + if err != nil { + log.Warnf(ctx, "Failed to patch whl: %s: %s", arg, err) + } else { + log.Warnf(ctx, "Patched whl: %s -> %s", arg, out) + } + } + }, + } +} diff --git a/cmd/selftest/selftest.go b/cmd/selftest/selftest.go index b3e8541d4..039285399 100644 --- a/cmd/selftest/selftest.go +++ b/cmd/selftest/selftest.go @@ -13,5 +13,6 @@ func New() *cobra.Command { cmd.AddCommand(newSendTelemetry()) cmd.AddCommand(newPanic()) + cmd.AddCommand(newPatchWhl()) return cmd } diff --git a/integration/libs/git/git_fetch_test.go b/integration/libs/git/git_fetch_test.go index 0998d775b..73e33de73 100644 --- a/integration/libs/git/git_fetch_test.go +++ b/integration/libs/git/git_fetch_test.go @@ -50,7 +50,7 @@ func TestFetchRepositoryInfoAPI_FromRepo(t *testing.T) { ctx, wt := acc.WorkspaceTest(t) targetPath := ensureWorkspacePrefix(acc.TemporaryRepo(wt, examplesRepoUrl)) - ctx = dbr.MockRuntime(ctx, true) + ctx = dbr.MockRuntime(ctx, dbr.Environment{IsDbr: true, Version: "15.4"}) for _, inputPath := range []string{ path.Join(targetPath, "knowledge_base/dashboard_nyc_taxi"), @@ -72,7 +72,7 @@ func TestFetchRepositoryInfoAPI_FromNonRepo(t *testing.T) { err := wt.W.Workspace.MkdirsByPath(ctx, path.Join(rootPath, "a/b/c")) require.NoError(t, err) - ctx = dbr.MockRuntime(ctx, true) + ctx = dbr.MockRuntime(ctx, dbr.Environment{IsDbr: true, Version: "15.4"}) tests := []struct { input string diff --git a/libs/dbr/context.go b/libs/dbr/context.go index 7512c0fe2..303e911fb 100644 --- a/libs/dbr/context.go +++ b/libs/dbr/context.go @@ -15,6 +15,11 @@ const ( dbrKey = key(1) ) +type Environment struct { + IsDbr bool + Version string +} + // DetectRuntime detects whether or not the current // process is running inside a Databricks Runtime environment. // It return a new context with the detection result set. @@ -27,11 +32,11 @@ func DetectRuntime(ctx context.Context) context.Context { // MockRuntime is a helper function to mock the detection result. // It returns a new context with the detection result set. -func MockRuntime(ctx context.Context, b bool) context.Context { +func MockRuntime(ctx context.Context, runtime Environment) context.Context { if v := ctx.Value(dbrKey); v != nil { panic("dbr.MockRuntime called twice on the same context") } - return context.WithValue(ctx, dbrKey, b) + return context.WithValue(ctx, dbrKey, runtime) } // RunsOnRuntime returns the detection result from the context. @@ -45,5 +50,14 @@ func RunsOnRuntime(ctx context.Context) bool { if v == nil { panic("dbr.RunsOnRuntime called without calling dbr.DetectRuntime first") } - return v.(bool) + return v.(Environment).IsDbr +} + +func RuntimeVersion(ctx context.Context) string { + v := ctx.Value(dbrKey) + if v == nil { + panic("dbr.RuntimeVersion called without calling dbr.DetectRuntime first") + } + + return v.(Environment).Version } diff --git a/libs/dbr/context_test.go b/libs/dbr/context_test.go index fc53cf130..94e155ab5 100644 --- a/libs/dbr/context_test.go +++ b/libs/dbr/context_test.go @@ -23,11 +23,11 @@ func TestContext_MockRuntimePanics(t *testing.T) { ctx := context.Background() // Run detection. - ctx = MockRuntime(ctx, true) + ctx = MockRuntime(ctx, Environment{IsDbr: true, Version: "15.4"}) // Expect a panic if the mock function is run twice. assert.Panics(t, func() { - MockRuntime(ctx, true) + MockRuntime(ctx, Environment{IsDbr: true, Version: "15.4"}) }) } @@ -40,6 +40,15 @@ func TestContext_RunsOnRuntimePanics(t *testing.T) { }) } +func TestContext_RuntimeVersionPanics(t *testing.T) { + ctx := context.Background() + + // Expect a panic if the detection is not run. + assert.Panics(t, func() { + RuntimeVersion(ctx) + }) +} + func TestContext_RunsOnRuntime(t *testing.T) { ctx := context.Background() @@ -52,8 +61,26 @@ func TestContext_RunsOnRuntime(t *testing.T) { }) } +func TestContext_RuntimeVersion(t *testing.T) { + ctx := context.Background() + + // Run detection. + ctx = DetectRuntime(ctx) + + // Expect no panic because detection has run. + assert.NotPanics(t, func() { + RuntimeVersion(ctx) + }) +} + func TestContext_RunsOnRuntimeWithMock(t *testing.T) { ctx := context.Background() - assert.True(t, RunsOnRuntime(MockRuntime(ctx, true))) - assert.False(t, RunsOnRuntime(MockRuntime(ctx, false))) + assert.True(t, RunsOnRuntime(MockRuntime(ctx, Environment{IsDbr: true, Version: "15.4"}))) + assert.False(t, RunsOnRuntime(MockRuntime(ctx, Environment{}))) +} + +func TestContext_RuntimeVersionWithMock(t *testing.T) { + ctx := context.Background() + assert.Equal(t, "15.4", RuntimeVersion(MockRuntime(ctx, Environment{IsDbr: true, Version: "15.4"}))) + assert.Empty(t, RuntimeVersion(MockRuntime(ctx, Environment{}))) } diff --git a/libs/dbr/detect.go b/libs/dbr/detect.go index 12b83dc8c..a66f2fb77 100644 --- a/libs/dbr/detect.go +++ b/libs/dbr/detect.go @@ -11,27 +11,29 @@ import ( // Dereference [os.Stat] to allow mocking in tests. var statFunc = os.Stat -const EnvVarName = "DATABRICKS_RUNTIME_VERSION" - // detect returns true if the current process is running on a Databricks Runtime. // Its return value is meant to be cached in the context. -func detect(ctx context.Context) bool { +func detect(ctx context.Context) Environment { // Databricks Runtime implies Linux. // Return early on other operating systems. if runtime.GOOS != "linux" { - return false + return Environment{} } // Databricks Runtime always has the DATABRICKS_RUNTIME_VERSION environment variable set. - if value, ok := env.Lookup(ctx, EnvVarName); !ok || value == "" { - return false + version, ok := env.Lookup(ctx, "DATABRICKS_RUNTIME_VERSION") + if !ok || version == "" { + return Environment{} } // Expect to see a "/databricks" directory. if fi, err := statFunc("/databricks"); err != nil || !fi.IsDir() { - return false + return Environment{} } // All checks passed. - return true + return Environment{ + IsDbr: true, + Version: version, + } } diff --git a/libs/dbr/detect_test.go b/libs/dbr/detect_test.go index 3a4a43a73..9716b50f6 100644 --- a/libs/dbr/detect_test.go +++ b/libs/dbr/detect_test.go @@ -35,7 +35,7 @@ func TestDetect_NotLinux(t *testing.T) { } ctx := context.Background() - assert.False(t, detect(ctx)) + assert.Equal(t, Environment{}, detect(ctx)) } func TestDetect_Env(t *testing.T) { @@ -46,17 +46,23 @@ func TestDetect_Env(t *testing.T) { t.Run("empty", func(t *testing.T) { ctx := env.Set(context.Background(), "DATABRICKS_RUNTIME_VERSION", "") - assert.False(t, detect(ctx)) + assert.Equal(t, Environment{}, detect(ctx)) }) t.Run("non-empty cluster", func(t *testing.T) { ctx := env.Set(context.Background(), "DATABRICKS_RUNTIME_VERSION", "15.4") - assert.True(t, detect(ctx)) + assert.Equal(t, Environment{ + IsDbr: true, + Version: "15.4", + }, detect(ctx)) }) t.Run("non-empty serverless", func(t *testing.T) { ctx := env.Set(context.Background(), "DATABRICKS_RUNTIME_VERSION", "client.1.13") - assert.True(t, detect(ctx)) + assert.Equal(t, Environment{ + IsDbr: true, + Version: "client.1.13", + }, detect(ctx)) }) } @@ -68,16 +74,19 @@ func TestDetect_Stat(t *testing.T) { t.Run("error", func(t *testing.T) { configureStatFunc(t, nil, fs.ErrNotExist) - assert.False(t, detect(ctx)) + assert.Equal(t, Environment{}, detect(ctx)) }) t.Run("not a directory", func(t *testing.T) { configureStatFunc(t, fakefs.FileInfo{}, nil) - assert.False(t, detect(ctx)) + assert.Equal(t, Environment{}, detect(ctx)) }) t.Run("directory", func(t *testing.T) { configureStatFunc(t, fakefs.FileInfo{FakeDir: true}, nil) - assert.True(t, detect(ctx)) + assert.Equal(t, Environment{ + IsDbr: true, + Version: "non-empty", + }, detect(ctx)) }) } diff --git a/libs/patchwheel/parse.go b/libs/patchwheel/parse.go new file mode 100644 index 000000000..3b183e420 --- /dev/null +++ b/libs/patchwheel/parse.go @@ -0,0 +1,64 @@ +package patchwheel + +import ( + "fmt" + "strconv" + "strings" + "time" +) + +// WheelInfo contains information extracted from a wheel filename +type WheelInfo struct { + Distribution string // Package distribution name + Version string // Package version + Tags []string // [optional build tag,] python_tag, abi_tag, platform_tag +} + +// calculateNewVersion generates a new version string and filename based on the wheel info and modification time. +// The version is updated according to the following rules: +// - if there is an existing part after + it is dropped +// - append + to version +// +// Example version transform: "1.2.3" -> "1.2.3+1741091696780123321" +func calculateNewVersion(info WheelInfo, mtime time.Time) (newVersion, newFilename string) { + baseVersion, _, _ := strings.Cut(info.Version, "+") + + newVersion = baseVersion + "+" + strconv.FormatInt(mtime.UnixNano(), 10) + + newFilename = fmt.Sprintf("%s-%s-%s.whl", + info.Distribution, + newVersion, + strings.Join(info.Tags, "-")) + + return newVersion, newFilename +} + +// ParseWheelFilename parses a wheel filename and extracts its components. +// Wheel filenames follow the pattern: {distribution}-{version}(-{build tag})?-{python_tag}-{abi_tag}-{platform_tag}.whl +// https://peps.python.org/pep-0491 +// The function does not try hard to validate if the format is correct, it tries to parse whatever is available. +// However, it does require 5 or 6 components in the filename. +func ParseWheelFilename(filename string) (WheelInfo, error) { + parts := strings.Split(filename, "-") + if len(parts) < 5 { + return WheelInfo{}, fmt.Errorf("invalid wheel filename format: not enough parts in %s", filename) + } + + if len(parts) > 6 { + return WheelInfo{}, fmt.Errorf("invalid wheel filename format: too many parts in %s", filename) + } + + trimmedLastTag, foundWhl := strings.CutSuffix(parts[len(parts)-1], ".whl") + + if !foundWhl { + return WheelInfo{}, fmt.Errorf("invalid wheel filename format: missing .whl extension in %s", filename) + } + + parts[len(parts)-1] = trimmedLastTag + + return WheelInfo{ + Distribution: parts[0], + Version: parts[1], + Tags: parts[2:], + }, nil +} diff --git a/libs/patchwheel/parse_test.go b/libs/patchwheel/parse_test.go new file mode 100644 index 000000000..99eac4281 --- /dev/null +++ b/libs/patchwheel/parse_test.go @@ -0,0 +1,297 @@ +package patchwheel + +import ( + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +func TestCalculateNewVersion(t *testing.T) { + tests := []struct { + name string + info WheelInfo + mtime time.Time + expectedVersion string + expectedFilename string + }{ + { + name: "basic version", + info: WheelInfo{ + Distribution: "mypkg", + Version: "1.2.3", + Tags: []string{"py3", "none", "any"}, + }, + mtime: time.Date(2025, 3, 4, 12, 34, 56, 780_123_321, time.UTC), + expectedVersion: "1.2.3+1741091696780123321", + expectedFilename: "mypkg-1.2.3+1741091696780123321-py3-none-any.whl", + }, + { + name: "existing plus version", + info: WheelInfo{ + Distribution: "mypkg", + Version: "1.2.3+local", + Tags: []string{"py3", "none", "any"}, + }, + mtime: time.Date(2025, 3, 4, 12, 34, 56, 100_000_005, time.UTC), + expectedVersion: "1.2.3+1741091696100000005", + expectedFilename: "mypkg-1.2.3+1741091696100000005-py3-none-any.whl", + }, + { + name: "complex distribution name", + info: WheelInfo{ + Distribution: "my-pkg-name", + Version: "1.2.3", + Tags: []string{"py3", "none", "any"}, + }, + mtime: time.Date(2025, 3, 4, 12, 34, 56, 0, time.UTC), + expectedVersion: "1.2.3+1741091696000000000", + expectedFilename: "my-pkg-name-1.2.3+1741091696000000000-py3-none-any.whl", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + newVersion, newFilename := calculateNewVersion(tt.info, tt.mtime) + if newVersion != tt.expectedVersion { + t.Errorf("expected version %s, got %s", tt.expectedVersion, newVersion) + } + if newFilename != tt.expectedFilename { + t.Errorf("expected filename %s, got %s", tt.expectedFilename, newFilename) + } + }) + } +} + +func TestParseWheelFilename(t *testing.T) { + tests := []struct { + filename string + wantDistribution string + wantVersion string + wantTags []string + wantErr bool + }{ + { + filename: "myproj-0.1.0-py3-none-any.whl", + wantDistribution: "myproj", + wantVersion: "0.1.0", + wantTags: []string{"py3", "none", "any"}, + wantErr: false, + }, + { + filename: "myproj-0.1.0+20240303123456-py3-none-any.whl", + wantDistribution: "myproj", + wantVersion: "0.1.0+20240303123456", + wantTags: []string{"py3", "none", "any"}, + wantErr: false, + }, + { + filename: "my_proj_with_parts-0.1.0-py3-none-any.whl", + wantDistribution: "my_proj_with_parts", + wantVersion: "0.1.0", + wantTags: []string{"py3", "none", "any"}, + wantErr: false, + }, + // Test cases adapted from wheelodex/wheel-filename + // https://github.com/wheelodex/wheel-filename/blob/f5a72b560d016cc9663c8b5a094c96dc338a2209/test/test_parse.py + // MIT License: https://github.com/wheelodex/wheel-filename/blob/f5a72b560d016cc9663c8b5a094c96dc338a2209/LICENSE + { + filename: "astrocats-0.3.2-universal-none-any.whl", + wantDistribution: "astrocats", + wantVersion: "0.3.2", + wantTags: []string{"universal", "none", "any"}, + wantErr: false, + }, + { + filename: "bencoder.pyx-1.1.2-pp226-pp226-win32.whl", + wantDistribution: "bencoder.pyx", + wantVersion: "1.1.2", + wantTags: []string{"pp226", "pp226", "win32"}, + wantErr: false, + }, + { + filename: "brotlipy-0.1.2-pp27-none-macosx_10_10_x86_64.whl", + wantDistribution: "brotlipy", + wantVersion: "0.1.2", + wantTags: []string{"pp27", "none", "macosx_10_10_x86_64"}, + wantErr: false, + }, + { + filename: "brotlipy-0.3.0-pp226-pp226u-macosx_10_10_x86_64.whl", + wantDistribution: "brotlipy", + wantVersion: "0.3.0", + wantTags: []string{"pp226", "pp226u", "macosx_10_10_x86_64"}, + wantErr: false, + }, + { + filename: "carbonara_archinfo-7.7.9.14.post1-py2-none-any.whl", + wantDistribution: "carbonara_archinfo", + wantVersion: "7.7.9.14.post1", + wantTags: []string{"py2", "none", "any"}, + wantErr: false, + }, + { + filename: "coremltools-0.3.0-py2.7-none-any.whl", + wantDistribution: "coremltools", + wantVersion: "0.3.0", + wantTags: []string{"py2.7", "none", "any"}, + wantErr: false, + }, + { + filename: "cvxopt-1.2.0-001-cp34-cp34m-macosx_10_6_intel.macosx_10_9_intel.macosx_10_9_x86_64.macosx_10_10_intel.macosx_10_10_x86_64.whl", + wantDistribution: "cvxopt", + wantVersion: "1.2.0", + wantTags: []string{"001", "cp34", "cp34m", "macosx_10_6_intel.macosx_10_9_intel.macosx_10_9_x86_64.macosx_10_10_intel.macosx_10_10_x86_64"}, + wantErr: false, + }, + { + filename: "django_mbrowse-0.0.1-10-py2-none-any.whl", + wantDistribution: "django_mbrowse", + wantVersion: "0.0.1", + wantTags: []string{"10", "py2", "none", "any"}, + wantErr: false, + }, + { + filename: "efilter-1!1.2-py2-none-any.whl", + wantDistribution: "efilter", + wantVersion: "1!1.2", + wantTags: []string{"py2", "none", "any"}, + wantErr: false, + }, + { + filename: "line.sep-0.2.0.dev1-py2.py3-none-any.whl", + wantDistribution: "line.sep", + wantVersion: "0.2.0.dev1", + wantTags: []string{"py2.py3", "none", "any"}, + wantErr: false, + }, + { + filename: "mayan_edms-1.1.0-1502100955-py2-none-any.whl", + wantDistribution: "mayan_edms", + wantVersion: "1.1.0", + wantTags: []string{"1502100955", "py2", "none", "any"}, + wantErr: false, + }, + { + filename: "mxnet_model_server-1.0a5-20180816-py2.py3-none-any.whl", + wantDistribution: "mxnet_model_server", + wantVersion: "1.0a5", + wantTags: []string{"20180816", "py2.py3", "none", "any"}, + wantErr: false, + }, + { + filename: "pip-18.0-py2.py3-none-any.whl", + wantDistribution: "pip", + wantVersion: "18.0", + wantTags: []string{"py2.py3", "none", "any"}, + wantErr: false, + }, + { + filename: "polarTransform-2-1.0.0-py3-none-any.whl", + wantDistribution: "polarTransform", + wantVersion: "2", + wantTags: []string{"1.0.0", "py3", "none", "any"}, + wantErr: false, + }, + { + filename: "psycopg2-2.7.5-cp37-cp37m-macosx_10_6_intel.macosx_10_9_intel.macosx_10_9_x86_64.macosx_10_10_intel.macosx_10_10_x86_64.whl", + wantDistribution: "psycopg2", + wantVersion: "2.7.5", + wantTags: []string{"cp37", "cp37m", "macosx_10_6_intel.macosx_10_9_intel.macosx_10_9_x86_64.macosx_10_10_intel.macosx_10_10_x86_64"}, + wantErr: false, + }, + { + filename: "pyinterval-1.0.0-0-cp27-none-win32.whl", + wantDistribution: "pyinterval", + wantVersion: "1.0.0", + wantTags: []string{"0", "cp27", "none", "win32"}, + wantErr: false, + }, + { + filename: "pypi_simple-0.1.0.dev1-py2.py3-none-any.whl", + wantDistribution: "pypi_simple", + wantVersion: "0.1.0.dev1", + wantTags: []string{"py2.py3", "none", "any"}, + wantErr: false, + }, + { + filename: "PyQt3D-5.7.1-5.7.1-cp34.cp35.cp36-abi3-macosx_10_6_intel.whl", + wantDistribution: "PyQt3D", + wantVersion: "5.7.1", + wantTags: []string{"5.7.1", "cp34.cp35.cp36", "abi3", "macosx_10_6_intel"}, + wantErr: false, + }, + { + filename: "qypi-0.4.1-py3-none-any.whl", + wantDistribution: "qypi", + wantVersion: "0.4.1", + wantTags: []string{"py3", "none", "any"}, + wantErr: false, + }, + { + filename: "SimpleSteem-1.1.9-3.0-none-any.whl", + wantDistribution: "SimpleSteem", + wantVersion: "1.1.9", + wantTags: []string{"3.0", "none", "any"}, + wantErr: false, + }, + { + filename: "simple_workflow-0.1.47-pypy-none-any.whl", + wantDistribution: "simple_workflow", + wantVersion: "0.1.47", + wantTags: []string{"pypy", "none", "any"}, + wantErr: false, + }, + { + filename: "tables-3.4.2-3-cp27-cp27m-manylinux1_i686.whl", + wantDistribution: "tables", + wantVersion: "3.4.2", + wantTags: []string{"3", "cp27", "cp27m", "manylinux1_i686"}, + wantErr: false, + }, + { + filename: "invalid-filename.txt", + wantDistribution: "", + wantVersion: "", + wantTags: nil, + wantErr: true, + }, + { + filename: "not-enough-parts-py3.whl", + wantDistribution: "", + wantVersion: "", + wantTags: nil, + wantErr: true, + }, + } + + for _, tt := range tests { + t.Run(tt.filename, func(t *testing.T) { + info, err := ParseWheelFilename(tt.filename) + if tt.wantErr { + require.Error(t, err) + } else { + require.NoError(t, err) + require.Equal(t, tt.wantDistribution, info.Distribution, "distribution mismatch") + require.Equal(t, tt.wantVersion, info.Version, "version mismatch") + require.Equal(t, tt.wantTags, info.Tags, "tags mismatch") + } + }) + } +} + +func TestParseError(t *testing.T) { + tests := []string{ + "", + "a-b", + "a-b-c.whl", + "a-b-c-d.whl", + "a-b-c-d-e-f-g.whl", + } + for _, tt := range tests { + t.Run(tt, func(t *testing.T) { + info, err := ParseWheelFilename(tt) + require.Error(t, err, "info=%v", info) + }) + } +} diff --git a/libs/patchwheel/patch.go b/libs/patchwheel/patch.go new file mode 100644 index 000000000..672a17c81 --- /dev/null +++ b/libs/patchwheel/patch.go @@ -0,0 +1,268 @@ +package patchwheel + +import ( + "archive/zip" + "bufio" + "bytes" + "context" + "crypto/sha256" + "encoding/base64" + "errors" + "fmt" + "io" + "os" + "path/filepath" + "strconv" + "strings" + + "github.com/databricks/cli/libs/log" +) + +var versionKey []byte = []byte("Version:") + +func findFile(r *zip.ReadCloser, filename string) *zip.File { + for _, f := range r.File { + if f.Name == filename { + return f + } + } + return nil +} + +// patchMetadata returns new METADATA content with an updated "Version:" field and validates that previous version matches oldVersion +func patchMetadata(r io.Reader, oldVersion, newVersion string) ([]byte, error) { + scanner := bufio.NewScanner(r) + var buf bytes.Buffer + for scanner.Scan() { + line := scanner.Bytes() + if versionValue, ok := bytes.CutPrefix(line, versionKey); ok { + foundVersion := string(bytes.TrimSpace(versionValue)) + if foundVersion != oldVersion { + return nil, fmt.Errorf("Unexpected version in METADATA: %s (expected %s)", strings.TrimSpace(string(line)), oldVersion) + } + buf.WriteString(string(versionKey) + " " + newVersion + "\n") + } else { + buf.Write(line) + buf.WriteString("\n") + } + } + if err := scanner.Err(); err != nil { + return nil, err + } + return buf.Bytes(), nil +} + +// patchRecord updates RECORD content: it replaces the old dist-info prefix with the new one +// in all file paths and, for the METADATA entry, updates the hash and size. +func patchRecord(r io.Reader, oldDistInfoPrefix, newDistInfoPrefix, metadataHash string, metadataSize int) ([]byte, error) { + metadataPath := newDistInfoPrefix + "METADATA" + scanner := bufio.NewScanner(r) + var buf bytes.Buffer + for scanner.Scan() { + line := scanner.Bytes() + if len(bytes.TrimSpace(line)) == 0 { + continue + } + + parts := strings.Split(string(line), ",") + + if len(parts) < 3 { + // If the line doesn't have enough parts, preserve it as-is + buf.Write(line) + buf.WriteString("\n") + continue + } + + origPath := parts[0] + pathSuffix, hasDistPrefix := strings.CutPrefix(origPath, oldDistInfoPrefix) + if hasDistPrefix { + parts[0] = newDistInfoPrefix + pathSuffix + } + + if metadataPath == parts[0] { + parts[1] = "sha256=" + metadataHash + parts[2] = strconv.Itoa(metadataSize) + } + + buf.WriteString(strings.Join(parts, ",") + "\n") + } + if err := scanner.Err(); err != nil { + return nil, err + } + return buf.Bytes(), nil +} + +// PatchWheel reads an existing wheel file path and outputs a new one in outputDir, +// with a version modified according to the following rules: +// - if there is an existing part after + it is dropped +// - append + to version +// +// All parts of wheel are modified to ensure the wheel is in correct format: +// METADATA: Version field is updated +// RECORD: METADATA entry is updated with correct hash and size +// -.dist-info directory is renamed to -.dist-info +// +// The function is idempotent: repeated calls with the same input will produce the same output. +// If the target wheel already exists, it returns the path to the existing wheel without redoing the patching. +func PatchWheel(ctx context.Context, path, outputDir string) (string, error) { + fileInfo, err := os.Stat(path) + if err != nil { + return "", err + } + wheelMtime := fileInfo.ModTime().UTC() + + filename := filepath.Base(path) + wheelInfo, err := ParseWheelFilename(filename) + if err != nil { + return "", err + } + + newVersion, newFilename := calculateNewVersion(wheelInfo, wheelMtime) + outpath := filepath.Join(outputDir, newFilename) + + if _, err := os.Stat(outpath); err == nil { + log.Debugf(ctx, "Skipping patching of %s, already exists: %s", path, outpath) + return outpath, nil + } + + tmpFilename := outpath + fmt.Sprintf(".tmp%d", os.Getpid()) + + needRemoval := true + + defer func() { + if needRemoval { + _ = os.Remove(tmpFilename) + } + }() + + r, err := zip.OpenReader(path) + if err != nil { + return "", err + } + defer r.Close() + + oldDistInfoPrefix := wheelInfo.Distribution + "-" + wheelInfo.Version + ".dist-info/" + metadataFile := findFile(r, oldDistInfoPrefix+"METADATA") + if metadataFile == nil { + return "", fmt.Errorf("wheel %s missing %sMETADATA", path, oldDistInfoPrefix) + } + + recordFile := findFile(r, oldDistInfoPrefix+"RECORD") + if recordFile == nil { + return "", fmt.Errorf("wheel %s missing %sRECORD file", path, oldDistInfoPrefix) + } + + metadataReader, err := metadataFile.Open() + if err != nil { + return "", err + } + defer metadataReader.Close() + + newMetadata, err := patchMetadata(metadataReader, wheelInfo.Version, newVersion) + if err != nil { + return "", err + } + + h := sha256.New() + h.Write(newMetadata) + metadataHash := base64.URLEncoding.WithPadding(base64.NoPadding).EncodeToString(h.Sum(nil)) + metadataSize := len(newMetadata) + + // Compute the new dist-info directory prefix. + newDistInfoPrefix := strings.Replace(oldDistInfoPrefix, wheelInfo.Version, newVersion, 1) + if newDistInfoPrefix == oldDistInfoPrefix { + return "", fmt.Errorf("unexpected dist-info directory format: %s (version=%s)", oldDistInfoPrefix, wheelInfo.Version) + } + + recordReader, err := recordFile.Open() + if err != nil { + return "", err + } + defer recordReader.Close() + + newRecord, err := patchRecord(recordReader, oldDistInfoPrefix, newDistInfoPrefix, metadataHash, metadataSize) + if err != nil { + return "", err + } + + outFile, err := os.Create(tmpFilename) + if err != nil { + return "", err + } + defer outFile.Close() + + metadataUpdated := 0 + recordUpdated := 0 + + zipw := zip.NewWriter(outFile) + for _, f := range r.File { + // If the file is inside the old dist-info directory, update its name. + newName := f.Name + if strings.HasPrefix(f.Name, oldDistInfoPrefix) { + newName = newDistInfoPrefix + f.Name[len(oldDistInfoPrefix):] + } + header := &zip.FileHeader{ + Name: newName, + Method: f.Method, + } + + header.Modified = f.ModTime() + header.Comment = f.Comment + if f.FileInfo().IsDir() && !strings.HasSuffix(header.Name, "/") { + header.Name += "/" + } + + writer, err := zipw.CreateHeader(header) + if err != nil { + return "", err + } + + if f.Name == metadataFile.Name { + _, err = writer.Write(newMetadata) + if err != nil { + return "", err + } + metadataUpdated += 1 + } else if f.Name == recordFile.Name { + _, err = writer.Write(newRecord) + if err != nil { + return "", err + } + recordUpdated += 1 + } else { + rc, err := f.Open() + if err != nil { + return "", err + } + _, err = io.Copy(writer, rc) + if err != nil { + rc.Close() + return "", err + } + if err := rc.Close(); err != nil { + return "", err + } + } + } + + if err := zipw.Close(); err != nil { + return "", err + } + + outFile.Close() + + if metadataUpdated != 1 { + return "", errors.New("Could not update METADATA") + } + + if recordUpdated != 1 { + return "", errors.New("Could not update RECORD") + } + + if err := os.Rename(tmpFilename, outpath); err != nil { + return "", err + } + needRemoval = false + + return outpath, nil +} diff --git a/libs/patchwheel/patch_test.go b/libs/patchwheel/patch_test.go new file mode 100644 index 000000000..fae09aaed --- /dev/null +++ b/libs/patchwheel/patch_test.go @@ -0,0 +1,218 @@ +package patchwheel + +import ( + "bytes" + "context" + "os" + "os/exec" + "path/filepath" + "runtime" + "strings" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +var ( + scriptsDir = getPythonScriptsDir() + prebuiltWheel = "testdata/my_test_code-0.0.1-py3-none-any.whl" + emptyZip = "testdata/empty.zip" +) + +func getPythonScriptsDir() string { + if runtime.GOOS == "windows" { + return "Scripts" + } + return "bin" +} + +func getPythonVersions() []string { + return []string{ + "python3.9", + "python3.10", + "python3.11", + "python3.12", + "python3.13", + } +} + +func verifyVersion(t *testing.T, tempDir, wheelPath string) { + wheelInfo, err := ParseWheelFilename(wheelPath) + require.NoError(t, err) + expectedVersion := wheelInfo.Version + + pyExec := filepath.Join(tempDir, ".venv", scriptsDir, "python") + cmdOut := captureOutput(t, tempDir, pyExec, "-c", "import myproj; myproj.print_version()") + actualVersion := strings.TrimSpace(cmdOut) + t.Logf("Verified installed version: %s", actualVersion) + assert.Equal(t, expectedVersion, actualVersion, "Installed version doesn't match expected version from wheel filename") +} + +func minimalPythonProject() map[string]string { + return map[string]string{ + "pyproject.toml": `[project] +name = "myproj" +version = "0.1.0" + +[build-system] +requires = ["setuptools>=61.0.0", "wheel"] +build-backend = "setuptools.build_meta" + +[tool.setuptools.packages.find] +where = ["src"] +`, + "src/myproj/__init__.py": ` +def hello(): + return "Hello, world!" + +def print_version(): + from importlib.metadata import version + print(version("myproj")) +`, + } +} + +func writeProjectFiles(t *testing.T, baseDir string, files map[string]string) { + for path, content := range files { + fullPath := filepath.Join(baseDir, path) + if err := os.MkdirAll(filepath.Dir(fullPath), 0o755); err != nil { + t.Fatalf("Failed to create directory %s: %v", filepath.Dir(fullPath), err) + } + if err := os.WriteFile(fullPath, []byte(content), 0o644); err != nil { + t.Fatalf("Failed to write file %s: %v", fullPath, err) + } + } +} + +func runCmd(t *testing.T, dir, name string, args ...string) { + out := captureOutput(t, dir, name, args...) + if len(out) > 0 { + t.Errorf("Output from %s %s:\n%s", name, args, out) + } +} + +func captureOutput(t *testing.T, dir, name string, args ...string) string { + cmd := exec.Command(name, args...) + cmd.Dir = dir + var out bytes.Buffer + cmd.Stdout = &out + cmd.Stderr = &out + err := cmd.Run() + if err != nil { + t.Logf("Command failed: %s %s", name, strings.Join(args, " ")) + t.Logf("Output:\n%s", out.String()) + t.Fatal(err) + } + return out.String() +} + +func getWheel(t *testing.T, dir string) string { + pattern := filepath.Join(dir, "*.whl") + matches, err := filepath.Glob(pattern) + if err != nil { + t.Fatalf("Error matching pattern %s: %v", pattern, err) + } + + if len(matches) == 0 { + t.Fatalf("No files found matching %s", pattern) + return "" + } + + if len(matches) != 1 { + t.Fatalf("Too many matches %s: %v", pattern, matches) + return "" + } + + return matches[0] +} + +func TestPatchWheel(t *testing.T) { + pythonVersions := getPythonVersions() + for _, py := range pythonVersions { + t.Run(py, func(t *testing.T) { + t.Parallel() + tempDir := t.TempDir() + + projFiles := minimalPythonProject() + writeProjectFiles(t, tempDir, projFiles) + + runCmd(t, tempDir, "uv", "venv", "-q", "--python", py) + + runCmd(t, tempDir, "uv", "build", "-q", "--wheel") + distDir := filepath.Join(tempDir, "dist") + origWheel := getWheel(t, distDir) + + patchedWheel, err := PatchWheel(context.Background(), origWheel, distDir) + require.NoError(t, err) + + patchedInfo, err := os.Stat(patchedWheel) + require.NoError(t, err) + patchedTime := patchedInfo.ModTime() + + // Test idempotency - patching the same wheel again should produce the same result + // and should not recreate the file (file modification time should remain the same) + patchedWheel2, err := PatchWheel(context.Background(), origWheel, distDir) + require.NoError(t, err) + require.Equal(t, patchedWheel, patchedWheel2, "PatchWheel is not idempotent") + + patchedInfo2, err := os.Stat(patchedWheel2) + require.NoError(t, err) + require.Equal(t, patchedTime, patchedInfo2.ModTime(), "File was recreated when it shouldn't have been") + + runCmd(t, tempDir, "uv", "pip", "install", "-q", patchedWheel) + verifyVersion(t, tempDir, patchedWheel) + + newTime := patchedInfo.ModTime().Add(10 * time.Millisecond) + + err = os.Chtimes(origWheel, newTime, newTime) + require.NoError(t, err) + + patchedWheel3, err := PatchWheel(context.Background(), origWheel, distDir) + require.NoError(t, err) + require.Greater(t, patchedWheel3, patchedWheel) + + // Now use regular pip to re-install the wheel. First install pip. + runCmd(t, tempDir, "uv", "pip", "install", "-q", "pip") + + pippath := filepath.Join(".venv", getPythonScriptsDir(), "pip") + runCmd(t, tempDir, pippath, "install", "-q", patchedWheel3) + verifyVersion(t, tempDir, patchedWheel3) + }) + } +} + +func errPatchWheel(t *testing.T, name, out string) { + ctx := context.Background() + outname, err := PatchWheel(ctx, name, out) + assert.Error(t, err, "PatchWheel(%s, %s) expected to error", name, out) + assert.Empty(t, outname) +} + +func TestError(t *testing.T) { + // empty name and dir + errPatchWheel(t, "", "") + + // empty name + errPatchWheel(t, "", ".") + + // file not found + errPatchWheel(t, "not-found.txt", ".") + + // output directory not found + errPatchWheel(t, prebuiltWheel, "not-found/a/b/c") +} + +func TestEmptyZip(t *testing.T) { + tempDir := t.TempDir() + errPatchWheel(t, emptyZip, tempDir) +} + +func TestNonZip(t *testing.T) { + tempDir := t.TempDir() + + _, err := os.Stat("patch.go") + require.NoError(t, err, "file must exist for this test") + errPatchWheel(t, "patch.go", tempDir) +} diff --git a/libs/patchwheel/pkg.go b/libs/patchwheel/pkg.go new file mode 100644 index 000000000..97349a588 --- /dev/null +++ b/libs/patchwheel/pkg.go @@ -0,0 +1,22 @@ +package patchwheel + +/* + +Patching whl file with a dynamic version suffix. + +When developing a DAB, users want to redeploy a wheel without updating a version in pyproject.toml / setup.py manually. + +However, installing the same version with pip causes pip to skip the install. Databricks envs follow this behaviour. + +For this reason, we've modified default-python template to auto-update the version https://github.com/databricks/cli/pull/1034 + +However, that makes it tied to setup.py / setuptools and puts onus on users to keep this behaviour. + +This package removes the constraint on how the wheel is built and allows adding dynamic part as a post-build step. + +PatchWheel(ctx, path, outputDir) takes existing whl file and creates a new patched one with a version that includes + mtime of the original wheel as a suffix. + METADATA, directory names, RECORD are all updated to ensure the correct format. + +ParseWheelFilename(filename) extracts version from the filename, according to WHL format rules. +*/ diff --git a/libs/patchwheel/testdata/empty.zip b/libs/patchwheel/testdata/empty.zip new file mode 100644 index 000000000..15cb0ecb3 Binary files /dev/null and b/libs/patchwheel/testdata/empty.zip differ diff --git a/libs/template/renderer_test.go b/libs/template/renderer_test.go index f9306ca31..41439846b 100644 --- a/libs/template/renderer_test.go +++ b/libs/template/renderer_test.go @@ -43,7 +43,7 @@ func init() { } func assertBuiltinTemplateValid(t *testing.T, template string, settings map[string]any, target string, isServicePrincipal, build bool, tempDir string) { - ctx := dbr.MockRuntime(context.Background(), false) + ctx := dbr.MockRuntime(context.Background(), dbr.Environment{}) templateFS, err := fs.Sub(builtinTemplates, path.Join("templates", template)) require.NoError(t, err) diff --git a/libs/template/writer_test.go b/libs/template/writer_test.go index 9d57966ee..a3e10d4fe 100644 --- a/libs/template/writer_test.go +++ b/libs/template/writer_test.go @@ -31,7 +31,7 @@ func TestDefaultWriterConfigureOnDBR(t *testing.T) { t.Skip("Skipping test on Windows") } - ctx := dbr.MockRuntime(context.Background(), true) + ctx := dbr.MockRuntime(context.Background(), dbr.Environment{IsDbr: true, Version: "15.4"}) ctx = root.SetWorkspaceClient(ctx, &databricks.WorkspaceClient{ Config: &workspaceConfig.Config{Host: "https://myhost.com"}, })