prunev2: Add labels for objects that we apply

As we apply objects when using apply/prune v2, we want to be sure they
include the label that ties them back to the applyset they are part
of.

Co-Authored-By: Katrina Verey <katrina.verey@shopify.com>

Kubernetes-commit: ab058308401b35b4865424cfa43ed75a554af2a3

pkg/cmd/apply/apply.go

@@ -452,7 +452,15 @@ func (o *ApplyOptions) GetObjects() ([]*resource.Info, error) {
 			LabelSelectorParam(o.Selector).
 			Flatten().
 			Do()
+
 		o.objects, err = r.Infos()
+
+		if o.ApplySet != nil {
+			if err := o.ApplySet.addLabels(o.objects); err != nil {
+				return nil, err
+			}
+		}
+
 		o.objectsCached = true
 	}
 	return o.objects, err

pkg/cmd/apply/apply_test.go

@@ -28,6 +28,7 @@ import (
 	"strings"
 	"testing"
 
+	"github.com/google/go-cmp/cmp"
 	"github.com/spf13/cobra"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -57,6 +58,7 @@ import (
 	"k8s.io/kubectl/pkg/util/openapi"
 	utilpointer "k8s.io/utils/pointer"
 	"k8s.io/utils/strings/slices"
+	"sigs.k8s.io/yaml"
 )
 
 var (
@@ -2212,3 +2214,58 @@ func TestApplySetParentValidation(t *testing.T) {
 		}
 	})
 }
+
+func TestLoadObjects(t *testing.T) {
+	f := cmdtesting.NewTestFactory()
+	defer f.Cleanup()
+
+	testdirs := []string{"testdata/prune/simple"}
+	for _, testdir := range testdirs {
+		t.Run(testdir, func(t *testing.T) {
+			cmdtesting.WithAlphaEnvs([]cmdutil.FeatureGate{cmdutil.ApplySet}, t, func(t *testing.T) {
+				cmd := &cobra.Command{}
+				flags := NewApplyFlags(genericclioptions.NewTestIOStreamsDiscard())
+				flags.AddFlags(cmd)
+				cmd.Flags().Set("filename", filepath.Join(testdir, "manifest1.yaml"))
+				cmd.Flags().Set("applyset", filepath.Base(testdir))
+				cmd.Flags().Set("prune", "true")
+
+				o, err := flags.ToOptions(f, cmd, "kubectl", []string{})
+				if err != nil {
+					t.Fatalf("unexpected error creating apply options: %v", err)
+				}
+
+				// TODO(justinsb): Enable validation once we unblock --applyset
+				// err = o.Validate()
+				// if err != nil {
+				// 	t.Fatalf("unexpected error from validate: %v", err)
+				// }
+
+				resources, err := o.GetObjects()
+				if err != nil {
+					t.Fatalf("GetObjects gave unexpected error %v", err)
+				}
+
+				var objectYAMLs []string
+				for _, obj := range resources {
+					y, err := yaml.Marshal(obj.Object)
+					if err != nil {
+						t.Fatalf("error marshaling object: %v", err)
+					}
+					objectYAMLs = append(objectYAMLs, string(y))
+				}
+				got := strings.Join(objectYAMLs, "\n---\n\n")
+
+				p := filepath.Join(testdir, "expected-manifest1-getobjects.yaml")
+				wantBytes, err := os.ReadFile(p)
+				if err != nil {
+					t.Fatalf("error reading file %q: %v", p, err)
+				}
+				want := string(wantBytes)
+				if diff := cmp.Diff(want, got); diff != "" {
+					t.Errorf("GetObjects returned unexpected diff (-want +got):\n%s", diff)
+				}
+			})
+		})
+	}
+}

pkg/cmd/apply/applyset.go

@@ -23,6 +23,7 @@ import (
 	"k8s.io/apimachinery/pkg/api/meta"
 	"k8s.io/apimachinery/pkg/runtime/schema"
 	utilerrors "k8s.io/apimachinery/pkg/util/errors"
+	"k8s.io/cli-runtime/pkg/resource"
 )
 
 var defaultApplySetParentGVR = schema.GroupVersionResource{Version: "v1", Resource: "secrets"}
@@ -86,6 +87,36 @@ func (a ApplySet) Validate() error {
 	return utilerrors.NewAggregate(errors)
 }
 
+func (a *ApplySet) LabelsForMember() map[string]string {
+	return map[string]string{
+		"applyset.k8s.io/part-of": a.ID(),
+	}
+}
+
+// addLabels sets our tracking labels on each object; this should be called as part of loading the objects.
+func (a *ApplySet) addLabels(objects []*resource.Info) error {
+	applysetLabels := a.LabelsForMember()
+	for _, obj := range objects {
+		accessor, err := meta.Accessor(obj.Object)
+		if err != nil {
+			return fmt.Errorf("getting accessor: %w", err)
+		}
+		labels := accessor.GetLabels()
+		if labels == nil {
+			labels = make(map[string]string)
+		}
+		for k, v := range applysetLabels {
+			if _, found := labels[k]; found {
+				return fmt.Errorf("applyset label %q already set in input data", k)
+			}
+			labels[k] = v
+		}
+		accessor.SetLabels(labels)
+	}
+
+	return nil
+}
+
 func (p *ApplySetParentRef) IsNamespaced() bool {
 	return p.RESTMapping.Scope.Name() == meta.RESTScopeNameNamespace
 }

pkg/cmd/apply/testdata/prune/simple/expected-manifest1-getobjects.yaml

@@ -0,0 +1,15 @@
+apiVersion: v1
+kind: Namespace
+metadata:
+  labels:
+    applyset.k8s.io/part-of: placeholder-todo
+  name: foo
+
+---
+
+apiVersion: v1
+kind: Namespace
+metadata:
+  labels:
+    applyset.k8s.io/part-of: placeholder-todo
+  name: bar

pkg/cmd/apply/testdata/prune/simple/manifest1.yaml

@@ -0,0 +1,11 @@
+apiVersion: v1
+kind: Namespace
+metadata:
+  name: foo
+
+---
+
+apiVersion: v1
+kind: Namespace
+metadata:
+  name: bar