More flexible resource type to StatusReader mapping

pkg/kstatus/polling/engine/engine.go

@@ -24,7 +24,7 @@ type ClusterReaderFactoryFunc func(reader client.Reader, mapper meta.RESTMapper,
 type PollerEngine struct {
 	Reader              client.Reader
 	Mapper              meta.RESTMapper
-	StatusReaders       map[schema.GroupKind]StatusReader
+	StatusReaders       []StatusReader
 	DefaultStatusReader StatusReader
 }
 
@@ -144,7 +144,7 @@ type statusPollerRunner struct {
 	// statusReaders contains the resource specific statusReaders. These will contain logic for how to
 	// compute status for specific GroupKinds. These will use an ClusterReader to fetch
 	// status of a resource and any generated resources.
-	statusReaders map[schema.GroupKind]StatusReader
+	statusReaders []StatusReader
 
 	// defaultStatusReader is the generic engine that is used for all GroupKinds that
 	// doesn't have a specific engine in the statusReaders map.
@@ -237,11 +237,12 @@ func (r *statusPollerRunner) pollStatusForAllResources() {
 }
 
 func (r *statusPollerRunner) statusReaderForGroupKind(gk schema.GroupKind) StatusReader {
-	statusReader, ok := r.statusReaders[gk]
-	if !ok {
-		return r.defaultStatusReader
+	for _, sr := range r.statusReaders {
+		if sr.Supports(gk) {
+			return sr
+		}
 	}
-	return statusReader
+	return r.defaultStatusReader
 }
 
 func (r *statusPollerRunner) isUpdatedResourceStatus(resourceStatus *event.ResourceStatus) bool {

pkg/kstatus/polling/engine/engine_test.go

@@ -111,7 +111,7 @@ func TestStatusPollerRunner(t *testing.T) {
 			engine := PollerEngine{
 				Mapper:              fakeMapper,
 				DefaultStatusReader: tc.defaultStatusReader,
-				StatusReaders:       map[schema.GroupKind]StatusReader{},
+				StatusReaders:       []StatusReader{},
 			}
 
 			options := Options{
@@ -216,6 +216,10 @@ type fakeStatusReader struct {
 	resourceStatusCount map[schema.GroupKind]int
 }
 
+func (f *fakeStatusReader) Supports(schema.GroupKind) bool {
+	return true
+}
+
 func (f *fakeStatusReader) ReadStatus(_ context.Context, _ ClusterReader, identifier object.ObjMetadata) *event.ResourceStatus {
 	count := f.resourceStatusCount[identifier.GroupKind]
 	resourceStatusSlice := f.resourceStatuses[identifier.GroupKind]

pkg/kstatus/polling/engine/status_reader.go

@@ -7,6 +7,7 @@ import (
 	"context"
 
 	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
+	"k8s.io/apimachinery/pkg/runtime/schema"
 	"sigs.k8s.io/cli-utils/pkg/kstatus/polling/event"
 	"sigs.k8s.io/cli-utils/pkg/object"
 )
@@ -18,6 +19,10 @@ import (
 // how to identify these generated resources and how to compute status for
 // these generated resources.
 type StatusReader interface {
+	// Supports tells the caller whether the StatusReader can compute status for
+	// the provided GroupKind.
+	Supports(schema.GroupKind) bool
+
 	// ReadStatus will fetch the resource identified by the given identifier
 	// from the cluster and return an ResourceStatus that will contain
 	// information about the latest state of the resource, its computed status

pkg/kstatus/polling/polling.go

@@ -7,9 +7,7 @@ import (
 	"context"
 	"time"
 
-	appsv1 "k8s.io/api/apps/v1"
 	"k8s.io/apimachinery/pkg/api/meta"
-	"k8s.io/apimachinery/pkg/runtime/schema"
 	"sigs.k8s.io/cli-utils/pkg/kstatus/polling/clusterreader"
 	"sigs.k8s.io/cli-utils/pkg/kstatus/polling/engine"
 	"sigs.k8s.io/cli-utils/pkg/kstatus/polling/event"
@@ -21,11 +19,14 @@ import (
 
 // NewStatusPoller creates a new StatusPoller using the given clusterreader and mapper. The StatusPoller
 // will use the client for all calls to the cluster.
-func NewStatusPoller(reader client.Reader, mapper meta.RESTMapper, customStatusReaders map[schema.GroupKind]engine.StatusReader) *StatusPoller {
-	statusReaders, defaultStatusReader := createStatusReaders(mapper)
-	for gk, sr := range customStatusReaders {
-		statusReaders[gk] = sr
-	}
+func NewStatusPoller(reader client.Reader, mapper meta.RESTMapper, customStatusReaders []engine.StatusReader) *StatusPoller {
+	var statusReaders []engine.StatusReader
+
+	statusReaders = append(statusReaders, customStatusReaders...)
+
+	srs, defaultStatusReader := createStatusReaders(mapper)
+	statusReaders = append(statusReaders, srs...)
+
 	return &StatusPoller{
 		engine: &engine.PollerEngine{
 			Reader:              reader,
@@ -69,17 +70,17 @@ type Options struct {
 // a specific statusreaders.
 // TODO: We should consider making the registration more automatic instead of having to create each of them
 // here. Also, it might be worth creating them on demand.
-func createStatusReaders(mapper meta.RESTMapper) (map[schema.GroupKind]engine.StatusReader, engine.StatusReader) {
+func createStatusReaders(mapper meta.RESTMapper) ([]engine.StatusReader, engine.StatusReader) {
 	defaultStatusReader := statusreaders.NewGenericStatusReader(mapper, status.Compute)
 
 	replicaSetStatusReader := statusreaders.NewReplicaSetStatusReader(mapper, defaultStatusReader)
 	deploymentStatusReader := statusreaders.NewDeploymentResourceReader(mapper, replicaSetStatusReader)
 	statefulSetStatusReader := statusreaders.NewStatefulSetResourceReader(mapper, defaultStatusReader)
 
-	statusReaders := map[schema.GroupKind]engine.StatusReader{
-		appsv1.SchemeGroupVersion.WithKind("Deployment").GroupKind():  deploymentStatusReader,
-		appsv1.SchemeGroupVersion.WithKind("StatefulSet").GroupKind(): statefulSetStatusReader,
-		appsv1.SchemeGroupVersion.WithKind("ReplicaSet").GroupKind():  replicaSetStatusReader,
+	statusReaders := []engine.StatusReader{
+		deploymentStatusReader,
+		statefulSetStatusReader,
+		replicaSetStatusReader,
 	}
 
 	return statusReaders, defaultStatusReader

pkg/kstatus/polling/statusreaders/common.go

@@ -44,9 +44,14 @@ type baseStatusReader struct {
 // resourceTypeStatusReader is an interface that can be implemented differently
 // for each resource type.
 type resourceTypeStatusReader interface {
+	Supports(gk schema.GroupKind) bool
 	ReadStatusForObject(ctx context.Context, reader engine.ClusterReader, object *unstructured.Unstructured) *event.ResourceStatus
 }
 
+func (b *baseStatusReader) Supports(gk schema.GroupKind) bool {
+	return b.resourceStatusReader.Supports(gk)
+}
+
 // ReadStatus reads the object identified by the passed-in identifier and computes it's status. It reads
 // the resource here, but computing status is delegated to the ReadStatusForObject function.
 func (b *baseStatusReader) ReadStatus(ctx context.Context, reader engine.ClusterReader, identifier object.ObjMetadata) *event.ResourceStatus {

pkg/kstatus/polling/statusreaders/deployment.go

@@ -9,6 +9,7 @@ import (
 	appsv1 "k8s.io/api/apps/v1"
 	"k8s.io/apimachinery/pkg/api/meta"
 	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
+	"k8s.io/apimachinery/pkg/runtime/schema"
 	"sigs.k8s.io/cli-utils/pkg/kstatus/polling/engine"
 	"sigs.k8s.io/cli-utils/pkg/kstatus/polling/event"
 	"sigs.k8s.io/cli-utils/pkg/kstatus/status"
@@ -38,6 +39,10 @@ type deploymentResourceReader struct {
 
 var _ resourceTypeStatusReader = &deploymentResourceReader{}
 
+func (d *deploymentResourceReader) Supports(gk schema.GroupKind) bool {
+	return gk == appsv1.SchemeGroupVersion.WithKind("Deployment").GroupKind()
+}
+
 func (d *deploymentResourceReader) ReadStatusForObject(ctx context.Context, reader engine.ClusterReader, deployment *unstructured.Unstructured) *event.ResourceStatus {
 	identifier := object.UnstructuredToObjMetaOrDie(deployment)
 

pkg/kstatus/polling/statusreaders/generic.go

@@ -8,6 +8,7 @@ import (
 
 	"k8s.io/apimachinery/pkg/api/meta"
 	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
+	"k8s.io/apimachinery/pkg/runtime/schema"
 	"sigs.k8s.io/cli-utils/pkg/kstatus/polling/engine"
 	"sigs.k8s.io/cli-utils/pkg/kstatus/polling/event"
 	"sigs.k8s.io/cli-utils/pkg/kstatus/status"
@@ -43,6 +44,10 @@ type genericStatusReader struct {
 
 var _ resourceTypeStatusReader = &genericStatusReader{}
 
+func (g *genericStatusReader) Supports(schema.GroupKind) bool {
+	return true
+}
+
 func (g *genericStatusReader) ReadStatusForObject(_ context.Context, _ engine.ClusterReader, resource *unstructured.Unstructured) *event.ResourceStatus {
 	identifier := object.UnstructuredToObjMetaOrDie(resource)
 

pkg/kstatus/polling/statusreaders/replicaset.go

@@ -6,8 +6,10 @@ package statusreaders
 import (
 	"context"
 
+	appsv1 "k8s.io/api/apps/v1"
 	"k8s.io/apimachinery/pkg/api/meta"
 	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
+	"k8s.io/apimachinery/pkg/runtime/schema"
 	"sigs.k8s.io/cli-utils/pkg/kstatus/polling/engine"
 	"sigs.k8s.io/cli-utils/pkg/kstatus/polling/event"
 )
@@ -33,6 +35,10 @@ type replicaSetStatusReader struct {
 
 var _ resourceTypeStatusReader = &replicaSetStatusReader{}
 
+func (r *replicaSetStatusReader) Supports(gk schema.GroupKind) bool {
+	return gk == appsv1.SchemeGroupVersion.WithKind("ReplicaSet").GroupKind()
+}
+
 func (r *replicaSetStatusReader) ReadStatusForObject(ctx context.Context, reader engine.ClusterReader, rs *unstructured.Unstructured) *event.ResourceStatus {
 	return newPodControllerStatusReader(r.mapper, r.podStatusReader).readStatus(ctx, reader, rs)
 }

pkg/kstatus/polling/statusreaders/statefulset.go

@@ -6,8 +6,10 @@ package statusreaders
 import (
 	"context"
 
+	appsv1 "k8s.io/api/apps/v1"
 	"k8s.io/apimachinery/pkg/api/meta"
 	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
+	"k8s.io/apimachinery/pkg/runtime/schema"
 	"sigs.k8s.io/cli-utils/pkg/kstatus/polling/engine"
 	"sigs.k8s.io/cli-utils/pkg/kstatus/polling/event"
 )
@@ -33,6 +35,10 @@ type statefulSetResourceReader struct {
 
 var _ resourceTypeStatusReader = &statefulSetResourceReader{}
 
+func (s *statefulSetResourceReader) Supports(gk schema.GroupKind) bool {
+	return gk == appsv1.SchemeGroupVersion.WithKind("StatefulSet").GroupKind()
+}
+
 func (s *statefulSetResourceReader) ReadStatusForObject(ctx context.Context, reader engine.ClusterReader, statefulSet *unstructured.Unstructured) *event.ResourceStatus {
 	return newPodControllerStatusReader(s.mapper, s.podResourceReader).readStatus(ctx, reader, statefulSet)
 }

pkg/kstatus/polling/statusreaders/testing.go

@@ -43,6 +43,10 @@ func (f *fakeClusterReader) ListNamespaceScoped(_ context.Context, list *unstruc
 
 type fakeStatusReader struct{}
 
+func (f *fakeStatusReader) Supports(schema.GroupKind) bool {
+	return true
+}
+
 func (f *fakeStatusReader) ReadStatus(_ context.Context, _ engine.ClusterReader, _ object.ObjMetadata) *event.ResourceStatus {
 	return nil
 }

pkg/util/factory/statuspoller.go

@@ -6,7 +6,6 @@ package factory
 import (
 	"fmt"
 
-	"k8s.io/apimachinery/pkg/runtime/schema"
 	cmdutil "k8s.io/kubectl/pkg/cmd/util"
 	"k8s.io/kubectl/pkg/scheme"
 	"sigs.k8s.io/cli-utils/pkg/kstatus/polling"
@@ -32,5 +31,5 @@ func NewStatusPoller(f cmdutil.Factory) (*polling.StatusPoller, error) {
 		return nil, fmt.Errorf("error creating client: %w", err)
 	}
 
-	return polling.NewStatusPoller(c, mapper, map[schema.GroupKind]engine.StatusReader{}), nil
+	return polling.NewStatusPoller(c, mapper, []engine.StatusReader{}), nil
 }