From 3d8446bbdbb1797a87879af11cb5eb7d4a041aa6 Mon Sep 17 00:00:00 2001
From: Andrew Nester <andrew.nester@databricks.com>
Date: Fri, 5 Jul 2024 12:58:28 +0200
Subject: [PATCH] Rewrite local path for libraries in foreach tasks (#1569)

## Changes
Now local library path in `libraries` section of foreach each tasks are
correctly replaced with remote path for this library when it's uploaded
to Databricks

## Tests
Added unit test
---
 bundle/artifacts/artifacts.go      | 66 ++++++++++++++++++++----------
 bundle/artifacts/artifacts_test.go | 16 ++++++++
 2 files changed, 60 insertions(+), 22 deletions(-)

diff --git a/bundle/artifacts/artifacts.go b/bundle/artifacts/artifacts.go
index 470c329a..a5f41ae4 100644
--- a/bundle/artifacts/artifacts.go
+++ b/bundle/artifacts/artifacts.go
@@ -12,6 +12,7 @@ import (
 	"github.com/databricks/cli/bundle"
 	"github.com/databricks/cli/bundle/artifacts/whl"
 	"github.com/databricks/cli/bundle/config"
+	"github.com/databricks/cli/bundle/config/resources"
 	"github.com/databricks/cli/libs/cmdio"
 	"github.com/databricks/cli/libs/diag"
 	"github.com/databricks/cli/libs/filer"
@@ -135,36 +136,57 @@ func uploadArtifact(ctx context.Context, b *bundle.Bundle, a *config.Artifact, u
 		remotePath := path.Join(wsfsBase, f.RemotePath)
 
 		for _, job := range b.Config.Resources.Jobs {
-			for i := range job.Tasks {
-				task := &job.Tasks[i]
-				for j := range task.Libraries {
-					lib := &task.Libraries[j]
-					if lib.Whl != "" && isArtifactMatchLibrary(f, lib.Whl, b) {
-						lib.Whl = remotePath
-					}
-					if lib.Jar != "" && isArtifactMatchLibrary(f, lib.Jar, b) {
-						lib.Jar = remotePath
-					}
-				}
+			rewriteArtifactPath(b, f, job, remotePath)
+
+		}
+	}
+
+	return nil
+}
+
+func rewriteArtifactPath(b *bundle.Bundle, f *config.ArtifactFile, job *resources.Job, remotePath string) {
+	// Rewrite artifact path in job task libraries
+	for i := range job.Tasks {
+		task := &job.Tasks[i]
+		for j := range task.Libraries {
+			lib := &task.Libraries[j]
+			if lib.Whl != "" && isArtifactMatchLibrary(f, lib.Whl, b) {
+				lib.Whl = remotePath
 			}
+			if lib.Jar != "" && isArtifactMatchLibrary(f, lib.Jar, b) {
+				lib.Jar = remotePath
+			}
+		}
 
-			for i := range job.Environments {
-				env := &job.Environments[i]
-				if env.Spec == nil {
-					continue
+		// Rewrite artifact path in job task libraries for ForEachTask
+		if task.ForEachTask != nil {
+			forEachTask := task.ForEachTask
+			for j := range forEachTask.Task.Libraries {
+				lib := &forEachTask.Task.Libraries[j]
+				if lib.Whl != "" && isArtifactMatchLibrary(f, lib.Whl, b) {
+					lib.Whl = remotePath
 				}
-
-				for j := range env.Spec.Dependencies {
-					lib := env.Spec.Dependencies[j]
-					if isArtifactMatchLibrary(f, lib, b) {
-						env.Spec.Dependencies[j] = remotePath
-					}
+				if lib.Jar != "" && isArtifactMatchLibrary(f, lib.Jar, b) {
+					lib.Jar = remotePath
 				}
 			}
 		}
 	}
 
-	return nil
+	// Rewrite artifact path in job environments
+	for i := range job.Environments {
+		env := &job.Environments[i]
+		if env.Spec == nil {
+			continue
+		}
+
+		for j := range env.Spec.Dependencies {
+			lib := env.Spec.Dependencies[j]
+			if isArtifactMatchLibrary(f, lib, b) {
+				env.Spec.Dependencies[j] = remotePath
+			}
+		}
+	}
 }
 
 func isArtifactMatchLibrary(f *config.ArtifactFile, libPath string, b *bundle.Bundle) bool {
diff --git a/bundle/artifacts/artifacts_test.go b/bundle/artifacts/artifacts_test.go
index ca0e578b..53c2798e 100644
--- a/bundle/artifacts/artifacts_test.go
+++ b/bundle/artifacts/artifacts_test.go
@@ -52,6 +52,20 @@ func TestArtifactUpload(t *testing.T) {
 										},
 									},
 								},
+								{
+									ForEachTask: &jobs.ForEachTask{
+										Task: jobs.Task{
+											Libraries: []compute.Library{
+												{
+													Whl: filepath.Join("whl", "*.whl"),
+												},
+												{
+													Whl: "/Workspace/Users/foo@bar.com/mywheel.whl",
+												},
+											},
+										},
+									},
+								},
 							},
 							Environments: []jobs.JobEnvironment{
 								{
@@ -88,4 +102,6 @@ func TestArtifactUpload(t *testing.T) {
 	require.Equal(t, "/Workspace/Users/foo@bar.com/mywheel.whl", b.Config.Resources.Jobs["job"].JobSettings.Tasks[0].Libraries[1].Whl)
 	require.Equal(t, "/Workspace/foo/bar/artifacts/source.whl", b.Config.Resources.Jobs["job"].JobSettings.Environments[0].Spec.Dependencies[0])
 	require.Equal(t, "/Workspace/Users/foo@bar.com/mywheel.whl", b.Config.Resources.Jobs["job"].JobSettings.Environments[0].Spec.Dependencies[1])
+	require.Equal(t, "/Workspace/foo/bar/artifacts/source.whl", b.Config.Resources.Jobs["job"].JobSettings.Tasks[1].ForEachTask.Task.Libraries[0].Whl)
+	require.Equal(t, "/Workspace/Users/foo@bar.com/mywheel.whl", b.Config.Resources.Jobs["job"].JobSettings.Tasks[1].ForEachTask.Task.Libraries[1].Whl)
 }