From 192f33bb13a156bebf9d7d2c2b06092d8ae9775d Mon Sep 17 00:00:00 2001
From: Witold Czaplewski <witi83@users.noreply.github.com>
Date: Wed, 21 Aug 2024 12:03:56 +0200
Subject: [PATCH] [DAB] Add support for requirements libraries in Job Tasks
 (#1543)

## Changes
While experimenting with DAB I discovered that requirements libraries
are being ignored.

One thing worth mentioning is that `bundle validate` runs successfully,
but `bundle deploy` fails. This PR only covers the second part.


## Tests
<!-- How is this tested? -->
Added a unit test
---
 bundle/config/mutator/translate_paths_jobs.go | 5 +++++
 bundle/config/mutator/translate_paths_test.go | 9 +++++++++
 bundle/libraries/helpers.go                   | 3 +++
 bundle/libraries/helpers_test.go              | 1 +
 4 files changed, 18 insertions(+)

diff --git a/bundle/config/mutator/translate_paths_jobs.go b/bundle/config/mutator/translate_paths_jobs.go
index 6febf4f8..e34eeb2f 100644
--- a/bundle/config/mutator/translate_paths_jobs.go
+++ b/bundle/config/mutator/translate_paths_jobs.go
@@ -50,6 +50,11 @@ func rewritePatterns(t *translateContext, base dyn.Pattern) []jobRewritePattern
 			t.translateNoOp,
 			noSkipRewrite,
 		},
+		{
+			base.Append(dyn.Key("libraries"), dyn.AnyIndex(), dyn.Key("requirements")),
+			t.translateFilePath,
+			noSkipRewrite,
+		},
 	}
 }
 
diff --git a/bundle/config/mutator/translate_paths_test.go b/bundle/config/mutator/translate_paths_test.go
index 780a540d..fd64593b 100644
--- a/bundle/config/mutator/translate_paths_test.go
+++ b/bundle/config/mutator/translate_paths_test.go
@@ -110,6 +110,7 @@ func TestTranslatePaths(t *testing.T) {
 	touchNotebookFile(t, filepath.Join(dir, "my_pipeline_notebook.py"))
 	touchEmptyFile(t, filepath.Join(dir, "my_python_file.py"))
 	touchEmptyFile(t, filepath.Join(dir, "dist", "task.jar"))
+	touchEmptyFile(t, filepath.Join(dir, "requirements.txt"))
 
 	b := &bundle.Bundle{
 		RootPath:   dir,
@@ -140,6 +141,9 @@ func TestTranslatePaths(t *testing.T) {
 									NotebookTask: &jobs.NotebookTask{
 										NotebookPath: "./my_job_notebook.py",
 									},
+									Libraries: []compute.Library{
+										{Requirements: "./requirements.txt"},
+									},
 								},
 								{
 									PythonWheelTask: &jobs.PythonWheelTask{
@@ -232,6 +236,11 @@ func TestTranslatePaths(t *testing.T) {
 		"/bundle/my_job_notebook",
 		b.Config.Resources.Jobs["job"].Tasks[2].NotebookTask.NotebookPath,
 	)
+	assert.Equal(
+		t,
+		"/bundle/requirements.txt",
+		b.Config.Resources.Jobs["job"].Tasks[2].Libraries[0].Requirements,
+	)
 	assert.Equal(
 		t,
 		"/bundle/my_python_file.py",
diff --git a/bundle/libraries/helpers.go b/bundle/libraries/helpers.go
index 89679c91..b7e707cc 100644
--- a/bundle/libraries/helpers.go
+++ b/bundle/libraries/helpers.go
@@ -12,5 +12,8 @@ func libraryPath(library *compute.Library) string {
 	if library.Egg != "" {
 		return library.Egg
 	}
+	if library.Requirements != "" {
+		return library.Requirements
+	}
 	return ""
 }
diff --git a/bundle/libraries/helpers_test.go b/bundle/libraries/helpers_test.go
index adc20a24..e4bd3277 100644
--- a/bundle/libraries/helpers_test.go
+++ b/bundle/libraries/helpers_test.go
@@ -13,5 +13,6 @@ func TestLibraryPath(t *testing.T) {
 	assert.Equal(t, path, libraryPath(&compute.Library{Whl: path}))
 	assert.Equal(t, path, libraryPath(&compute.Library{Jar: path}))
 	assert.Equal(t, path, libraryPath(&compute.Library{Egg: path}))
+	assert.Equal(t, path, libraryPath(&compute.Library{Requirements: path}))
 	assert.Equal(t, "", libraryPath(&compute.Library{}))
 }