Add dynamic completion for a resource query (#209)

README.md

@@ -234,8 +234,8 @@ stern -p
 Stern supports command-line auto completion for bash, zsh or fish. `stern
 --completion=(bash|zsh|fish)` outputs the shell completion code which work by being
 evaluated in `.bashrc`, etc for the specified shell. In addition, Stern
-supports dynamic completion for `--namespace` and `--context`. In order to use
-that, kubectl must be installed on your environment.
+supports dynamic completion for `--namespace`, `--context` and a resource query
+in the form `<resource>/<name>`.
 
 If you use bash, stern bash completion code depends on the
 [bash-completion](https://github.com/scop/bash-completion). On the macOS, you

cmd/cmd.go

@@ -389,6 +389,7 @@ func NewSternCmd(stream genericclioptions.IOStreams) (*cobra.Command, error) {
 
 			return o.Run(cmd)
 		},
+		ValidArgsFunction: queryCompletionFunc(o),
 	}
 
 	o.AddFlags(cmd.Flags())

cmd/flag_completion.go

@@ -21,9 +21,12 @@ import (
 	"io"
 	"strings"
 
+	"github.com/pkg/errors"
 	"github.com/spf13/cobra"
 	"github.com/stern/stern/kubernetes"
+	"github.com/stern/stern/stern"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	clientset "k8s.io/client-go/kubernetes"
 )
 
 func runCompletion(shell string, cmd *cobra.Command, out io.Writer) error {
@@ -119,7 +122,137 @@ func contextCompletionFunc(o *options) func(cmd *cobra.Command, args []string, t
 	}
 }
 
+// queryCompletionFunc is a completion function that completes a resource
+// that match the toComplete prefix.
+func queryCompletionFunc(o *options) func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
+	return func(_ *cobra.Command, _ []string, toComplete string) ([]string, cobra.ShellCompDirective) {
+		var comps []string
+		parts := strings.Split(toComplete, "/")
+		if len(parts) != 2 {
+			// list available resources in the form "<resource>/"
+			for _, matcher := range stern.ResourceMatchers {
+				if strings.HasPrefix(matcher.Name(), toComplete) {
+					comps = append(comps, matcher.Name()+"/")
+				}
+			}
+			return comps, cobra.ShellCompDirectiveNoFileComp | cobra.ShellCompDirectiveNoSpace
+		}
+
+		// list available names in the resources in the form "<resource>/<name>"
+		uniqueNamespaces := makeUnique(o.namespaces)
+		if o.allNamespaces || len(uniqueNamespaces) > 1 {
+			// do not support multiple namespaces for simplicity
+			return compError(errors.New("multiple namespaces are not supported"))
+		}
+
+		clientConfig := kubernetes.NewClientConfig(o.kubeConfig, o.context)
+		clientset, err := kubernetes.NewClientSet(clientConfig)
+		if err != nil {
+			return compError(err)
+		}
+		var namespace string
+		if len(uniqueNamespaces) == 1 {
+			namespace = uniqueNamespaces[0]
+		} else {
+			n, _, err := clientConfig.Namespace()
+			if err != nil {
+				return compError(err)
+			}
+			namespace = n
+		}
+
+		kind, name := parts[0], parts[1]
+		names, err := retrieveNamesFromResource(context.TODO(), clientset, namespace, kind)
+		if err != nil {
+			return compError(err)
+		}
+		for _, n := range names {
+			if strings.HasPrefix(n, name) {
+				comps = append(comps, kind+"/"+n)
+			}
+		}
+		return comps, cobra.ShellCompDirectiveNoFileComp
+	}
+}
+
 func compError(err error) ([]string, cobra.ShellCompDirective) {
 	cobra.CompError(err.Error())
 	return nil, cobra.ShellCompDirectiveError
 }
+
+func retrieveNamesFromResource(ctx context.Context, client clientset.Interface, namespace, kind string) ([]string, error) {
+	opt := metav1.ListOptions{}
+	var names []string
+	switch {
+	// core
+	case stern.PodMatcher.Matches(kind):
+		l, err := client.CoreV1().Pods(namespace).List(ctx, opt)
+		if err != nil {
+			return nil, err
+		}
+		for _, item := range l.Items {
+			names = append(names, item.GetName())
+		}
+	case stern.ReplicationControllerMatcher.Matches(kind):
+		l, err := client.CoreV1().ReplicationControllers(namespace).List(ctx, opt)
+		if err != nil {
+			return nil, err
+		}
+		for _, item := range l.Items {
+			names = append(names, item.GetName())
+		}
+	case stern.ServiceMatcher.Matches(kind):
+		l, err := client.CoreV1().Services(namespace).List(ctx, opt)
+		if err != nil {
+			return nil, err
+		}
+		for _, item := range l.Items {
+			names = append(names, item.GetName())
+		}
+	// apps
+	case stern.DeploymentMatcher.Matches(kind):
+		l, err := client.AppsV1().Deployments(namespace).List(ctx, opt)
+		if err != nil {
+			return nil, err
+		}
+		for _, item := range l.Items {
+			names = append(names, item.GetName())
+		}
+	case stern.DaemonSetMatcher.Matches(kind):
+		l, err := client.AppsV1().DaemonSets(namespace).List(ctx, opt)
+		if err != nil {
+			return nil, err
+		}
+		for _, item := range l.Items {
+			names = append(names, item.GetName())
+		}
+	case stern.ReplicaSetMatcher.Matches(kind):
+		l, err := client.AppsV1().ReplicaSets(namespace).List(ctx, opt)
+		if err != nil {
+			return nil, err
+		}
+		for _, item := range l.Items {
+			names = append(names, item.GetName())
+		}
+	case stern.StatefulSetMatcher.Matches(kind):
+		l, err := client.AppsV1().StatefulSets(namespace).List(ctx, opt)
+		if err != nil {
+			return nil, err
+		}
+		for _, item := range l.Items {
+			names = append(names, item.GetName())
+		}
+	// batch
+	case stern.JobMatcher.Matches(kind):
+		l, err := client.BatchV1().Jobs(namespace).List(ctx, opt)
+		if err != nil {
+			return nil, err
+		}
+		for _, item := range l.Items {
+			names = append(names, item.GetName())
+		}
+	default:
+		return nil, fmt.Errorf("resource type %s is not supported", kind)
+	}
+	return names, nil
+}

