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

@@ -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 {
+			rewriteArtifactPath(b, f, job, remotePath)
-				task := &job.Tasks[i]
+
-				for j := range task.Libraries {
+		}
-					lib := &task.Libraries[j]
+	}
-					if lib.Whl != "" && isArtifactMatchLibrary(f, lib.Whl, b) {
+
-						lib.Whl = remotePath
+	return nil
-					}
+}
-					if lib.Jar != "" && isArtifactMatchLibrary(f, lib.Jar, b) {
+
-						lib.Jar = remotePath
+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 {
+		// Rewrite artifact path in job task libraries for ForEachTask
-				env := &job.Environments[i]
+		if task.ForEachTask != nil {
-				if env.Spec == nil {
+			forEachTask := task.ForEachTask
-					continue
+			for j := range forEachTask.Task.Libraries {
+				lib := &forEachTask.Task.Libraries[j]
+				if lib.Whl != "" && isArtifactMatchLibrary(f, lib.Whl, b) {
+					lib.Whl = remotePath
 				}
-
+				if lib.Jar != "" && isArtifactMatchLibrary(f, lib.Jar, b) {
-				for j := range env.Spec.Dependencies {
+					lib.Jar = remotePath
-					lib := env.Spec.Dependencies[j]
-					if isArtifactMatchLibrary(f, lib, b) {
-						env.Spec.Dependencies[j] = 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 {

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)
 }