Add env_name to Secrets

Signed-off-by: Thomas J. Fan <thomasjpfan@gmail.com>
flyteorg · Jan 10, 2025 · ef360e8 · ef360e8
1 parent dd7a215
commit ef360e8
Show file tree

Hide file tree

Showing 3 changed files with 116 additions and 0 deletions.
diff --git a/flytepropeller/pkg/webhook/k8s_secrets.go b/flytepropeller/pkg/webhook/k8s_secrets.go
@@ -75,11 +75,26 @@ func (i K8sSecretInjector) Inject(ctx context.Context, secret *core.Secret, p *c
 
 		p.Spec.InitContainers = AppendEnvVars(p.Spec.InitContainers, prefixEnvVar)
 		p.Spec.Containers = AppendEnvVars(p.Spec.Containers, prefixEnvVar)
+
+		if secret.GetEnvName() != "" {
+			extraEnvVar := CreateVolumeMountEnvVarForSecretWithEnvName(secret)
+			p.Spec.InitContainers = AppendEnvVars(p.Spec.InitContainers, extraEnvVar)
+			p.Spec.Containers = AppendEnvVars(p.Spec.Containers, extraEnvVar)
+		}
+
 	case core.Secret_ENV_VAR:
 		envVar := CreateEnvVarForSecret(secret)
 		p.Spec.InitContainers = AppendEnvVars(p.Spec.InitContainers, envVar)
 		p.Spec.Containers = AppendEnvVars(p.Spec.Containers, envVar)
 
+		if secret.GetEnvName() != "" {
+			extraEnvVar := *envVar.DeepCopy()
+			extraEnvVar.Name = secret.GetEnvName()
+
+			p.Spec.InitContainers = AppendEnvVars(p.Spec.InitContainers, extraEnvVar)
+			p.Spec.Containers = AppendEnvVars(p.Spec.Containers, extraEnvVar)
+		}
+
 		prefixEnvVar := corev1.EnvVar{
 			Name:  SecretEnvVarPrefix,
 			Value: K8sDefaultEnvVarPrefix,

diff --git a/flytepropeller/pkg/webhook/k8s_secrets_test.go b/flytepropeller/pkg/webhook/k8s_secrets_test.go
@@ -182,6 +182,95 @@ func TestK8sSecretInjector_Inject(t *testing.T) {
 		},
 	}
 
+	successPodEnvWithEnvName := corev1.Pod{
+		Spec: corev1.PodSpec{
+			InitContainers: []corev1.Container{},
+			Containers: []corev1.Container{
+				{
+					Name: "container1",
+					Env: []corev1.EnvVar{
+						{
+							Name: "_FSEC_GROUP_HELLO",
+							ValueFrom: &corev1.EnvVarSource{
+								SecretKeyRef: &corev1.SecretKeySelector{
+									Key: "HELLO",
+									LocalObjectReference: corev1.LocalObjectReference{
+										Name: "grOUP",
+									},
+									Optional: &optional,
+								},
+							},
+						},
+						{
+							Name: "MY_CUSTOM_ENV",
+							ValueFrom: &corev1.EnvVarSource{
+								SecretKeyRef: &corev1.SecretKeySelector{
+									Key: "HELLO",
+									LocalObjectReference: corev1.LocalObjectReference{
+										Name: "grOUP",
+									},
+									Optional: &optional,
+								},
+							},
+						},
+						{
+							Name:  "FLYTE_SECRETS_ENV_PREFIX",
+							Value: "_FSEC_",
+						},
+					},
+				},
+			},
+		},
+	}
+
+	successPodFileWithName := corev1.Pod{
+		Spec: corev1.PodSpec{
+			Volumes: []corev1.Volume{
+				{
+					Name: "m4ze5vkql3",
+					VolumeSource: corev1.VolumeSource{
+						Secret: &corev1.SecretVolumeSource{
+							SecretName: "grOUP",
+							Items: []corev1.KeyToPath{
+								{
+									Key:  "HELLO",
+									Path: "hello",
+								},
+							},
+							Optional: &optional,
+						},
+					},
+				},
+			},
+			InitContainers: []corev1.Container{},
+			Containers: []corev1.Container{
+				{
+					Name: "container1",
+					VolumeMounts: []corev1.VolumeMount{
+						{
+							Name:      "m4ze5vkql3",
+							MountPath: "/etc/flyte/secrets/group",
+							ReadOnly:  true,
+						},
+					},
+					Env: []corev1.EnvVar{
+						{
+							Name:  "FLYTE_SECRETS_DEFAULT_DIR",
+							Value: "/etc/flyte/secrets",
+						},
+						{
+							Name: "FLYTE_SECRETS_FILE_PREFIX",
+						},
+						{
+							Name:  "MY_CUSTOM_ENV",
+							Value: "/etc/flyte/secrets/group/hello",
+						},
+					},
+				},
+			},
+		},
+	}
+
 	ctx := context.Background()
 	type args struct {
 		secret *coreIdl.Secret
@@ -197,9 +286,14 @@ func TestK8sSecretInjector_Inject(t *testing.T) {
 			want: &corev1.Pod{}, wantErr: true},
 		{name: "simple", args: args{secret: &coreIdl.Secret{Group: "grOUP", Key: "HELLO", MountRequirement: coreIdl.Secret_ENV_VAR}, p: inputPod.DeepCopy()},
 			want: &successPodEnv, wantErr: false},
+		{name: "simple with env_name", args: args{secret: &coreIdl.Secret{Group: "grOUP", Key: "HELLO", MountRequirement: coreIdl.Secret_ENV_VAR, EnvName: "MY_CUSTOM_ENV"}, p: inputPod.DeepCopy()},
+			want: &successPodEnvWithEnvName, wantErr: false},
 		{name: "require file single", args: args{secret: &coreIdl.Secret{Group: "grOUP", Key: "HELLO", MountRequirement: coreIdl.Secret_FILE},
 			p: inputPod.DeepCopy()},
 			want: &successPodFile, wantErr: false},
+		{name: "require file single with name", args: args{secret: &coreIdl.Secret{Group: "grOUP", Key: "HELLO", MountRequirement: coreIdl.Secret_FILE, EnvName: "MY_CUSTOM_ENV"},
+			p: inputPod.DeepCopy()},
+			want: &successPodFileWithName, wantErr: false},
 		{name: "require file multiple from same secret group", args: args{secret: &coreIdl.Secret{Group: "grOUP", Key: "world", MountRequirement: coreIdl.Secret_FILE},
 			p: successPodFile.DeepCopy()},
 			want: &successPodMultiFiles, wantErr: false},

diff --git a/flytepropeller/pkg/webhook/utils.go b/flytepropeller/pkg/webhook/utils.go
@@ -67,6 +67,13 @@ func CreateVolumeMountForSecret(volumeName string, secret *core.Secret) corev1.V
 	}
 }
 
+func CreateVolumeMountEnvVarForSecretWithEnvName(secret *core.Secret) corev1.EnvVar {
+	return corev1.EnvVar{
+		Name:  secret.GetEnvName(),
+		Value: filepath.Join(filepath.Join(K8sSecretPathPrefix...), strings.ToLower(secret.GetGroup()), strings.ToLower(secret.GetKey())),
+	}
+}
+
 func AppendVolumeMounts(containers []corev1.Container, mount corev1.VolumeMount) []corev1.Container {
 	res := make([]corev1.Container, 0, len(containers))
 	for _, c := range containers {