Updates k8s_secrets

Signed-off-by: Geert Pingen <[email protected]>
flyteorg · Jul 8, 2023 · d36a30b · d36a30b
1 parent f6d99e7
commit d36a30b
Show file tree

Hide file tree

Showing 2 changed files with 61 additions and 35 deletions.
diff --git a/flytepropeller/pkg/webhook/global_secrets.go b/flytepropeller/pkg/webhook/global_secrets.go
@@ -36,37 +36,40 @@ func (g GlobalSecrets) Inject(ctx context.Context, secret *coreIdl.Secret, p *co
 		return p, false, err
 	}
 
-	switch secret.MountRequirement {
-	case coreIdl.Secret_FILE:
-		return nil, false, fmt.Errorf("global secrets can only be injected as environment "+
-			"variables [%v/%v]", secret.Group, secret.Key)
-	case coreIdl.Secret_ANY:
-		fallthrough
-	case coreIdl.Secret_ENV_VAR:
-		if len(secret.Group) == 0 {
-			return nil, false, fmt.Errorf("mounting a secret to env var requires selecting the "+
-				"secret and a single key within. Key [%v]", secret.Key)
+	if secret.MountTarget != nil {
+		switch secret.MountTarget.(type) {
+		case *coreIdl.Secret_EnvVar:
+			target, ok := secret.GetMountTarget().(*coreIdl.Secret_EnvVar)
+			if ok {
+				InjectEnvVar(p, secret, &target.EnvVar.Name, v)
+			}
+		case *coreIdl.Secret_File:
+			return nil, false, fmt.Errorf("global secrets can only be injected as environment "+
+				"variables [%v/%v]", secret.Group, secret.Key)
+		default:
+			err := fmt.Errorf("unrecognized mount target [%v] for secret [%v]", secret.GetMountTarget(), secret.Key)
+			logger.Error(ctx, err)
+			return p, false, err
 		}
-
-		envVar := corev1.EnvVar{
-			Name:  strings.ToUpper(K8sDefaultEnvVarPrefix + secret.Group + EnvVarGroupKeySeparator + secret.Key),
-			Value: v,
-		}
-
-		prefixEnvVar := corev1.EnvVar{
-			Name:  SecretEnvVarPrefix,
-			Value: K8sDefaultEnvVarPrefix,
+	} else {
+		switch secret.MountRequirement {
+		case coreIdl.Secret_FILE:
+			return nil, false, fmt.Errorf("global secrets can only be injected as environment "+
+				"variables [%v/%v]", secret.Group, secret.Key)
+		case coreIdl.Secret_ANY:
+			fallthrough
+		case coreIdl.Secret_ENV_VAR:
+			if len(secret.Group) == 0 {
+				return nil, false, fmt.Errorf("mounting a secret to env var requires selecting the "+
+					"secret and a single key within. Key [%v]", secret.Key)
+			}
+
+			InjectEnvVar(p, secret, nil, v)
+		default:
+			err := fmt.Errorf("unrecognized mount requirement [%v] for secret [%v]", secret.MountRequirement.String(), secret.Key)
+			logger.Error(ctx, err)
+			return p, false, err
 		}
-
-		p.Spec.InitContainers = AppendEnvVars(p.Spec.InitContainers, prefixEnvVar)
-		p.Spec.Containers = AppendEnvVars(p.Spec.Containers, prefixEnvVar)
-
-		p.Spec.InitContainers = AppendEnvVars(p.Spec.InitContainers, envVar)
-		p.Spec.Containers = AppendEnvVars(p.Spec.Containers, envVar)
-	default:
-		err := fmt.Errorf("unrecognized mount requirement [%v] for secret [%v]", secret.MountRequirement.String(), secret.Key)
-		logger.Error(ctx, err)
-		return p, false, err
 	}
 
 	return p, true, nil
@@ -77,3 +80,26 @@ func NewGlobalSecrets(provider GlobalSecretProvider) GlobalSecrets {
 		envSecretManager: provider,
 	}
 }
+
+func InjectEnvVar(p *corev1.Pod, secret *coreIdl.Secret, envVarName *string, value string) {
+	_envVarName := strings.ToUpper(K8sDefaultEnvVarPrefix + secret.Group + EnvVarGroupKeySeparator + secret.Key)
+	if envVarName != nil {
+		_envVarName = *envVarName
+	}
+
+	envVar := corev1.EnvVar{
+		Name:  _envVarName,
+		Value: value,
+	}
+
+	prefixEnvVar := corev1.EnvVar{
+		Name:  SecretEnvVarPrefix,
+		Value: K8sDefaultEnvVarPrefix,
+	}
+
+	p.Spec.InitContainers = AppendEnvVars(p.Spec.InitContainers, prefixEnvVar)
+	p.Spec.Containers = AppendEnvVars(p.Spec.Containers, prefixEnvVar)
+
+	p.Spec.InitContainers = AppendEnvVars(p.Spec.InitContainers, envVar)
+	p.Spec.Containers = AppendEnvVars(p.Spec.Containers, envVar)
+}
diff --git a/flytepropeller/pkg/webhook/k8s_secrets.go b/flytepropeller/pkg/webhook/k8s_secrets.go
@@ -50,12 +50,12 @@ func (i K8sSecretInjector) Inject(ctx context.Context, secret *core.Secret, p *c
 		case *core.Secret_EnvVar:
 			target, ok := secret.GetMountTarget().(*core.Secret_EnvVar)
 			if ok {
-				injectSecretAsEnvVar(p, secret, &target.EnvVar.Name)
+				InjectSecretAsEnvVar(p, secret, &target.EnvVar.Name)
 			}
 		case *core.Secret_File:
 			target, ok := secret.GetMountTarget().(*core.Secret_File)
 			if ok {
-				injectSecretAsFile(p, secret, &target.File.Path)
+				InjectSecretAsFile(p, secret, &target.File.Path)
 			}
 		default:
 			err := fmt.Errorf("unrecognized mount target [%v] for secret [%v]", secret.GetMountTarget(), secret.Key)
@@ -67,9 +67,9 @@ func (i K8sSecretInjector) Inject(ctx context.Context, secret *core.Secret, p *c
 		case core.Secret_ANY:
 			fallthrough
 		case core.Secret_FILE:
-			injectSecretAsFile(p, secret, nil)
+			InjectSecretAsFile(p, secret, nil)
 		case core.Secret_ENV_VAR:
-			injectSecretAsEnvVar(p, secret, nil)
+			InjectSecretAsEnvVar(p, secret, nil)
 		default:
 			err := fmt.Errorf("unrecognized mount requirement [%v] for secret [%v]", secret.MountRequirement.String(), secret.Key)
 			logger.Error(ctx, err)
@@ -84,7 +84,7 @@ func NewK8sSecretsInjector() K8sSecretInjector {
 	return K8sSecretInjector{}
 }
 
-func injectSecretAsFile(p *corev1.Pod, secret *core.Secret, mountPath *string) {
+func InjectSecretAsFile(p *corev1.Pod, secret *core.Secret, mountPath *string) {
 	// Inject a Volume that to the pod and all of its containers and init containers that mounts the secret into a
 	// file.
 	volume := CreateVolumeForSecret(secret)
@@ -117,7 +117,7 @@ func injectSecretAsFile(p *corev1.Pod, secret *core.Secret, mountPath *string) {
 	p.Spec.Containers = AppendEnvVars(p.Spec.Containers, prefixEnvVar)
 }
 
-func injectSecretAsEnvVar(p *corev1.Pod, secret *core.Secret, envVarName *string) {
+func InjectSecretAsEnvVar(p *corev1.Pod, secret *core.Secret, envVarName *string) {
 	envVar := CreateEnvVarForSecret(secret)
 	if envVarName != nil {
 		envVar = CreateNamedEnvVarForSecret(secret, *envVarName)