Allow env from resource with keys & updated tests

hack/make-rules/test-cmd-util.sh

@@ -3174,12 +3174,22 @@ run_deployment_tests() {
   kube::test::get_object_assert deployment "{{range.items}}{{$id_field}}:{{end}}" 'nginx-deployment:'
   kube::test::get_object_assert configmap "{{range.items}}{{$id_field}}:{{end}}" 'test-set-env-config:'
   kube::test::get_object_assert secret "{{range.items}}{{$id_field}}:{{end}}" 'test-set-env-secret:'
+  # Set env of deployments by configmap from keys
+  kubectl set env deployment nginx-deployment --keys=key-2 --from=configmap/test-set-env-config "${kube_flags[@]}"
+  # Assert correct value in deployment env
+  kube::test::get_object_assert 'deploy nginx-deployment' "{{ (index (index .spec.template.spec.containers 0).env 0).name}}" 'KEY_2'
+  # Assert single value in deployment env
+  kube::test::get_object_assert 'deploy nginx-deployment' "{{ len (index .spec.template.spec.containers 0).env }}" '1'
+  # Set env of deployments by configmap
+  kubectl set env deployment nginx-deployment --from=configmap/test-set-env-config "${kube_flags[@]}"
+  # Assert all values in deployment env
+  kube::test::get_object_assert 'deploy nginx-deployment' "{{ len (index .spec.template.spec.containers 0).env }}" '2'
   # Set env of deployments for all container
   kubectl set env deployment nginx-deployment env=prod "${kube_flags[@]}"
   # Set env of deployments for specific container
   kubectl set env deployment nginx-deployment superenv=superprod -c=nginx "${kube_flags[@]}"
-  # Set env of deployments by configmap
-  kubectl set env deployment nginx-deployment --from=configmap/test-set-env-config "${kube_flags[@]}"
+  # Set env of deployments by secret from keys
+  kubectl set env deployment nginx-deployment --keys=username --from=secret/test-set-env-secret "${kube_flags[@]}"
   # Set env of deployments by secret
   kubectl set env deployment nginx-deployment --from=secret/test-set-env-secret "${kube_flags[@]}"
   # Remove specific env of deployment

pkg/kubectl/cmd/set/set_env.go

@@ -61,7 +61,7 @@ var (
 		` + envResources)
 
 	envExample = templates.Examples(`
-		# Update deployment 'registry' with a new environment variable
+          # Update deployment 'registry' with a new environment variable
 	  kubectl set env deployment/registry STORAGE_DIR=/local
 
 	  # List the environment variables defined on a deployments 'sample-build'
@@ -82,6 +82,9 @@ var (
 	  # Import environment from a config map with a prefix
 	  kubectl set env --from=configmap/myconfigmap --prefix=MYSQL_ deployment/myapp
 
+          # Import specific keys from a config map
+          kubectl set env --keys=my-example-key --from=configmap/myconfigmap deployment/myapp
+
 	  # Remove the environment variable ENV from container 'c1' in all deployment configs
 	  kubectl set env deployments --all --containers="c1" ENV-
 
@@ -107,6 +110,7 @@ type EnvOptions struct {
 	Selector          string
 	From              string
 	Prefix            string
+	Keys              []string
 
 	PrintObj printers.ResourcePrinterFunc
 
@@ -157,6 +161,7 @@ func NewCmdEnv(f cmdutil.Factory, streams genericclioptions.IOStreams) *cobra.Co
 	cmd.Flags().StringVarP(&o.From, "from", "", "", "The name of a resource from which to inject environment variables")
 	cmd.Flags().StringVarP(&o.Prefix, "prefix", "", "", "Prefix to append to variable names")
 	cmd.Flags().StringArrayVarP(&o.EnvParams, "env", "e", o.EnvParams, "Specify a key-value pair for an environment variable to set into each container.")
+	cmd.Flags().StringSliceVarP(&o.Keys, "keys", "", o.Keys, "Comma-separated list of keys to import from specified resource")
 	cmd.Flags().BoolVar(&o.List, "list", o.List, "If true, display the environment and any changes in the standard format. this flag will removed when we have kubectl view env.")
 	cmd.Flags().BoolVar(&o.Resolve, "resolve", o.Resolve, "If true, show secret or configmap references when listing variables")
 	cmd.Flags().StringVarP(&o.Selector, "selector", "l", o.Selector, "Selector (label query) to filter on")
@@ -183,6 +188,19 @@ func keyToEnvName(key string) string {
 	return strings.ToUpper(validEnvNameRegexp.ReplaceAllString(key, "_"))
 }
 
+func contains(key string, keyList []string) bool {
+	if len(keyList) == 0 {
+		return true
+	}
+
+	for _, k := range keyList {
+		if k == key {
+			return true
+		}
+	}
+	return false
+}
+
 func (o *EnvOptions) Complete(f cmdutil.Factory, cmd *cobra.Command, args []string) error {
 	if o.All && len(o.Selector) > 0 {
 		return fmt.Errorf("cannot set --all and --selector at the same time")
@@ -230,6 +248,9 @@ func (o *EnvOptions) Validate() error {
 	if o.List && len(o.output) > 0 {
 		return fmt.Errorf("--list and --output may not be specified together")
 	}
+	if len(o.Keys) > 0 && len(o.From) == 0 {
+		return fmt.Errorf("when specifying --keys, a configmap or secret must be provided with --from")
+	}
 	return nil
 }
 
@@ -265,33 +286,37 @@ func (o *EnvOptions) RunEnv() error {
 			switch from := info.Object.(type) {
 			case *v1.Secret:
 				for key := range from.Data {
-					envVar := v1.EnvVar{
-						Name: keyToEnvName(key),
-						ValueFrom: &v1.EnvVarSource{
-							SecretKeyRef: &v1.SecretKeySelector{
-								LocalObjectReference: v1.LocalObjectReference{
-									Name: from.Name,
+					if contains(key, o.Keys) {
+						envVar := v1.EnvVar{
+							Name: keyToEnvName(key),
+							ValueFrom: &v1.EnvVarSource{
+								SecretKeyRef: &v1.SecretKeySelector{
+									LocalObjectReference: v1.LocalObjectReference{
+										Name: from.Name,
+									},
+									Key: key,
 								},
-								Key: key,
 							},
-						},
+						}
+						env = append(env, envVar)
 					}
-					env = append(env, envVar)
 				}
 			case *v1.ConfigMap:
 				for key := range from.Data {
-					envVar := v1.EnvVar{
-						Name: keyToEnvName(key),
-						ValueFrom: &v1.EnvVarSource{
-							ConfigMapKeyRef: &v1.ConfigMapKeySelector{
-								LocalObjectReference: v1.LocalObjectReference{
-									Name: from.Name,
+					if contains(key, o.Keys) {
+						envVar := v1.EnvVar{
+							Name: keyToEnvName(key),
+							ValueFrom: &v1.EnvVarSource{
+								ConfigMapKeyRef: &v1.ConfigMapKeySelector{
+									LocalObjectReference: v1.LocalObjectReference{
+										Name: from.Name,
+									},
+									Key: key,
 								},
-								Key: key,
 							},
-						},
+						}
+						env = append(env, envVar)
 					}
-					env = append(env, envVar)
 				}
 			default:
 				return fmt.Errorf("unsupported resource specified in --from")

pkg/kubectl/cmd/set/set_env_test.go

@@ -492,3 +492,150 @@ func TestSetEnvRemote(t *testing.T) {
 		})
 	}
 }
+
+func TestSetEnvFromResource(t *testing.T) {
+	mockConfigMap := &v1.ConfigMap{
+		ObjectMeta: metav1.ObjectMeta{Name: "testconfigmap"},
+		Data: map[string]string{
+			"env":          "prod",
+			"test-key":     "testValue",
+			"test-key-two": "testValueTwo",
+		},
+	}
+
+	mockSecret := &v1.Secret{
+		ObjectMeta: metav1.ObjectMeta{Name: "testsecret"},
+		Data: map[string][]byte{
+			"env":          []byte("prod"),
+			"test-key":     []byte("testValue"),
+			"test-key-two": []byte("testValueTwo"),
+		},
+	}
+
+	inputs := []struct {
+		name           string
+		args           []string
+		from           string
+		keys           []string
+		assertIncludes []string
+		assertExcludes []string
+	}{
+		{
+			name: "test from configmap",
+			args: []string{"deployment", "nginx"},
+			from: "configmap/testconfigmap",
+			keys: []string{},
+			assertIncludes: []string{
+				`{"name":"ENV","valueFrom":{"configMapKeyRef":{"key":"env","name":"testconfigmap"}}}`,
+				`{"name":"TEST_KEY","valueFrom":{"configMapKeyRef":{"key":"test-key","name":"testconfigmap"}}}`,
+				`{"name":"TEST_KEY_TWO","valueFrom":{"configMapKeyRef":{"key":"test-key-two","name":"testconfigmap"}}}`,
+			},
+			assertExcludes: []string{},
+		},
+		{
+			name: "test from secret",
+			args: []string{"deployment", "nginx"},
+			from: "secret/testsecret",
+			keys: []string{},
+			assertIncludes: []string{
+				`{"name":"ENV","valueFrom":{"secretKeyRef":{"key":"env","name":"testsecret"}}}`,
+				`{"name":"TEST_KEY","valueFrom":{"secretKeyRef":{"key":"test-key","name":"testsecret"}}}`,
+				`{"name":"TEST_KEY_TWO","valueFrom":{"secretKeyRef":{"key":"test-key-two","name":"testsecret"}}}`,
+			},
+			assertExcludes: []string{},
+		},
+		{
+			name: "test from configmap with keys",
+			args: []string{"deployment", "nginx"},
+			from: "configmap/testconfigmap",
+			keys: []string{"env", "test-key-two"},
+			assertIncludes: []string{
+				`{"name":"ENV","valueFrom":{"configMapKeyRef":{"key":"env","name":"testconfigmap"}}}`,
+				`{"name":"TEST_KEY_TWO","valueFrom":{"configMapKeyRef":{"key":"test-key-two","name":"testconfigmap"}}}`,
+			},
+			assertExcludes: []string{`{"name":"TEST_KEY","valueFrom":{"configMapKeyRef":{"key":"test-key","name":"testconfigmap"}}}`},
+		},
+		{
+			name: "test from secret with keys",
+			args: []string{"deployment", "nginx"},
+			from: "secret/testsecret",
+			keys: []string{"env", "test-key-two"},
+			assertIncludes: []string{
+				`{"name":"ENV","valueFrom":{"secretKeyRef":{"key":"env","name":"testsecret"}}}`,
+				`{"name":"TEST_KEY_TWO","valueFrom":{"secretKeyRef":{"key":"test-key-two","name":"testsecret"}}}`,
+			},
+			assertExcludes: []string{`{"name":"TEST_KEY","valueFrom":{"secretKeyRef":{"key":"test-key","name":"testsecret"}}}`},
+		},
+	}
+
+	for _, input := range inputs {
+		mockDeployment := &appsv1.Deployment{
+			ObjectMeta: metav1.ObjectMeta{Name: "nginx"},
+			Spec: appsv1.DeploymentSpec{
+				Template: v1.PodTemplateSpec{
+					Spec: v1.PodSpec{
+						Containers: []v1.Container{
+							{
+								Name:  "nginx",
+								Image: "nginx",
+							},
+						},
+					},
+				},
+			},
+		}
+		t.Run(input.name, func(t *testing.T) {
+			tf := cmdtesting.NewTestFactory()
+			defer tf.Cleanup()
+
+			tf.Namespace = "test"
+			tf.ClientConfigVal = &restclient.Config{ContentConfig: restclient.ContentConfig{GroupVersion: &schema.GroupVersion{Version: ""}}}
+			tf.Client = &fake.RESTClient{
+				GroupVersion:         schema.GroupVersion{Group: "", Version: "v1"},
+				NegotiatedSerializer: serializer.DirectCodecFactory{CodecFactory: scheme.Codecs},
+				Client: fake.CreateHTTPClient(func(req *http.Request) (*http.Response, error) {
+					switch p, m := req.URL.Path, req.Method; {
+					case p == "/namespaces/test/configmaps/testconfigmap" && m == http.MethodGet:
+						return &http.Response{StatusCode: http.StatusOK, Header: defaultHeader(), Body: objBody(mockConfigMap)}, nil
+					case p == "/namespaces/test/secrets/testsecret" && m == http.MethodGet:
+						return &http.Response{StatusCode: http.StatusOK, Header: defaultHeader(), Body: objBody(mockSecret)}, nil
+					case p == "/namespaces/test/deployments/nginx" && m == http.MethodGet:
+						return &http.Response{StatusCode: http.StatusOK, Header: defaultHeader(), Body: objBody(mockDeployment)}, nil
+					case p == "/namespaces/test/deployments/nginx" && m == http.MethodPatch:
+						stream, err := req.GetBody()
+						if err != nil {
+							return nil, err
+						}
+						bytes, err := ioutil.ReadAll(stream)
+						if err != nil {
+							return nil, err
+						}
+						for _, include := range input.assertIncludes {
+							assert.Contains(t, string(bytes), include)
+						}
+						for _, exclude := range input.assertExcludes {
+							assert.NotContains(t, string(bytes), exclude)
+						}
+						return &http.Response{StatusCode: http.StatusOK, Header: defaultHeader(), Body: objBody(mockDeployment)}, nil
+					default:
+						t.Errorf("%s: unexpected request: %#v\n%#v", input.name, req.URL, req)
+						return nil, nil
+					}
+				}),
+			}
+
+			outputFormat := "yaml"
+			streams := genericclioptions.NewTestIOStreamsDiscard()
+			opts := NewEnvOptions(streams)
+			opts.From = input.from
+			opts.Keys = input.keys
+			opts.PrintFlags = genericclioptions.NewPrintFlags("").WithDefaultOutput(outputFormat).WithTypeSetter(scheme.Scheme)
+			opts.Local = false
+			opts.IOStreams = streams
+			err := opts.Complete(tf, NewCmdEnv(tf, streams), input.args)
+			assert.NoError(t, err)
+			err = opts.RunEnv()
+			assert.NoError(t, err)
+		})
+	}
+}