From 18a5b05d82ce589c4bfb26ea50c8b99d52fe3f72 Mon Sep 17 00:00:00 2001
From: Andrew Nester <andrew.nester@databricks.com>
Date: Fri, 8 Sep 2023 15:45:21 +0200
Subject: [PATCH] Apply Python wheel trampoline if workspace library is used
 (#755)

## Changes
Workspace library will be detected by trampoline in 2 cases:
- User defined to use local wheel file
- User defined to use remote wheel file from Workspace file system

In both of these cases we should correctly apply Python trampoline

## Tests
Added a regression test (also covered by Python e2e test)
---
 bundle/artifacts/artifacts_test.go |  2 ++
 bundle/config/artifact.go          |  8 ++++++--
 bundle/libraries/libraries.go      | 11 +++++++++++
 bundle/python/transform.go         | 11 +++++++++--
 bundle/python/transform_test.go    |  2 +-
 5 files changed, 29 insertions(+), 5 deletions(-)

diff --git a/bundle/artifacts/artifacts_test.go b/bundle/artifacts/artifacts_test.go
index 4c0a18f3..bbae44ef 100644
--- a/bundle/artifacts/artifacts_test.go
+++ b/bundle/artifacts/artifacts_test.go
@@ -105,6 +105,7 @@ func TestUploadArtifactFileToCorrectRemotePath(t *testing.T) {
 
 	b.WorkspaceClient().Workspace.WithImpl(MockWorkspaceService{})
 	artifact := &config.Artifact{
+		Type: "whl",
 		Files: []config.ArtifactFile{
 			{
 				Source: whlPath,
@@ -118,4 +119,5 @@ func TestUploadArtifactFileToCorrectRemotePath(t *testing.T) {
 	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)
+	require.Regexp(t, regexp.MustCompile("/Workspace/Users/test@databricks.com/whatever/.internal/[a-z0-9]+/test.whl"), artifact.Files[0].Libraries[0].Whl)
 }
diff --git a/bundle/config/artifact.go b/bundle/config/artifact.go
index 1955e265..d7048a02 100644
--- a/bundle/config/artifact.go
+++ b/bundle/config/artifact.go
@@ -78,9 +78,13 @@ func (a *Artifact) NormalisePaths() {
 		remotePath := path.Join(wsfsBase, f.RemotePath)
 		for i := range f.Libraries {
 			lib := f.Libraries[i]
-			switch a.Type {
-			case ArtifactPythonWheel:
+			if lib.Whl != "" {
 				lib.Whl = remotePath
+				continue
+			}
+			if lib.Jar != "" {
+				lib.Jar = remotePath
+				continue
 			}
 		}
 
diff --git a/bundle/libraries/libraries.go b/bundle/libraries/libraries.go
index d9a257db..8e2e504c 100644
--- a/bundle/libraries/libraries.go
+++ b/bundle/libraries/libraries.go
@@ -78,6 +78,17 @@ func IsTaskWithLocalLibraries(task *jobs.Task) bool {
 	return false
 }
 
+func IsTaskWithWorkspaceLibraries(task *jobs.Task) bool {
+	for _, l := range task.Libraries {
+		path := libPath(&l)
+		if isWorkspacePath(path) {
+			return true
+		}
+	}
+
+	return false
+}
+
 func isMissingRequiredLibraries(task *jobs.Task) bool {
 	if task.Libraries != nil {
 		return false
diff --git a/bundle/python/transform.go b/bundle/python/transform.go
index 3d744df9..d8eb33f5 100644
--- a/bundle/python/transform.go
+++ b/bundle/python/transform.go
@@ -73,8 +73,11 @@ func (t *pythonTrampoline) GetTasks(b *bundle.Bundle) []mutator.TaskWithJobKey {
 		for i := range tasks {
 			task := &tasks[i]
 
-			// Keep only Python wheel tasks with local libraries referenced
-			if task.PythonWheelTask == nil || !libraries.IsTaskWithLocalLibraries(task) {
+			// Keep only Python wheel tasks with workspace libraries referenced.
+			// At this point of moment we don't have local paths in Libraries sections anymore
+			// Local paths have been replaced with the remote when the artifacts where uploaded
+			// in artifacts.UploadAll mutator.
+			if task.PythonWheelTask == nil || !needsTrampoline(task) {
 				continue
 			}
 
@@ -87,6 +90,10 @@ func (t *pythonTrampoline) GetTasks(b *bundle.Bundle) []mutator.TaskWithJobKey {
 	return result
 }
 
+func needsTrampoline(task *jobs.Task) bool {
+	return libraries.IsTaskWithWorkspaceLibraries(task)
+}
+
 func (t *pythonTrampoline) GetTemplateData(task *jobs.Task) (map[string]any, error) {
 	params, err := t.generateParameters(task.PythonWheelTask)
 	if err != nil {
diff --git a/bundle/python/transform_test.go b/bundle/python/transform_test.go
index 99d3129d..a7448f23 100644
--- a/bundle/python/transform_test.go
+++ b/bundle/python/transform_test.go
@@ -84,7 +84,7 @@ func TestTransformFiltersWheelTasksOnly(t *testing.T) {
 									TaskKey:         "key1",
 									PythonWheelTask: &jobs.PythonWheelTask{},
 									Libraries: []compute.Library{
-										{Whl: "./dist/test.whl"},
+										{Whl: "/Workspace/Users/test@test.com/bundle/dist/test.whl"},
 									},
 								},
 								{