From b512ca8afea9b26e003fb54ecadbe6ec4a49025a Mon Sep 17 00:00:00 2001
From: Tomas <40318863+tomasmik@users.noreply.github.com>
Date: Tue, 23 May 2023 14:09:47 +0300
Subject: [PATCH] Allow to pass env vars for local preview (#145)

* Allow to pass env vars for local preview
* comment fixes
---
 internal/cmd/stack/flags.go         | 10 +++++
 internal/cmd/stack/local_preview.go | 58 +++++++++++++++++++++++++++--
 internal/cmd/stack/stack.go         |  2 +
 3 files changed, 67 insertions(+), 3 deletions(-)

diff --git a/internal/cmd/stack/flags.go b/internal/cmd/stack/flags.go
index 1fbf9f3..b00bcc9 100644
--- a/internal/cmd/stack/flags.go
+++ b/internal/cmd/stack/flags.go
@@ -114,3 +114,13 @@ var flagShowLabels = &cli.BoolFlag{
 	Required: false,
 	Value:    false,
 }
+
+var flagOverrideEnvVarsTF = &cli.StringSliceFlag{
+	Name:  "tf-env-var-override",
+	Usage: "[Optional] Terraform environment variables injected into the run at runtime, they will be prefixed with TF_ by default, example: --tf-env-var-override 'foo=bar,bar=baz'",
+}
+
+var flagOverrideEnvVars = &cli.StringSliceFlag{
+	Name:  "env-var-override",
+	Usage: "[Optional] Environment variables injected into the run at runtime, example: --env-var-override 'foo=bar,bar=baz'",
+}
diff --git a/internal/cmd/stack/local_preview.go b/internal/cmd/stack/local_preview.go
index ad545e4..4caace4 100644
--- a/internal/cmd/stack/local_preview.go
+++ b/internal/cmd/stack/local_preview.go
@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"os"
 	"path/filepath"
+	"strings"
 
 	"github.com/mholt/archiver/v3"
 	"github.com/shurcooL/graphql"
@@ -16,10 +17,16 @@ import (
 
 func localPreview() cli.ActionFunc {
 	return func(cliCtx *cli.Context) error {
+		envVars, err := parseEnvVariablesForLocalPreview(cliCtx)
+		if err != nil {
+			return err
+		}
+
 		stackID, err := getStackID(cliCtx)
 		if err != nil {
 			return err
 		}
+
 		ctx := context.Background()
 
 		var packagePath *string = nil
@@ -88,12 +95,13 @@ func localPreview() cli.ActionFunc {
 		var triggerMutation struct {
 			RunProposeLocalWorkspace struct {
 				ID string `graphql:"id"`
-			} `graphql:"runProposeLocalWorkspace(stack: $stack, workspace: $workspace)"`
+			} `graphql:"runProposeLocalWorkspace(stack: $stack, workspace: $workspace, environmentVarsOverrides: $environmentVarsOverrides)"`
 		}
 
 		triggerVariables := map[string]interface{}{
-			"stack":     graphql.ID(stackID),
-			"workspace": graphql.ID(uploadMutation.UploadLocalWorkspace.ID),
+			"stack":                    graphql.ID(stackID),
+			"workspace":                graphql.ID(uploadMutation.UploadLocalWorkspace.ID),
+			"environmentVarsOverrides": envVars,
 		}
 
 		var requestOpts []graphql.RequestOption
@@ -125,3 +133,47 @@ func localPreview() cli.ActionFunc {
 		return terminal.Error()
 	}
 }
+
+// EnvironmentVariable represents a key-value pair of environment variables
+type EnvironmentVariable struct {
+	Key   graphql.String `json:"key"`
+	Value graphql.String `json:"value"`
+}
+
+func parseEnvVariablesForLocalPreview(cliCtx *cli.Context) ([]EnvironmentVariable, error) {
+	envVars := make([]EnvironmentVariable, 0)
+
+	parseFn := func(ev string, mutateKey func(string) string) error {
+		parts := strings.Split(ev, "=")
+		if len(parts) != 2 {
+			return fmt.Errorf("invalid environment variable %q, must be in the form of KEY=VALUE", ev)
+		}
+
+		if mutateKey != nil {
+			parts[0] = mutateKey(parts[0])
+		}
+
+		envVars = append(envVars, EnvironmentVariable{
+			Key:   graphql.String(parts[0]),
+			Value: graphql.String(parts[1]),
+		})
+
+		return nil
+	}
+
+	for _, ev := range cliCtx.StringSlice(flagOverrideEnvVars.Name) {
+		if err := parseFn(ev, nil); err != nil {
+			return nil, err
+		}
+	}
+
+	for _, ev := range cliCtx.StringSlice(flagOverrideEnvVarsTF.Name) {
+		if err := parseFn(ev, func(s string) string {
+			return strings.Join([]string{"TF_", s}, "")
+		}); err != nil {
+			return nil, err
+		}
+	}
+
+	return envVars, nil
+}
diff --git a/internal/cmd/stack/stack.go b/internal/cmd/stack/stack.go
index 010226c..e7c1724 100644
--- a/internal/cmd/stack/stack.go
+++ b/internal/cmd/stack/stack.go
@@ -118,6 +118,8 @@ func Command() *cli.Command {
 					flagRunMetadata,
 					flagNoTail,
 					flagNoUpload,
+					flagOverrideEnvVars,
+					flagOverrideEnvVarsTF,
 				},
 				Action:    localPreview(),
 				Before:    authenticated.Ensure,