From 86c30dd3289751b30a232b415fcc2b4d76232187 Mon Sep 17 00:00:00 2001
From: Andrew Nester <andrew.nester@databricks.com>
Date: Thu, 31 Aug 2023 16:10:32 +0200
Subject: [PATCH] Fixed artifact file uploading on Windows and wheel execution
 on DBR 13.3 (#722)

## Changes
Fixed artifact file uploading on Windows and wheel execution on DBR 13.3

Fixes #719, #720

## Tests
Added regression test for Windows
---
 bundle/artifacts/artifacts.go      |  5 +-
 bundle/artifacts/artifacts_test.go | 89 ++++++++++++++++++++++++++++++
 bundle/python/transform.go         |  2 +
 3 files changed, 94 insertions(+), 2 deletions(-)
 create mode 100644 bundle/artifacts/artifacts_test.go

diff --git a/bundle/artifacts/artifacts.go b/bundle/artifacts/artifacts.go
index c5413121..0331adb7 100644
--- a/bundle/artifacts/artifacts.go
+++ b/bundle/artifacts/artifacts.go
@@ -8,6 +8,7 @@ import (
 	"fmt"
 	"os"
 	"path"
+	"path/filepath"
 
 	"github.com/databricks/cli/bundle"
 	"github.com/databricks/cli/bundle/artifacts/whl"
@@ -107,7 +108,7 @@ func uploadArtifact(ctx context.Context, a *config.Artifact, b *bundle.Bundle) e
 	for i := range a.Files {
 		f := &a.Files[i]
 		if f.NeedsUpload() {
-			filename := path.Base(f.Source)
+			filename := filepath.Base(f.Source)
 			cmdio.LogString(ctx, fmt.Sprintf("artifacts.Upload(%s): Uploading...", filename))
 			remotePath, err := uploadArtifactFile(ctx, f.Source, b)
 			if err != nil {
@@ -136,7 +137,7 @@ func uploadArtifactFile(ctx context.Context, file string, b *bundle.Bundle) (str
 	}
 
 	fileHash := sha256.Sum256(raw)
-	remotePath := path.Join(uploadPath, fmt.Sprintf("%x", fileHash), path.Base(file))
+	remotePath := path.Join(uploadPath, fmt.Sprintf("%x", fileHash), filepath.Base(file))
 	// Make sure target directory exists.
 	err = b.WorkspaceClient().Workspace.MkdirsByPath(ctx, path.Dir(remotePath))
 	if err != nil {
diff --git a/bundle/artifacts/artifacts_test.go b/bundle/artifacts/artifacts_test.go
new file mode 100644
index 00000000..65a1950a
--- /dev/null
+++ b/bundle/artifacts/artifacts_test.go
@@ -0,0 +1,89 @@
+package artifacts
+
+import (
+	"context"
+	"os"
+	"path/filepath"
+	"regexp"
+	"testing"
+
+	"github.com/databricks/cli/bundle"
+	"github.com/databricks/cli/bundle/config"
+	"github.com/databricks/databricks-sdk-go/service/compute"
+	"github.com/databricks/databricks-sdk-go/service/workspace"
+	"github.com/stretchr/testify/require"
+)
+
+func touchEmptyFile(t *testing.T, path string) {
+	err := os.MkdirAll(filepath.Dir(path), 0700)
+	require.NoError(t, err)
+	f, err := os.Create(path)
+	require.NoError(t, err)
+	f.Close()
+}
+
+type MockWorkspaceService struct {
+}
+
+// Delete implements workspace.WorkspaceService.
+func (MockWorkspaceService) Delete(ctx context.Context, request workspace.Delete) error {
+	panic("unimplemented")
+}
+
+// Export implements workspace.WorkspaceService.
+func (MockWorkspaceService) Export(ctx context.Context, request workspace.ExportRequest) (*workspace.ExportResponse, error) {
+	panic("unimplemented")
+}
+
+// GetStatus implements workspace.WorkspaceService.
+func (MockWorkspaceService) GetStatus(ctx context.Context, request workspace.GetStatusRequest) (*workspace.ObjectInfo, error) {
+	panic("unimplemented")
+}
+
+// Import implements workspace.WorkspaceService.
+func (MockWorkspaceService) Import(ctx context.Context, request workspace.Import) error {
+	return nil
+}
+
+// List implements workspace.WorkspaceService.
+func (MockWorkspaceService) List(ctx context.Context, request workspace.ListWorkspaceRequest) (*workspace.ListResponse, error) {
+	panic("unimplemented")
+}
+
+// Mkdirs implements workspace.WorkspaceService.
+func (MockWorkspaceService) Mkdirs(ctx context.Context, request workspace.Mkdirs) error {
+	return nil
+}
+
+func TestUploadArtifactFileToCorrectRemotePath(t *testing.T) {
+	dir := t.TempDir()
+	whlPath := filepath.Join(dir, "dist", "test.whl")
+	touchEmptyFile(t, whlPath)
+	b := &bundle.Bundle{
+		Config: config.Root{
+			Path: dir,
+			Bundle: config.Bundle{
+				Target: "whatever",
+			},
+			Workspace: config.Workspace{
+				ArtifactsPath: "/Users/test@databricks.com/whatever",
+			},
+		},
+	}
+
+	b.WorkspaceClient().Workspace.WithImpl(MockWorkspaceService{})
+	artifact := &config.Artifact{
+		Files: []config.ArtifactFile{
+			{
+				Source: whlPath,
+				Libraries: []*compute.Library{
+					{Whl: "dist\\test.whl"},
+				},
+			},
+		},
+	}
+
+	err := uploadArtifact(context.Background(), artifact, b)
+	require.NoError(t, err)
+	require.Regexp(t, regexp.MustCompile("/Users/test@databricks.com/whatever/.internal/[a-z0-9]+/test.whl"), artifact.Files[0].RemotePath)
+}
diff --git a/bundle/python/transform.go b/bundle/python/transform.go
index 6ec75a03..53db450b 100644
--- a/bundle/python/transform.go
+++ b/bundle/python/transform.go
@@ -16,6 +16,8 @@ const NOTEBOOK_TEMPLATE = `# Databricks notebook source
 %pip install --force-reinstall {{.Whl}}
 {{end}}
 
+dbutils.library.restartPython()
+
 try:
 	from importlib import metadata
 except ImportError: # for Python<3.8