From 9298db470808a0f9b3cb347aac9d999b9b70ff30 Mon Sep 17 00:00:00 2001
From: Jeff Wong <awole20@gmail.com>
Date: Fri, 13 Sep 2024 08:56:06 -0700
Subject: [PATCH] Add support for build arg files in yaml format.

Adds key/value yaml files as possible input for build arg file.

This allows for multiline args support in files.
---
 go.mod                                   |  1 +
 task.go                                  | 22 +++++++++++++++++-----
 task_test.go                             |  9 +++++++++
 testdata/build-args/build_args_file.yaml |  2 ++
 4 files changed, 29 insertions(+), 5 deletions(-)
 create mode 100644 testdata/build-args/build_args_file.yaml

diff --git a/go.mod b/go.mod
index 7d45872..62e7600 100644
--- a/go.mod
+++ b/go.mod
@@ -16,4 +16,5 @@ require (
 	github.com/vbauerster/mpb v3.4.0+incompatible
 	github.com/vrischmann/envconfig v1.3.0
 	golang.org/x/sys v0.0.0-20220722155257-8c9f86f7a55f // indirect
+	gopkg.in/yaml.v3 v3.0.0-20210107192922-496545a6307b
 )
diff --git a/task.go b/task.go
index c5778b8..316eb31 100644
--- a/task.go
+++ b/task.go
@@ -3,6 +3,7 @@ package task
 import (
 	"encoding/json"
 	"fmt"
+	"gopkg.in/yaml.v3"
 	"io"
 	"io/ioutil"
 	"os"
@@ -374,13 +375,24 @@ func sanitize(cfg *Config) error {
 			return errors.Wrap(err, "read build args file")
 		}
 
-		for _, arg := range strings.Split(string(buildArgs), "\n") {
-			if len(arg) == 0 {
-				// skip blank lines
-				continue
+		if strings.HasSuffix(cfg.BuildArgsFile, ".yml") || strings.HasSuffix(cfg.BuildArgsFile, ".yaml") {
+			var buildArgsData map[string]string
+			err = yaml.Unmarshal(buildArgs, &buildArgsData)
+			if err != nil {
+				return errors.Wrap(err, "read build args yaml file")
 			}
+			for key, arg := range buildArgsData {
+				cfg.BuildArgs = append(cfg.BuildArgs, key + "=" + arg)
+			}
+		} else {
+			for _, arg := range strings.Split(string(buildArgs), "\n") {
+				if len(arg) == 0 {
+					// skip blank lines
+					continue
+				}
 
-			cfg.BuildArgs = append(cfg.BuildArgs, arg)
+				cfg.BuildArgs = append(cfg.BuildArgs, arg)
+			}
 		}
 	}
 
diff --git a/task_test.go b/task_test.go
index 6b3f04e..1975029 100644
--- a/task_test.go
+++ b/task_test.go
@@ -152,6 +152,15 @@ func (s *TaskSuite) TestBuildArgsFile() {
 	s.NoError(err)
 }
 
+func (s *TaskSuite) TestBuildArgsYamlFile() {
+	s.req.Config.ContextDir = "testdata/build-args"
+	s.req.Config.BuildArgsFile = "testdata/build-args/build_args_file.yaml"
+
+	// the Dockerfile itself asserts that the arg has been received
+	_, err := s.build()
+	s.NoError(err)
+}
+
 func (s *TaskSuite) TestBuildArgsStaticAndFile() {
 	s.req.Config.ContextDir = "testdata/build-args"
 	s.req.Config.BuildArgs = []string{"some_arg=some_value"}
diff --git a/testdata/build-args/build_args_file.yaml b/testdata/build-args/build_args_file.yaml
new file mode 100644
index 0000000..d1fb8ac
--- /dev/null
+++ b/testdata/build-args/build_args_file.yaml
@@ -0,0 +1,2 @@
+some_arg: some_value
+some_other_arg: some_other_value