Update use_secret_as_env to accept secret name dynamically at runtime

Signed-off-by: ddalvi <[email protected]>

backend/src/v2/driver/driver.go

@@ -17,10 +17,12 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
-	"github.com/kubeflow/pipelines/backend/src/v2/objectstore"
 	"strconv"
+	"strings"
 	"time"
 
+	"github.com/kubeflow/pipelines/backend/src/v2/objectstore"
+
 	"github.com/golang/glog"
 	"github.com/golang/protobuf/ptypes/timestamp"
 	"github.com/google/uuid"
@@ -553,6 +555,31 @@ func extendPodSpecPatch(
 	// Get secret env information
 	for _, secretAsEnv := range kubernetesExecutorConfig.GetSecretAsEnv() {
 		for _, keyToEnv := range secretAsEnv.GetKeyToEnv() {
+			secretName := secretAsEnv.GetSecretName()
+
+			// Check if the secret name is dynamic (wrapped in "{{}}")
+			if strings.HasPrefix(secretName, "{{") && strings.HasSuffix(secretName, "}}") {
+				// Strip the braces
+				key := secretName[2 : len(secretName)-2]
+
+				// Check if the key exists in the parameter inputs map
+				inputParams, _, err := dag.Execution.GetParameters()
+				if err != nil {
+					return fmt.Errorf("failed to get input parameters: %v", err)
+				}
+
+				val, ok := inputParams[key]
+				if !ok {
+					return fmt.Errorf("dynamic secret name key '%s' not found in input parameters", key)
+				}
+
+				secretName = val.GetStringValue()
+				if secretName == "" {
+					return fmt.Errorf("secret name for key '%s' is empty", key)
+				}
+			} else if strings.TrimSpace(secretName) == "" {
+				return fmt.Errorf("secret name is empty or invalid")
+			}
 			secretEnvVar := k8score.EnvVar{
 				Name: keyToEnv.GetEnvVar(),
 				ValueFrom: &k8score.EnvVarSource{
@@ -561,7 +588,7 @@ func extendPodSpecPatch(
 					},
 				},
 			}
-			secretEnvVar.ValueFrom.SecretKeyRef.LocalObjectReference.Name = secretAsEnv.GetSecretName()
+			secretEnvVar.ValueFrom.SecretKeyRef.LocalObjectReference.Name = secretName
 			podSpec.Containers[0].Env = append(podSpec.Containers[0].Env, secretEnvVar)
 		}
 	}

kubernetes_platform/python/kfp/kubernetes/secret.py

@@ -16,13 +16,15 @@
 
 from google.protobuf import json_format
 from kfp.dsl import PipelineTask
+from typing import Union
+from kfp.dsl.pipeline_channel import PipelineParameterChannel
 from kfp.kubernetes import common
 from kfp.kubernetes import kubernetes_executor_config_pb2 as pb
 
 
 def use_secret_as_env(
     task: PipelineTask,
-    secret_name: str,
+    secret_name: Union[str, PipelineParameterChannel],
     secret_key_to_env: Dict[str, str],
 ) -> PipelineTask:
     """Use a Kubernetes Secret as an environment variable as described by the `Kubernetes documentation
@@ -39,14 +41,20 @@ def use_secret_as_env(
 
     msg = common.get_existing_kubernetes_config_as_message(task)
 
+    val = secret_name
+    # if secret_name is a PipelineParameterChannel, then we don't know what secret to mount until RUNTIME
+    # so, treat is as a map KEY instead of a secret name
+    if isinstance(secret_name, PipelineParameterChannel):
+        val = "{{" + secret_name.name + "}}"
+
     key_to_env = [
         pb.SecretAsEnv.SecretKeyToEnvMap(
             secret_key=secret_key,
             env_var=env_var,
         ) for secret_key, env_var in secret_key_to_env.items()
     ]
     secret_as_env = pb.SecretAsEnv(
-        secret_name=secret_name,
+        secret_name=val,
         key_to_env=key_to_env,
     )
 

kubernetes_platform/python/test/unit/test_secret.py

@@ -319,6 +319,36 @@ def my_pipeline():
             }
         }
 
+    def test_with_secret_name_param_env(self):
+        @dsl.pipeline
+        def my_pipeline(secret_name: str = 'my-secret'):
+            task = comp()
+            kubernetes.use_secret_as_env(
+                task,
+                secret_name=secret_name,
+                secret_key_to_env={'password': 'PASSWORD'}
+            )
+
+        assert json_format.MessageToDict(my_pipeline.platform_spec) == {
+            'platforms': {
+                'kubernetes': {
+                    'deploymentSpec': {
+                        'executors': {
+                            'exec-comp': {
+                                'secretAsEnv': [{
+                                    'secretName': '{{secret_name}}',
+                                    'keyToEnv': [{
+                                        'secretKey': 'password',
+                                        'envVar': 'PASSWORD'
+                                    }]
+                                }]
+                            }
+                        }
+                    }
+                }
+            }
+        }
+
     def test_preserves_secret_as_volume(self):
         # checks that use_secret_as_env respects previously set secrets as vol