From 77ff994d1ba2cc3c3c1a7dd055d063391e89b07e Mon Sep 17 00:00:00 2001
From: Andrew Nester <andrew.nester@databricks.com>
Date: Fri, 5 Apr 2024 17:52:39 +0200
Subject: [PATCH] Correctly transform libraries in for_each_task block (#1340)

## Changes
Now DABs correctly transforms and deploys libraries in for_each_task
block

```
tasks:
  - task_key: my_loop
     for_each_task:
       inputs: "[1,2,3]"
         task:
           task_key: my_loop_iteration
           libraries:
             - pypi:
                  package: my_package
```

## Tests
Added regression test
---
 bundle/deploy/terraform/convert.go           | 10 +++++
 bundle/deploy/terraform/convert_test.go      | 44 ++++++++++++++++++++
 bundle/deploy/terraform/tfdyn/convert_job.go | 13 +++++-
 3 files changed, 66 insertions(+), 1 deletion(-)

diff --git a/bundle/deploy/terraform/convert.go b/bundle/deploy/terraform/convert.go
index f2fb77e1..0ae6751d 100644
--- a/bundle/deploy/terraform/convert.go
+++ b/bundle/deploy/terraform/convert.go
@@ -100,6 +100,16 @@ func BundleToTerraform(config *config.Root) *schema.Root {
 					t.Library = append(t.Library, l)
 				}
 
+				// Convert for_each_task libraries
+				if v.ForEachTask != nil {
+					for _, v_ := range v.ForEachTask.Task.Libraries {
+						var l schema.ResourceJobTaskForEachTaskTaskLibrary
+						conv(v_, &l)
+						t.ForEachTask.Task.Library = append(t.ForEachTask.Task.Library, l)
+					}
+
+				}
+
 				dst.Task = append(dst.Task, t)
 			}
 
diff --git a/bundle/deploy/terraform/convert_test.go b/bundle/deploy/terraform/convert_test.go
index 9621a56a..986599a7 100644
--- a/bundle/deploy/terraform/convert_test.go
+++ b/bundle/deploy/terraform/convert_test.go
@@ -140,6 +140,50 @@ func TestBundleToTerraformJobTaskLibraries(t *testing.T) {
 	bundleToTerraformEquivalenceTest(t, &config)
 }
 
+func TestBundleToTerraformForEachTaskLibraries(t *testing.T) {
+	var src = resources.Job{
+		JobSettings: &jobs.JobSettings{
+			Name: "my job",
+			Tasks: []jobs.Task{
+				{
+					TaskKey: "key",
+					ForEachTask: &jobs.ForEachTask{
+						Inputs: "[1,2,3]",
+						Task: jobs.Task{
+							TaskKey: "iteration",
+							Libraries: []compute.Library{
+								{
+									Pypi: &compute.PythonPyPiLibrary{
+										Package: "mlflow",
+									},
+								},
+							},
+						},
+					},
+				},
+			},
+		},
+	}
+
+	var config = config.Root{
+		Resources: config.Resources{
+			Jobs: map[string]*resources.Job{
+				"my_job": &src,
+			},
+		},
+	}
+
+	out := BundleToTerraform(&config)
+	resource := out.Resource.Job["my_job"].(*schema.ResourceJob)
+
+	assert.Equal(t, "my job", resource.Name)
+	require.Len(t, resource.Task, 1)
+	require.Len(t, resource.Task[0].ForEachTask.Task.Library, 1)
+	assert.Equal(t, "mlflow", resource.Task[0].ForEachTask.Task.Library[0].Pypi.Package)
+
+	bundleToTerraformEquivalenceTest(t, &config)
+}
+
 func TestBundleToTerraformPipeline(t *testing.T) {
 	var src = resources.Pipeline{
 		PipelineSpec: &pipelines.PipelineSpec{
diff --git a/bundle/deploy/terraform/tfdyn/convert_job.go b/bundle/deploy/terraform/tfdyn/convert_job.go
index 778af1ad..65ac8b9b 100644
--- a/bundle/deploy/terraform/tfdyn/convert_job.go
+++ b/bundle/deploy/terraform/tfdyn/convert_job.go
@@ -45,7 +45,18 @@ func convertJobResource(ctx context.Context, vin dyn.Value) (dyn.Value, error) {
 
 	// Modify keys in the "task" blocks
 	vout, err = dyn.Map(vout, "task", dyn.Foreach(func(_ dyn.Path, v dyn.Value) (dyn.Value, error) {
-		return renameKeys(v, map[string]string{
+		// Modify "library" blocks for for_each_task
+		vout, err = dyn.Map(v, "for_each_task.task", func(_ dyn.Path, v dyn.Value) (dyn.Value, error) {
+			return renameKeys(v, map[string]string{
+				"libraries": "library",
+			})
+		})
+
+		if err != nil {
+			return dyn.InvalidValue, err
+		}
+
+		return renameKeys(vout, map[string]string{
 			"libraries": "library",
 		})
 	}))