From 327035a43a5f277e2b5ee7156f6217297023822b Mon Sep 17 00:00:00 2001
From: Donny Xia <donnyxia@google.com>
Date: Tue, 15 Sep 2020 11:49:26 -0700
Subject: [PATCH] Add --env/-e flag

---
 cmd/config/internal/commands/run-fns.go  |  5 +++++
 cmd/config/internal/commands/run_test.go | 13 +++++++++++++
 2 files changed, 18 insertions(+)

diff --git a/cmd/config/internal/commands/run-fns.go b/cmd/config/internal/commands/run-fns.go
index dc987e8cdf..b7c38624da 100644
--- a/cmd/config/internal/commands/run-fns.go
+++ b/cmd/config/internal/commands/run-fns.go
@@ -69,6 +69,9 @@ func GetRunFnRunner(name string) *RunFnRunner {
 		"a list of storage options read from the filesystem")
 	r.Command.Flags().BoolVar(
 		&r.LogSteps, "log-steps", false, "log steps to stderr")
+	r.Command.Flags().StringArrayVarP(
+		&r.Env, "env", "e", []string{},
+		"a list of environment variables to be used by functions")
 	return r
 }
 
@@ -96,6 +99,7 @@ type RunFnRunner struct {
 	NetworkName        string
 	Mounts             []string
 	LogSteps           bool
+	Env                []string
 }
 
 func (r *RunFnRunner) runE(c *cobra.Command, args []string) error {
@@ -311,6 +315,7 @@ func (r *RunFnRunner) preRunE(c *cobra.Command, args []string) error {
 		StorageMounts:  storageMounts,
 		ResultsDir:     r.ResultsDir,
 		LogSteps:       r.LogSteps,
+		Env:            r.Env,
 	}
 
 	// don't consider args for the function
diff --git a/cmd/config/internal/commands/run_test.go b/cmd/config/internal/commands/run_test.go
index 094a6e6eb6..7cd1c957f7 100644
--- a/cmd/config/internal/commands/run_test.go
+++ b/cmd/config/internal/commands/run_test.go
@@ -204,6 +204,7 @@ apiVersion: v1
 				Path:           "dir",
 				NetworkName:    "bridge",
 				EnableStarlark: true,
+				Env:            []string{},
 			},
 		},
 		{
@@ -257,6 +258,7 @@ apiVersion: v1
 				Path:        "dir",
 				NetworkName: "bridge",
 				ResultsDir:  "foo/",
+				Env:         []string{},
 			},
 			expected: `
 metadata:
@@ -292,6 +294,17 @@ apiVersion: v1
 				Path:        "dir",
 				NetworkName: "bridge",
 				LogSteps:    true,
+				Env:         []string{},
+			},
+		},
+		{
+			name: "envs",
+			args: []string{"run", "dir", "--env", "FOO=BAR", "-e", "BAR"},
+			path: "dir",
+			expectedStruct: &runfn.RunFns{
+				Path:        "dir",
+				NetworkName: "bridge",
+				Env:         []string{"FOO=BAR", "BAR"},
 			},
 		},
 	}