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..454a19bf2 --- /dev/null +++ b/acceptance/cmd/patchwhl/output.txt @@ -0,0 +1,6 @@ + +=== Test prebuilt wheel: +>>> setmtime.py 2025-03-05 15:07:33.123456789 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+2025030514073312-py3-none-any.whl diff --git a/acceptance/cmd/patchwhl/output/my_test_code-0.0.1+2025030514073312.dist-info/METADATA b/acceptance/cmd/patchwhl/output/my_test_code-0.0.1+2025030514073312.dist-info/METADATA new file mode 100644 index 000000000..c1ba293e1 --- /dev/null +++ b/acceptance/cmd/patchwhl/output/my_test_code-0.0.1+2025030514073312.dist-info/METADATA @@ -0,0 +1,9 @@ +Metadata-Version: 2.1 +Name: my-test-code +Version: 0.0.1+2025030514073312 +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+2025030514073312.dist-info/RECORD b/acceptance/cmd/patchwhl/output/my_test_code-0.0.1+2025030514073312.dist-info/RECORD new file mode 100644 index 000000000..2f6e7e51b --- /dev/null +++ b/acceptance/cmd/patchwhl/output/my_test_code-0.0.1+2025030514073312.dist-info/RECORD @@ -0,0 +1,8 @@ +src/__init__.py,sha256=BRmKeYehopKv4NG_SFa7t6wn248RrPHJivu7DM1R-Rw,48 +src/__main__.py,sha256=8TtsnLsaJEM35Y4L8ocrv-qfxusgYpRL2HPyYiabHng,242 +my_test_code-0.0.1+2025030514073312.dist-info/METADATA,sha256=nidQMSt6OxDHKdNCQGq1Kv_AmJa1ldwCMHRqPk2TFD8,214 +my_test_code-0.0.1+2025030514073312.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92 +my_test_code-0.0.1+2025030514073312.dist-info/entry_points.txt,sha256=oDWOW9SsBlk4Uejj1ftYPBxfhJ5ZJctb4JOUIG1rc-4,34 +my_test_code-0.0.1+2025030514073312.dist-info/top_level.txt,sha256=74rtVfumQlgAPzR5_2CgYN24MB0XARCg0t-gzk6gTrM,4 +my_test_code-0.0.1+2025030514073312.dist-info/RECORD,, + diff --git a/acceptance/cmd/patchwhl/output/my_test_code-0.0.1+2025030514073312.dist-info/WHEEL b/acceptance/cmd/patchwhl/output/my_test_code-0.0.1+2025030514073312.dist-info/WHEEL new file mode 100644 index 000000000..98c0d20b7 --- /dev/null +++ b/acceptance/cmd/patchwhl/output/my_test_code-0.0.1+2025030514073312.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+2025030514073312.dist-info/entry_points.txt b/acceptance/cmd/patchwhl/output/my_test_code-0.0.1+2025030514073312.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+2025030514073312.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+2025030514073312.dist-info/top_level.txt b/acceptance/cmd/patchwhl/output/my_test_code-0.0.1+2025030514073312.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+2025030514073312.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/script b/acceptance/cmd/patchwhl/script new file mode 100644 index 000000000..b9e0dd5d4 --- /dev/null +++ b/acceptance/cmd/patchwhl/script @@ -0,0 +1,9 @@ +title "Test prebuilt wheel:" + +trace setmtime.py "2025-03-05 15:07:33.123456789" 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 +cd output +unzip -q ../my_test_code-0.0.1+2025030514073312-py3-none-any.whl +rm ../my_test_code-0.0.1+2025030514073312-py3-none-any.whl 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 7d8cfcb76..5e2a2c30e 100644 --- a/cmd/selftest/selftest.go +++ b/cmd/selftest/selftest.go @@ -12,5 +12,6 @@ func New() *cobra.Command { } cmd.AddCommand(newPanic()) + cmd.AddCommand(newPatchWhl()) return cmd } diff --git a/libs/patchwheel/patch.go b/libs/patchwheel/patch.go index 956871eb6..d0967becd 100644 --- a/libs/patchwheel/patch.go +++ b/libs/patchwheel/patch.go @@ -54,7 +54,7 @@ func patchMetadata(r io.Reader, oldVersion, newVersion string) ([]byte, error) { if foundVersion != oldVersion { return nil, fmt.Errorf("Unexpected version in METADATA: %s (expected %s)", strings.TrimSpace(string(line)), oldVersion) } - buf.WriteString(string(versionKey) + newVersion) + buf.WriteString(string(versionKey) + " " + newVersion + "\n") } else { buf.Write(line) buf.WriteString("\n") diff --git a/libs/patchwheel/patch_test.go b/libs/patchwheel/patch_test.go index 5fe27e565..822c3d643 100644 --- a/libs/patchwheel/patch_test.go +++ b/libs/patchwheel/patch_test.go @@ -1,10 +1,8 @@ package patchwheel import ( - "archive/zip" "bytes" "context" - "io" "os" "os/exec" "path/filepath" @@ -182,63 +180,6 @@ func TestPatchWheel(t *testing.T) { } } -func TestPrebuilt(t *testing.T) { - tempDir := t.TempDir() - ctx := context.Background() - - // Set fixed modification time for deterministic testing - fixedTime := time.Date(2025, 3, 5, 14, 15, 55, 123456789, time.UTC) - err := os.Chtimes(prebuiltWheel, fixedTime, fixedTime) - require.NoError(t, err) - - // With the fixed time, we know exactly what the output filename will be - expectedVersion := "0.0.1+20250305141555.12" - expectedFilename := "my_test_code-" + expectedVersion + "-py3-none-any.whl" - expectedPath := filepath.Join(tempDir, expectedFilename) - - outname, err := PatchWheel(ctx, prebuiltWheel, tempDir) - require.NoError(t, err) - require.Equal(t, expectedPath, outname) - - _, err = os.Stat(outname) - require.NoError(t, err) - - // Verify the contents of the patched wheel - archive, err := zip.OpenReader(outname) - require.NoError(t, err) - defer archive.Close() - - // With fixed time, we know the exact dist-info directory name - distInfoPrefix := "my_test_code-" + expectedVersion + ".dist-info/" - - // Find METADATA and RECORD files - var metadataContent, recordContent []byte - for _, f := range archive.File { - if f.Name == distInfoPrefix+"METADATA" { - rc, err := f.Open() - require.NoError(t, err) - metadataContent, err = io.ReadAll(rc) - rc.Close() - require.NoError(t, err) - } else if f.Name == distInfoPrefix+"RECORD" { - rc, err := f.Open() - require.NoError(t, err) - recordContent, err = io.ReadAll(rc) - rc.Close() - require.NoError(t, err) - } - } - - // Verify METADATA contains the expected version - require.NotNil(t, metadataContent, "METADATA file not found in wheel") - assert.Contains(t, string(metadataContent), "Version: "+expectedVersion) - - // Verify RECORD contains entries with the correct dist-info prefix - require.NotNil(t, recordContent, "RECORD file not found in wheel") - assert.Contains(t, string(recordContent), distInfoPrefix+"METADATA") - assert.Contains(t, string(recordContent), distInfoPrefix+"RECORD") -} - func errPatchWheel(t *testing.T, name, out string) { ctx := context.Background() outname, err := PatchWheel(ctx, name, out) diff --git a/libs/patchwheel/testdata/my_test_code-0.0.1-py3-none-any.whl b/libs/patchwheel/testdata/my_test_code-0.0.1-py3-none-any.whl deleted file mode 100644 index 14702281d..000000000 Binary files a/libs/patchwheel/testdata/my_test_code-0.0.1-py3-none-any.whl and /dev/null differ