From 52a4f73b298a0a322795eca671664849ef43619c Mon Sep 17 00:00:00 2001
From: Yuvraj <code@evalsocket.dev>
Date: Fri, 18 Mar 2022 19:39:45 +0530
Subject: [PATCH] fix stow stroage config in copiliot container (#250)

* fix stow stroage config in copiliot

Signed-off-by: Yuvraj <code@evalsocket.dev>

* more changes

Signed-off-by: Yuvraj <code@evalsocket.dev>

* make lint

Signed-off-by: Yuvraj <code@evalsocket.dev>

* fix gcp flag issue

Signed-off-by: Yuvraj <code@evalsocket.dev>
---
 flyteplugins/go.mod                           |  2 +-
 flyteplugins/go.sum                           |  2 --
 .../tasks/pluginmachinery/flytek8s/copilot.go | 24 +++++++++----------
 .../pluginmachinery/flytek8s/copilot_test.go  | 12 ++++++++++
 4 files changed, 25 insertions(+), 15 deletions(-)

diff --git a/flyteplugins/go.mod b/flyteplugins/go.mod
index e90de1796e..e8c9bcfbeb 100644
--- a/flyteplugins/go.mod
+++ b/flyteplugins/go.mod
@@ -14,7 +14,7 @@ require (
 	github.com/aws/aws-sdk-go-v2/service/athena v1.0.0
 	github.com/coocood/freecache v1.1.1
 	github.com/flyteorg/flyteidl v0.23.0
-    github.com/flyteorg/flytestdlib v0.4.13
+	github.com/flyteorg/flytestdlib v0.4.13
 	github.com/go-logr/zapr v0.4.0 // indirect
 	github.com/go-test/deep v1.0.7
 	github.com/golang/protobuf v1.4.3
diff --git a/flyteplugins/go.sum b/flyteplugins/go.sum
index a34c7c98f5..ede5ef48f3 100644
--- a/flyteplugins/go.sum
+++ b/flyteplugins/go.sum
@@ -233,8 +233,6 @@ github.com/fatih/color v1.7.0/go.mod h1:Zm6kSWBoL9eyXnKyktHP6abPY2pDugNf5Kwzbycv
 github.com/fatih/color v1.10.0 h1:s36xzo75JdqLaaWoiEHk767eHiwo0598uUxyfiPkDsg=
 github.com/fatih/color v1.10.0/go.mod h1:ELkj/draVOlAH/xkhN6mQ50Qd0MPOk5AAr3maGEBuJM=
 github.com/fatih/structtag v1.2.0/go.mod h1:mBJUNpUnHmRKrKlQQlmCrh5PuhftFbNv8Ys4/aAZl94=
-github.com/flyteorg/flyteidl v0.21.23 h1:hzGIFNOt3VooW/NdnaicXijn3EKjNKTz1kY+tlHkED4=
-github.com/flyteorg/flyteidl v0.21.23/go.mod h1:576W2ViEyjTpT+kEVHAGbrTP3HARNUZ/eCwrNPmdx9U=
 github.com/flyteorg/flyteidl v0.23.0 h1:Pjl9Tq1pJfIK0au5PiqPVpl25xTYosN6BruZl+PgWAk=
 github.com/flyteorg/flyteidl v0.23.0/go.mod h1:576W2ViEyjTpT+kEVHAGbrTP3HARNUZ/eCwrNPmdx9U=
 github.com/flyteorg/flytestdlib v0.3.13/go.mod h1:Tz8JCECAbX6VWGwFT6cmEQ+RJpZ/6L9pswu3fzWs220=
diff --git a/flyteplugins/go/tasks/pluginmachinery/flytek8s/copilot.go b/flyteplugins/go/tasks/pluginmachinery/flytek8s/copilot.go
index aa153b082c..50e7b2d747 100644
--- a/flyteplugins/go/tasks/pluginmachinery/flytek8s/copilot.go
+++ b/flyteplugins/go/tasks/pluginmachinery/flytek8s/copilot.go
@@ -63,23 +63,23 @@ func CopilotCommandArgs(storageConfig *storage.Config) []string {
 		"/bin/flyte-copilot",
 		"--storage.limits.maxDownloadMBs=0",
 	}
+	if storageConfig.MultiContainerEnabled {
+		commands = append(commands, "--storage.enable-multicontainer")
+	}
+	if len(storageConfig.InitContainer) > 0 {
+		commands = append(commands, fmt.Sprintf("--storage.container=%s", storageConfig.InitContainer))
+
+	}
+	commands = append(commands, fmt.Sprintf("--storage.type=%s", storageConfig.Type))
+
 	if len(storageConfig.Stow.Config) > 0 && len(storageConfig.Stow.Kind) > 0 {
-		var cfg string
 		for key, val := range storageConfig.Stow.Config {
-			cfg += fmt.Sprintf("%s=%s,", key, val)
+			commands = append(commands, "--storage.stow.config")
+			commands = append(commands, fmt.Sprintf("%s=%s", key, val))
 		}
-		commands = append(commands, []string{
-			fmt.Sprintf("--storage.stow.config %s", cfg),
-			fmt.Sprintf("--storage.stow.kind=%s", storageConfig.Stow.Kind),
-		}...)
-	}
-	if storageConfig.MultiContainerEnabled {
-		commands = append(commands, "--storage.enable-multicontainer")
+		return append(commands, fmt.Sprintf("--storage.stow.kind=%s", storageConfig.Stow.Kind))
 	}
 	return append(commands, []string{
-		fmt.Sprintf("--storage.type=%s", storageConfig.Type),
-		fmt.Sprintf("--storage.enable-multicontainer=%v", storageConfig.MultiContainerEnabled),
-		fmt.Sprintf("--storage.container=%s", storageConfig.InitContainer),
 		fmt.Sprintf("--storage.connection.secret-key=%s", storageConfig.Connection.SecretKey),
 		fmt.Sprintf("--storage.connection.access-key=%s", storageConfig.Connection.AccessKey),
 		fmt.Sprintf("--storage.connection.auth-type=%s", storageConfig.Connection.AuthType),
diff --git a/flyteplugins/go/tasks/pluginmachinery/flytek8s/copilot_test.go b/flyteplugins/go/tasks/pluginmachinery/flytek8s/copilot_test.go
index a89803d437..587a29ac4f 100644
--- a/flyteplugins/go/tasks/pluginmachinery/flytek8s/copilot_test.go
+++ b/flyteplugins/go/tasks/pluginmachinery/flytek8s/copilot_test.go
@@ -80,6 +80,18 @@ func TestFlyteCoPilotContainer(t *testing.T) {
 		assert.Equal(t, 1, len(c.VolumeMounts))
 	})
 
+	t.Run("happy stow GCP backend", func(t *testing.T) {
+		storage.GetConfig().Type = storage.TypeStow
+		storage.GetConfig().InitContainer = "bucket"
+		storage.GetConfig().Stow.Kind = "google"
+		storage.GetConfig().Stow.Config = map[string]string{
+			"json":       "",
+			"project_id": "flyte-gcp",
+			"scope":      "read_write",
+		}
+		assert.Equal(t, 11, len(CopilotCommandArgs(storage.GetConfig())))
+	})
+
 	t.Run("bad-res-cpu", func(t *testing.T) {
 		old := cfg.CPU
 		cfg.CPU = "x"