From eb80584991ec370f454f28eda30ceace7d5ded69 Mon Sep 17 00:00:00 2001 From: Kavin <62742678+Kavinraja-G@users.noreply.github.com> Date: Tue, 28 May 2024 10:28:42 +0530 Subject: [PATCH] refactor get pods logic for multiple types --- cmd/kueuectl/app/list/helpers.go | 4 +- cmd/kueuectl/app/list/list_pods.go | 116 ++++++------ cmd/kueuectl/app/list/list_pods_test.go | 223 ++++-------------------- 3 files changed, 96 insertions(+), 247 deletions(-) diff --git a/cmd/kueuectl/app/list/helpers.go b/cmd/kueuectl/app/list/helpers.go index 7b0f9c5b08..bfd8ae8e25 100644 --- a/cmd/kueuectl/app/list/helpers.go +++ b/cmd/kueuectl/app/list/helpers.go @@ -50,7 +50,7 @@ func addActiveFilterFlagVar(cmd *cobra.Command, p *[]bool) { "Filter by active status. Valid values: 'true' and 'false'.") } -func addForFilterFlagVar(cmd *cobra.Command, p *string) { +func addForObjectFilterFlagVar(cmd *cobra.Command, p *string) { cmd.Flags().StringVar(p, "for", "", - "List pods belongs to a particular Job kind.") + "Filter pods to only those pertaining to the specified resource types.") } diff --git a/cmd/kueuectl/app/list/list_pods.go b/cmd/kueuectl/app/list/list_pods.go index f53e19c863..4b0bc912c2 100644 --- a/cmd/kueuectl/app/list/list_pods.go +++ b/cmd/kueuectl/app/list/list_pods.go @@ -21,6 +21,7 @@ import ( "errors" "fmt" "io" + "k8s.io/apimachinery/pkg/runtime/schema" "strings" "time" @@ -44,17 +45,23 @@ const ( the label selector or the field selector.` podExample = ` # List Pods kueuectl list pods --for job/job-name` - jobControllerUIDLabel = "batch.kubernetes.io/controller-uid" ) +type objectRef struct { + APIGroup string + Kind string + Name string +} + type PodOptions struct { PrintFlags *genericclioptions.PrintFlags - AllNamespaces bool - Namespace string - LabelSelector string - FieldSelector string - JobArg string + AllNamespaces bool + Namespace string + LabelSelector string + FieldSelector string + UserSpecifiedForObject string + ForObject objectRef Clientset k8s.Interface @@ -74,15 +81,14 @@ func NewPodCmd(clientGetter util.ClientGetter, streams genericiooptions.IOStream o := NewPodOptions(streams) cmd := &cobra.Command{ - Use: "pods --for [.]/]", + Use: "pods --for [[.]/]", DisableFlagsInUseLine: true, Aliases: []string{"po"}, Short: "List Pods belong to a Job Kind", Long: podLong, Example: podExample, Run: func(cmd *cobra.Command, args []string) { - cobra.CheckErr(o.Complete(clientGetter, cmd, args)) - cobra.CheckErr(o.Validate()) + cobra.CheckErr(o.Complete(clientGetter)) cobra.CheckErr(o.Run(cmd.Context())) }, } @@ -92,13 +98,13 @@ func NewPodCmd(clientGetter util.ClientGetter, streams genericiooptions.IOStream addAllNamespacesFlagVar(cmd, &o.AllNamespaces) addFieldSelectorFlagVar(cmd, &o.FieldSelector) addLabelSelectorFlagVar(cmd, &o.LabelSelector) - addForFilterFlagVar(cmd, &o.JobArg) + addForObjectFilterFlagVar(cmd, &o.UserSpecifiedForObject) return cmd } // Complete takes the command arguments and infers any remaining options. -func (o *PodOptions) Complete(clientGetter util.ClientGetter, cmd *cobra.Command, args []string) error { +func (o *PodOptions) Complete(clientGetter util.ClientGetter) error { var err error o.Namespace, _, err = clientGetter.ToRawKubeConfigLoader().Namespace() @@ -121,74 +127,49 @@ func (o *PodOptions) Complete(clientGetter util.ClientGetter, cmd *cobra.Command o.PrintObj = printer.PrintObj } - if len(args) > 0 { - jobArg, err := cmd.Flags().GetString("for") - if err != nil { - return err - } - o.JobArg = jobArg + err = o.parseForObject() + if err != nil { + return err } return nil } -func (o *PodOptions) Validate() error { - if !o.validJobFlagOptionProvided() { - return errors.New("not a valid --job flag. Please provide a valid job name in format job/job-name") - } - return nil -} - -func (o *PodOptions) validJobFlagOptionProvided() bool { - jobArgSlice := strings.Split(o.JobArg, "/") - - // jobArgSlice should be ["job", "job-name"] - if len(jobArgSlice) != 2 { - return false - } - - // valid only if the kind is a job - kind := strings.ToLower(jobArgSlice[0]) - return strings.Contains(kind, "job") -} - // Run prints the pods for a specific Job func (o *PodOptions) Run(ctx context.Context) error { - jobName := strings.Split(o.JobArg, "/")[1] - controllerUID, err := o.getJobControllerUID(ctx, jobName) + err := o.listPods(ctx) if err != nil { return err } - err = o.listPodsByControllerUID(ctx, controllerUID) - if err != nil { - return err - } return nil } -// getJobControllerUID fetches the controllerUID of the given Job -func (o *PodOptions) getJobControllerUID(ctx context.Context, jobName string) (string, error) { - job, err := o.Clientset.BatchV1().Jobs(o.Namespace).Get(ctx, jobName, metav1.GetOptions{}) - if err != nil { - return "", err +func (o *PodOptions) parseForObject() error { + if o.UserSpecifiedForObject == "" { + return nil + } + + parts := strings.Split(o.UserSpecifiedForObject, "/") + if len(parts) > 2 { + return errors.New(fmt.Sprintf("invalid value '%s' used in --for flag; value must be in the format [TYPE[.API-GROUP]/]NAME", o.UserSpecifiedForObject)) + } + + if len(parts) == 1 { + o.ForObject.Name = parts[0] + return nil } - return job.ObjectMeta.Labels[jobControllerUIDLabel], nil + o.ForObject.Name = parts[1] + o.ForObject.Kind, o.ForObject.APIGroup, _ = strings.Cut(parts[0], ".") + + return nil } -// listPodsByControllerUID lists the pods based on the given controllerUID linked to the pod -func (o *PodOptions) listPodsByControllerUID(ctx context.Context, controllerUID string) error { +// listPods lists the pods based on the given --for object +func (o *PodOptions) listPods(ctx context.Context) error { continueToken := "" - // assign or appends controllerUID label with the existing selector if any - // used for filtering the pods which belongs to the job controller - if o.LabelSelector != "" { - o.LabelSelector = fmt.Sprintf("%s,%s=%s", o.LabelSelector, jobControllerUIDLabel, controllerUID) - } else { - o.LabelSelector = fmt.Sprintf("%s=%s", jobControllerUIDLabel, controllerUID) - } - for { podList, err := o.Clientset.CoreV1().Pods(o.Namespace).List(ctx, metav1.ListOptions{ LabelSelector: o.LabelSelector, @@ -203,6 +184,23 @@ func (o *PodOptions) listPodsByControllerUID(ctx context.Context, controllerUID return nil } + filteredPods := make([]corev1.Pod, 0, len(podList.Items)) + for i := range podList.Items { + for _, ownerRef := range podList.Items[i].OwnerReferences { + gv, _ := schema.ParseGroupVersion(ownerRef.APIVersion) + + if strings.EqualFold(o.ForObject.Kind, ownerRef.Kind) && strings.EqualFold(o.ForObject.Name, ownerRef.Name) { + if o.ForObject.APIGroup == "" || strings.EqualFold(o.ForObject.APIGroup, gv.Group) { + filteredPods = append(filteredPods, podList.Items[i]) + break + } + } + } + } + + // replace the podList items with the new filtered pods + podList.Items = filteredPods + if err := o.PrintObj(podList, o.Out); err != nil { return err } diff --git a/cmd/kueuectl/app/list/list_pods_test.go b/cmd/kueuectl/app/list/list_pods_test.go index b116857eea..7406b19f4f 100644 --- a/cmd/kueuectl/app/list/list_pods_test.go +++ b/cmd/kueuectl/app/list/list_pods_test.go @@ -25,163 +25,50 @@ import ( "github.com/google/go-cmp/cmp" "github.com/google/go-cmp/cmp/cmpopts" - batchv1 "k8s.io/api/batch/v1" corev1 "k8s.io/api/core/v1" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/runtime" "k8s.io/cli-runtime/pkg/genericclioptions" "k8s.io/cli-runtime/pkg/genericiooptions" - "k8s.io/cli-runtime/pkg/printers" - k8s "k8s.io/client-go/kubernetes" "k8s.io/client-go/kubernetes/fake" ) -func TestPodOptions_getJobControllerUID(t *testing.T) { - type fields struct { - PrintFlags *genericclioptions.PrintFlags - AllNamespaces bool - Namespace string - LabelSelector string - FieldSelector string - JobArg string - Clientset k8s.Clientset - PrintObj printers.ResourcePrinterFunc - IOStreams genericiooptions.IOStreams - } - type args struct { - jobName string - } - tests := []struct { - name string - jobs []runtime.Object - fields fields - args args - want string - wantErr bool - }{ - { - name: "get valid controllerUID from a Job", - jobs: []runtime.Object{ - &batchv1.Job{ - ObjectMeta: metav1.ObjectMeta{ - Name: "test-job", - Namespace: "default", - Labels: map[string]string{ - jobControllerUIDLabel: "test-job-uid", - }, - }, - }, - }, - args: args{ - jobName: "test-job", - }, - fields: fields{ - Namespace: "default", - }, - want: "test-job-uid", - wantErr: false, - }, { - name: "missing controllerUID label from a Job", - jobs: []runtime.Object{ - &batchv1.Job{ - ObjectMeta: metav1.ObjectMeta{ - Name: "missing-uid-job", - Namespace: "default", - Labels: nil, - }, - }, - }, - fields: fields{ - Namespace: "default", - }, - args: args{ - jobName: "missing-uid-job", - }, - want: "", - wantErr: false, - }, { - name: "job not found in the namespace", - jobs: []runtime.Object{ - &batchv1.Job{ - ObjectMeta: metav1.ObjectMeta{ - Name: "test-job", - Namespace: "test-namespace", - Labels: nil, - }, - }, - }, - fields: fields{ - Namespace: "default", - }, - args: args{ - jobName: "test-job", - }, - want: "", - wantErr: true, - }, - } - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - t.Parallel() - fakeClientset := fake.NewSimpleClientset(tt.jobs...) - o := &PodOptions{ - PrintFlags: tt.fields.PrintFlags, - AllNamespaces: tt.fields.AllNamespaces, - Namespace: tt.fields.Namespace, - LabelSelector: tt.fields.LabelSelector, - FieldSelector: tt.fields.FieldSelector, - JobArg: tt.fields.JobArg, - Clientset: fakeClientset, - PrintObj: tt.fields.PrintObj, - IOStreams: tt.fields.IOStreams, - } - got, err := o.getJobControllerUID(context.TODO(), tt.args.jobName) - if (err != nil) != tt.wantErr { - t.Errorf("getJobControllerUID() error = %v, wantErr %v", err, tt.wantErr) - return - } - if got != tt.want { - t.Errorf("getJobControllerUID() got = %v, want %v", got, tt.want) - } - }) - } -} - -func TestPodOptions_listPodsByControllerUID(t *testing.T) { +func TestPodOptions_listPods(t *testing.T) { testStartTime := time.Now() type fields struct { - PrintFlags *genericclioptions.PrintFlags - AllNamespaces bool - Namespace string - LabelSelector string - FieldSelector string - JobArg string - } - type args struct { - controllerUID string + PrintFlags *genericclioptions.PrintFlags + AllNamespaces bool + Namespace string + LabelSelector string + FieldSelector string + UserSpecifiedForObject string + ForObject objectRef } tests := []struct { name string pods []runtime.Object fields fields - args args wantOut []string wantErr error }{ { - name: "list pods by valid controllerUID", + name: "list pods for batch/job type", pods: []runtime.Object{ &corev1.Pod{ ObjectMeta: metav1.ObjectMeta{ Name: "valid-pod-1", Namespace: "default", - Labels: map[string]string{ - jobControllerUIDLabel: "valid-controller-uid", - }, CreationTimestamp: metav1.Time{ Time: testStartTime.Add(-time.Hour).Truncate(time.Second), }, + OwnerReferences: []metav1.OwnerReference{ + { + APIVersion: "batch/v1", + Kind: "Job", + Name: "test-job", + }, + }, }, Status: corev1.PodStatus{ Phase: "RUNNING", @@ -190,37 +77,25 @@ func TestPodOptions_listPodsByControllerUID(t *testing.T) { ObjectMeta: metav1.ObjectMeta{ Name: "valid-pod-2", Namespace: "default", - Labels: map[string]string{ - jobControllerUIDLabel: "valid-controller-uid", - }, CreationTimestamp: metav1.Time{ Time: testStartTime.Add(-time.Hour).Truncate(time.Second), }, - }, - Status: corev1.PodStatus{ - Phase: "COMPLETED", - }, - }, &corev1.Pod{ - ObjectMeta: metav1.ObjectMeta{ - Name: "invalid-pod", - Namespace: "default", - Labels: map[string]string{ - jobControllerUIDLabel: "invalid-controller-uid", - }, - CreationTimestamp: metav1.Time{ - Time: testStartTime.Add(-time.Hour).Truncate(time.Second), + OwnerReferences: []metav1.OwnerReference{ + { + APIVersion: "batch/v1", + Kind: "Job", + Name: "test-job", + }, }, }, Status: corev1.PodStatus{ - Phase: "RUNNING", + Phase: "COMPLETED", }, }, }, fields: fields{ - Namespace: "default", - }, - args: args{ - controllerUID: "valid-controller-uid", + Namespace: "default", + UserSpecifiedForObject: "job/test-job", }, wantOut: []string{ "NAME STATUS AGE", @@ -228,52 +103,28 @@ func TestPodOptions_listPodsByControllerUID(t *testing.T) { "valid-pod-2 COMPLETED 60m", "", }, - }, { - name: "no valid pods matching controllerUID", - pods: []runtime.Object{ - &corev1.Pod{ - ObjectMeta: metav1.ObjectMeta{ - Name: "test-pod-1", - Namespace: "default", - Labels: map[string]string{ - jobControllerUIDLabel: "valid-controller-uid", - }, - CreationTimestamp: metav1.Time{ - Time: testStartTime.Add(-time.Hour).Truncate(time.Second), - }, - }, - Status: corev1.PodStatus{ - Phase: "RUNNING", - }, - }, - }, - fields: fields{ - Namespace: "default", - }, - args: args{ - controllerUID: "invalid-controller-uid", - }, - wantOut: []string{""}, }, } + for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { fakeClientset := fake.NewSimpleClientset(tt.pods...) s, _, out, _ := genericiooptions.NewTestIOStreams() o := &PodOptions{ - PrintFlags: tt.fields.PrintFlags, - AllNamespaces: tt.fields.AllNamespaces, - Namespace: tt.fields.Namespace, - LabelSelector: tt.fields.LabelSelector, - FieldSelector: tt.fields.FieldSelector, - JobArg: tt.fields.JobArg, - Clientset: fakeClientset, - PrintObj: printPodTable, - IOStreams: s, + PrintFlags: tt.fields.PrintFlags, + AllNamespaces: tt.fields.AllNamespaces, + Namespace: tt.fields.Namespace, + LabelSelector: tt.fields.LabelSelector, + FieldSelector: tt.fields.FieldSelector, + UserSpecifiedForObject: tt.fields.UserSpecifiedForObject, + Clientset: fakeClientset, + PrintObj: printPodTable, + IOStreams: s, } - gotErr := o.listPodsByControllerUID(context.TODO(), tt.args.controllerUID) + gotErr := o.parseForObject() + gotErr = o.listPods(context.TODO()) if diff := cmp.Diff(tt.wantErr, gotErr, cmpopts.EquateErrors()); diff != "" { t.Errorf("Unexpected error (-want/+got)\n%s", diff) }