cmd/flag_completion_test.go

@@ -0,0 +1,123 @@
+package cmd
+
+import (
+	"context"
+	"reflect"
+	"testing"
+
+	appsv1 "k8s.io/api/apps/v1"
+	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/client-go/kubernetes/fake"
+)
+
+func TestRetrieveNamesFromResource(t *testing.T) {
+	genMeta := func(name string) metav1.ObjectMeta {
+		return metav1.ObjectMeta{
+			Name:      name,
+			Namespace: "ns1",
+		}
+	}
+	objs := []runtime.Object{
+		&corev1.Pod{ObjectMeta: genMeta("pod1")},
+		&corev1.Pod{ObjectMeta: genMeta("pod2")},
+		&corev1.Pod{ObjectMeta: genMeta("pod3")},
+		&corev1.ReplicationController{ObjectMeta: genMeta("rc1")},
+		&corev1.Service{ObjectMeta: genMeta("svc1")},
+		&appsv1.Deployment{ObjectMeta: genMeta("deploy1")},
+		&appsv1.Deployment{ObjectMeta: genMeta("deploy2")},
+		&appsv1.DaemonSet{ObjectMeta: genMeta("ds1")},
+		&appsv1.DaemonSet{ObjectMeta: genMeta("ds2")},
+		&appsv1.ReplicaSet{ObjectMeta: genMeta("rs1")},
+		&appsv1.ReplicaSet{ObjectMeta: genMeta("rs2")},
+		&appsv1.StatefulSet{ObjectMeta: genMeta("sts1")},
+		&appsv1.StatefulSet{ObjectMeta: genMeta("sts2")},
+		&batchv1.Job{ObjectMeta: genMeta("job1")},
+		&batchv1.Job{ObjectMeta: genMeta("job2")},
+	}
+	client := fake.NewSimpleClientset(objs...)
+	tests := []struct {
+		desc      string
+		kinds     []string
+		expected  []string
+		wantError bool
+	}{
+		// core
+		{
+			desc:     "pods",
+			kinds:    []string{"po", "pods", "pod"},
+			expected: []string{"pod1", "pod2", "pod3"},
+		},
+		{
+			desc:     "replicationcontrollers",
+			kinds:    []string{"rc", "replicationcontrollers", "replicationcontroller"},
+			expected: []string{"rc1"},
+		},
+		// apps
+		{
+			desc:     "deployments",
+			kinds:    []string{"deploy", "deployments", "deployment"},
+			expected: []string{"deploy1", "deploy2"},
+		},
+		{
+			desc:     "daemonsets",
+			kinds:    []string{"ds", "daemonsets", "daemonset"},
+			expected: []string{"ds1", "ds2"},
+		},
+		{
+			desc:     "replicasets",
+			kinds:    []string{"rs", "replicasets", "replicaset"},
+			expected: []string{"rs1", "rs2"},
+		},
+		{
+			desc:     "statefulsets",
+			kinds:    []string{"sts", "statefulsets", "statefulset"},
+			expected: []string{"sts1", "sts2"},
+		},
+		// batch
+		{
+			desc:     "jobs",
+			kinds:    []string{"job", "jobs"},
+			expected: []string{"job1", "job2"},
+		},
+		// invalid
+		{
+			desc:      "invalid",
+			kinds:     []string{"", "unknown"},
+			wantError: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.desc, func(t *testing.T) {
+			for _, kind := range tt.kinds {
+				names, err := retrieveNamesFromResource(context.Background(), client, "ns1", kind)
+				if tt.wantError {
+					if err == nil {
+						t.Errorf("expected error, but got no error")
+					}
+					return
+				}
+				if err != nil {
+					t.Errorf("unexpected error: %v", err)
+					return
+				}
+				if !reflect.DeepEqual(tt.expected, names) {
+					t.Errorf("expected %v, but actual %v", tt.expected, names)
+				}
+				// expect empty slice with no error when no objects are found in the valid resource
+				names, err = retrieveNamesFromResource(context.Background(), client, "not-matched", kind)
+				if err != nil {
+					t.Errorf("unexpected error: %v", err)
+					return
+				}
+				if len(names) != 0 {
+					t.Errorf("expected empty slice, but got %v", names)
+					return
+				}
+			}
+		})
+	}
+}

stern/resource_matcher.go

@@ -49,4 +49,14 @@ var (
 	ReplicaSetMatcher            = ResourceMatcher{name: "replicaset", aliases: []string{"rs", "replicasets"}}
 	StatefulSetMatcher           = ResourceMatcher{name: "statefulset", aliases: []string{"sts", "statefulsets"}}
 	JobMatcher                   = ResourceMatcher{name: "job", aliases: []string{"jobs"}} // job does not have a short name
+	ResourceMatchers             = []ResourceMatcher{
+		PodMatcher,
+		ReplicationControllerMatcher,
+		ServiceMatcher,
+		DeploymentMatcher,
+		DaemonSetMatcher,
+		ReplicaSetMatcher,
+		StatefulSetMatcher,
+		JobMatcher,
+	}
 )