From ca12c36f501718d1a506e244c9b2bfdbb7580edb Mon Sep 17 00:00:00 2001
From: Daniel Rammer <daniel@union.ai>
Date: Mon, 12 Jun 2023 09:46:21 -0500
Subject: [PATCH] propagating environment variables through launchplans

Signed-off-by: Daniel Rammer <daniel@union.ai>
---
 pkg/controller/nodes/subworkflow/launchplan.go  | 17 +++++++++--------
 .../nodes/subworkflow/launchplan/admin.go       |  9 +++++++++
 .../nodes/subworkflow/launchplan/launchplan.go  | 11 ++++++-----
 3 files changed, 24 insertions(+), 13 deletions(-)

diff --git a/pkg/controller/nodes/subworkflow/launchplan.go b/pkg/controller/nodes/subworkflow/launchplan.go
index 5b764059d..d903fee8d 100644
--- a/pkg/controller/nodes/subworkflow/launchplan.go
+++ b/pkg/controller/nodes/subworkflow/launchplan.go
@@ -67,14 +67,15 @@ func (l *launchPlanHandler) StartLaunchPlan(ctx context.Context, nCtx handler.No
 	}
 
 	launchCtx := launchplan.LaunchContext{
-		ParentNodeExecution: parentNodeExecutionID,
-		MaxParallelism:      nCtx.ExecutionContext().GetExecutionConfig().MaxParallelism,
-		SecurityContext:     nCtx.ExecutionContext().GetSecurityContext(),
-		RawOutputDataConfig: nCtx.ExecutionContext().GetRawOutputDataConfig().RawOutputDataConfig,
-		Labels:              nCtx.ExecutionContext().GetLabels(),
-		Annotations:         nCtx.ExecutionContext().GetAnnotations(),
-		Interruptible:       nCtx.ExecutionContext().GetExecutionConfig().Interruptible,
-		OverwriteCache:      nCtx.ExecutionContext().GetExecutionConfig().OverwriteCache,
+		ParentNodeExecution:  parentNodeExecutionID,
+		MaxParallelism:       nCtx.ExecutionContext().GetExecutionConfig().MaxParallelism,
+		SecurityContext:      nCtx.ExecutionContext().GetSecurityContext(),
+		RawOutputDataConfig:  nCtx.ExecutionContext().GetRawOutputDataConfig().RawOutputDataConfig,
+		Labels:               nCtx.ExecutionContext().GetLabels(),
+		Annotations:          nCtx.ExecutionContext().GetAnnotations(),
+		Interruptible:        nCtx.ExecutionContext().GetExecutionConfig().Interruptible,
+		OverwriteCache:       nCtx.ExecutionContext().GetExecutionConfig().OverwriteCache,
+		EnvironmentVariables: nCtx.ExecutionContext().GetExecutionConfig().EnvironmentVariables,
 	}
 
 	if nCtx.ExecutionContext().GetExecutionConfig().RecoveryExecution.WorkflowExecutionIdentifier != nil {
diff --git a/pkg/controller/nodes/subworkflow/launchplan/admin.go b/pkg/controller/nodes/subworkflow/launchplan/admin.go
index 60499f230..58099adf1 100644
--- a/pkg/controller/nodes/subworkflow/launchplan/admin.go
+++ b/pkg/controller/nodes/subworkflow/launchplan/admin.go
@@ -102,6 +102,14 @@ func (a *adminLaunchPlanExecutor) Launch(ctx context.Context, launchCtx LaunchCo
 		}
 	}
 
+	environmentVariables := make([]*core.KeyValuePair, 0, len(launchCtx.EnvironmentVariables))
+	for k, v := range launchCtx.EnvironmentVariables {
+		environmentVariables = append(environmentVariables, &core.KeyValuePair{
+			Key:   k,
+			Value: v,
+		})
+	}
+
 	req := &admin.ExecutionCreateRequest{
 		Project: executionID.Project,
 		Domain:  executionID.Domain,
@@ -122,6 +130,7 @@ func (a *adminLaunchPlanExecutor) Launch(ctx context.Context, launchCtx LaunchCo
 			RawOutputDataConfig: launchCtx.RawOutputDataConfig,
 			Interruptible:       interruptible,
 			OverwriteCache:      launchCtx.OverwriteCache,
+			Envs:                &admin.Envs{Values: environmentVariables},
 		},
 	}
 
diff --git a/pkg/controller/nodes/subworkflow/launchplan/launchplan.go b/pkg/controller/nodes/subworkflow/launchplan/launchplan.go
index 7d6588ca3..2d2a25b60 100644
--- a/pkg/controller/nodes/subworkflow/launchplan/launchplan.go
+++ b/pkg/controller/nodes/subworkflow/launchplan/launchplan.go
@@ -25,11 +25,12 @@ type LaunchContext struct {
 	// MaxParallelism
 	MaxParallelism uint32
 	// RawOutputDataConfig
-	RawOutputDataConfig *admin.RawOutputDataConfig
-	Annotations         map[string]string
-	Labels              map[string]string
-	Interruptible       *bool
-	OverwriteCache      bool
+	RawOutputDataConfig  *admin.RawOutputDataConfig
+	Annotations          map[string]string
+	Labels               map[string]string
+	Interruptible        *bool
+	OverwriteCache       bool
+	EnvironmentVariables map[string]string
 }
 
 // Executor interface to be implemented by the remote system that can allow workflow launching capabilities