refactor(workload): centralize workload listing with registry-based listers and add Argo Rollouts support

2026-08-20 12:36:26 +00:00 · 2025-12-28 08:47:56 +01:00
parent c19058a66e
commit 3cf0119748
10 changed files with 808 additions and 123 deletions
@@ -3,6 +3,7 @@ package workload
 import (
 	"context"

+	argorolloutv1alpha1 "github.com/argoproj/argo-rollouts/pkg/apis/rollouts/v1alpha1"
 	appsv1 "k8s.io/api/apps/v1"
 	batchv1 "k8s.io/api/batch/v1"
 	"sigs.k8s.io/controller-runtime/pkg/client"
@@ -49,63 +50,81 @@ func (l *Lister) List(ctx context.Context, namespace string) ([]WorkloadAccessor
 }

 func (l *Lister) listByKind(ctx context.Context, namespace string, kind Kind) ([]WorkloadAccessor, error) {
-	switch kind {
-	case KindDeployment:
-		var list appsv1.DeploymentList
-		if err := l.Client.List(ctx, &list, client.InNamespace(namespace)); err != nil {
-			return nil, err
-		}
-		result := make([]WorkloadAccessor, len(list.Items))
-		for i := range list.Items {
-			result[i] = NewDeploymentWorkload(&list.Items[i])
-		}
-		return result, nil
-
-	case KindDaemonSet:
-		var list appsv1.DaemonSetList
-		if err := l.Client.List(ctx, &list, client.InNamespace(namespace)); err != nil {
-			return nil, err
-		}
-		result := make([]WorkloadAccessor, len(list.Items))
-		for i := range list.Items {
-			result[i] = NewDaemonSetWorkload(&list.Items[i])
-		}
-		return result, nil
-
-	case KindStatefulSet:
-		var list appsv1.StatefulSetList
-		if err := l.Client.List(ctx, &list, client.InNamespace(namespace)); err != nil {
-			return nil, err
-		}
-		result := make([]WorkloadAccessor, len(list.Items))
-		for i := range list.Items {
-			result[i] = NewStatefulSetWorkload(&list.Items[i])
-		}
-		return result, nil
-
-	case KindJob:
-		var list batchv1.JobList
-		if err := l.Client.List(ctx, &list, client.InNamespace(namespace)); err != nil {
-			return nil, err
-		}
-		result := make([]WorkloadAccessor, len(list.Items))
-		for i := range list.Items {
-			result[i] = NewJobWorkload(&list.Items[i])
-		}
-		return result, nil
-
-	case KindCronJob:
-		var list batchv1.CronJobList
-		if err := l.Client.List(ctx, &list, client.InNamespace(namespace)); err != nil {
-			return nil, err
-		}
-		result := make([]WorkloadAccessor, len(list.Items))
-		for i := range list.Items {
-			result[i] = NewCronJobWorkload(&list.Items[i])
-		}
-		return result, nil
-
-	default:
+	lister := l.Registry.ListerFor(kind)
+	if lister == nil {
 		return nil, nil
 	}
+	return lister(ctx, l.Client, namespace)
+}
+
+func listDeployments(ctx context.Context, c client.Client, namespace string) ([]WorkloadAccessor, error) {
+	var list appsv1.DeploymentList
+	if err := c.List(ctx, &list, client.InNamespace(namespace)); err != nil {
+		return nil, err
+	}
+	result := make([]WorkloadAccessor, len(list.Items))
+	for i := range list.Items {
+		result[i] = NewDeploymentWorkload(&list.Items[i])
+	}
+	return result, nil
+}
+
+func listDaemonSets(ctx context.Context, c client.Client, namespace string) ([]WorkloadAccessor, error) {
+	var list appsv1.DaemonSetList
+	if err := c.List(ctx, &list, client.InNamespace(namespace)); err != nil {
+		return nil, err
+	}
+	result := make([]WorkloadAccessor, len(list.Items))
+	for i := range list.Items {
+		result[i] = NewDaemonSetWorkload(&list.Items[i])
+	}
+	return result, nil
+}
+
+func listStatefulSets(ctx context.Context, c client.Client, namespace string) ([]WorkloadAccessor, error) {
+	var list appsv1.StatefulSetList
+	if err := c.List(ctx, &list, client.InNamespace(namespace)); err != nil {
+		return nil, err
+	}
+	result := make([]WorkloadAccessor, len(list.Items))
+	for i := range list.Items {
+		result[i] = NewStatefulSetWorkload(&list.Items[i])
+	}
+	return result, nil
+}
+
+func listJobs(ctx context.Context, c client.Client, namespace string) ([]WorkloadAccessor, error) {
+	var list batchv1.JobList
+	if err := c.List(ctx, &list, client.InNamespace(namespace)); err != nil {
+		return nil, err
+	}
+	result := make([]WorkloadAccessor, len(list.Items))
+	for i := range list.Items {
+		result[i] = NewJobWorkload(&list.Items[i])
+	}
+	return result, nil
+}
+
+func listCronJobs(ctx context.Context, c client.Client, namespace string) ([]WorkloadAccessor, error) {
+	var list batchv1.CronJobList
+	if err := c.List(ctx, &list, client.InNamespace(namespace)); err != nil {
+		return nil, err
+	}
+	result := make([]WorkloadAccessor, len(list.Items))
+	for i := range list.Items {
+		result[i] = NewCronJobWorkload(&list.Items[i])
+	}
+	return result, nil
+}
+
+func listRollouts(ctx context.Context, c client.Client, namespace string) ([]WorkloadAccessor, error) {
+	var list argorolloutv1alpha1.RolloutList
+	if err := c.List(ctx, &list, client.InNamespace(namespace)); err != nil {
+		return nil, err
+	}
+	result := make([]WorkloadAccessor, len(list.Items))
+	for i := range list.Items {
+		result[i] = NewRolloutWorkload(&list.Items[i])
+	}
+	return result, nil
 }
@@ -1,7 +1,9 @@
 package workload

 import (
+	"context"
 	"fmt"
+	"strings"

 	argorolloutv1alpha1 "github.com/argoproj/argo-rollouts/pkg/apis/rollouts/v1alpha1"
 	appsv1 "k8s.io/api/apps/v1"
@@ -9,16 +11,36 @@ import (
 	"sigs.k8s.io/controller-runtime/pkg/client"
 )

+// WorkloadLister is a function that lists workloads of a specific kind.
+type WorkloadLister func(ctx context.Context, c client.Client, namespace string) ([]WorkloadAccessor, error)
+
 // Registry provides factory methods for creating Workload instances.
 type Registry struct {
 	argoRolloutsEnabled bool
+	listers             map[Kind]WorkloadLister
 }

 // NewRegistry creates a new workload registry.
 func NewRegistry(argoRolloutsEnabled bool) *Registry {
-	return &Registry{
+	r := &Registry{
 		argoRolloutsEnabled: argoRolloutsEnabled,
+		listers: map[Kind]WorkloadLister{
+			KindDeployment:  listDeployments,
+			KindDaemonSet:   listDaemonSets,
+			KindStatefulSet: listStatefulSets,
+			KindJob:         listJobs,
+			KindCronJob:     listCronJobs,
+		},
 	}
+	if argoRolloutsEnabled {
+		r.listers[KindArgoRollout] = listRollouts
+	}
+	return r
+}
+
+// ListerFor returns the lister function for the given kind, or nil if not found.
+func (r *Registry) ListerFor(kind Kind) WorkloadLister {
+	return r.listers[kind]
 }

 // SupportedKinds returns all supported workload kinds.
@@ -59,22 +81,27 @@ func (r *Registry) FromObject(obj client.Object) (WorkloadAccessor, error) {
 	}
 }

+// kindAliases maps string representations to Kind constants.
+// Supports lowercase, title case, and plural forms for user convenience.
+var kindAliases = map[string]Kind{
+	"deployment":  KindDeployment,
+	"deployments": KindDeployment,
+	"daemonset":   KindDaemonSet,
+	"daemonsets":  KindDaemonSet,
+	"statefulset": KindStatefulSet,
+	"statefulsets": KindStatefulSet,
+	"rollout":     KindArgoRollout,
+	"rollouts":    KindArgoRollout,
+	"job":         KindJob,
+	"jobs":        KindJob,
+	"cronjob":     KindCronJob,
+	"cronjobs":    KindCronJob,
+}
+
 // KindFromString converts a string to a Kind.
 func KindFromString(s string) (Kind, error) {
-	switch s {
-	case "Deployment", "deployment", "deployments":
-		return KindDeployment, nil
-	case "DaemonSet", "daemonset", "daemonsets":
-		return KindDaemonSet, nil
-	case "StatefulSet", "statefulset", "statefulsets":
-		return KindStatefulSet, nil
-	case "Rollout", "rollout", "rollouts":
-		return KindArgoRollout, nil
-	case "Job", "job", "jobs":
-		return KindJob, nil
-	case "CronJob", "cronjob", "cronjobs":
-		return KindCronJob, nil
-	default:
-		return "", fmt.Errorf("unknown workload kind: %s", s)
+	if k, ok := kindAliases[strings.ToLower(s)]; ok {
+		return k, nil
 	}
+	return "", fmt.Errorf("unknown workload kind: %s", s)
 }