[Feat][kubectl-plugin] Add dynamic shell completion for kubectl ray session (#2390)

Signed-off-by: Chi-Sheng Liu <[email protected]>

Author: MortalHappiness

kubectl-plugin/pkg/cmd/session/session.go

@@ -7,6 +7,7 @@ import (
 
 	"github.com/ray-project/kuberay/kubectl-plugin/pkg/util"
 	"github.com/ray-project/kuberay/kubectl-plugin/pkg/util/client"
+	"github.com/ray-project/kuberay/kubectl-plugin/pkg/util/completion"
 	"github.com/spf13/cobra"
 	"k8s.io/cli-runtime/pkg/genericclioptions"
 	"k8s.io/cli-runtime/pkg/genericiooptions"
@@ -23,7 +24,6 @@ type appPort struct {
 type SessionOptions struct {
 	configFlags  *genericclioptions.ConfigFlags
 	ioStreams    *genericiooptions.IOStreams
-	client       client.Client
 	ResourceType util.ResourceType
 	ResourceName string
 	Namespace    string
@@ -76,20 +76,22 @@ func NewSessionOptions(streams genericiooptions.IOStreams) *SessionOptions {
 
 func NewSessionCommand(streams genericiooptions.IOStreams) *cobra.Command {
 	options := NewSessionOptions(streams)
+	factory := cmdutil.NewFactory(options.configFlags)
 
 	cmd := &cobra.Command{
-		Use:     "session (RAYCLUSTER | TYPE/NAME)",
-		Short:   "Forward local ports to the Ray resources.",
-		Long:    sessionLong,
-		Example: sessionExample,
+		Use:               "session (RAYCLUSTER | TYPE/NAME)",
+		Short:             "Forward local ports to the Ray resources.",
+		Long:              sessionLong,
+		Example:           sessionExample,
+		ValidArgsFunction: completion.RayClusterResourceNameCompletionFunc(factory),
 		RunE: func(cmd *cobra.Command, args []string) error {
 			if err := options.Complete(cmd, args); err != nil {
 				return err
 			}
 			if err := options.Validate(); err != nil {
 				return err
 			}
-			return options.Run(cmd.Context())
+			return options.Run(cmd.Context(), factory)
 		},
 	}
 	options.configFlags.AddFlags(cmd.Flags())
@@ -130,13 +132,6 @@ func (options *SessionOptions) Complete(cmd *cobra.Command, args []string) error
 		options.Namespace = *options.configFlags.Namespace
 	}
 
-	factory := cmdutil.NewFactory(options.configFlags)
-	k8sClient, err := client.NewClient(factory)
-	if err != nil {
-		return fmt.Errorf("failed to create client: %w", err)
-	}
-	options.client = k8sClient
-
 	return nil
 }
 
@@ -152,10 +147,13 @@ func (options *SessionOptions) Validate() error {
 	return nil
 }
 
-func (options *SessionOptions) Run(ctx context.Context) error {
-	factory := cmdutil.NewFactory(options.configFlags)
+func (options *SessionOptions) Run(ctx context.Context, factory cmdutil.Factory) error {
+	k8sClient, err := client.NewClient(factory)
+	if err != nil {
+		return fmt.Errorf("failed to create client: %w", err)
+	}
 
-	svcName, err := options.client.GetRayHeadSvcName(ctx, options.Namespace, options.ResourceType, options.ResourceName)
+	svcName, err := k8sClient.GetRayHeadSvcName(ctx, options.Namespace, options.ResourceType, options.ResourceName)
 	if err != nil {
 		return err
 	}

kubectl-plugin/pkg/util/completion/completion.go

@@ -0,0 +1,104 @@
+package completion
+
+import (
+	"fmt"
+	"strings"
+
+	"github.com/spf13/cobra"
+
+	cmdutil "k8s.io/kubectl/pkg/cmd/util"
+	"k8s.io/kubectl/pkg/util/completion"
+
+	"github.com/ray-project/kuberay/kubectl-plugin/pkg/util"
+)
+
+// RayResourceTypeCompletionFunc Returns a completion function that completes the Ray resource type.
+// That is, raycluster, rayjob, or rayservice.
+func RayResourceTypeCompletionFunc() func(*cobra.Command, []string, string) ([]string, cobra.ShellCompDirective) {
+	return func(_ *cobra.Command, _ []string, toComplete string) ([]string, cobra.ShellCompDirective) {
+		var comps []string
+		directive := cobra.ShellCompDirectiveNoFileComp
+		resourceTypes := getAllRayResourceType()
+		for _, resourceType := range resourceTypes {
+			if strings.HasPrefix(resourceType, toComplete) {
+				comps = append(comps, resourceType)
+			}
+		}
+		return comps, directive
+	}
+}
+
+// RayClusterCompletionFunc Returns a completion function that completes RayCluster resource names.
+func RayClusterCompletionFunc(f cmdutil.Factory) func(*cobra.Command, []string, string) ([]string, cobra.ShellCompDirective) {
+	return completion.ResourceNameCompletionFunc(f, string(util.RayCluster))
+}
+
+// RayJobCompletionFunc Returns a completion function that completes RayJob resource names.
+func RayJobCompletionFunc(f cmdutil.Factory) func(*cobra.Command, []string, string) ([]string, cobra.ShellCompDirective) {
+	return completion.ResourceNameCompletionFunc(f, string(util.RayJob))
+}
+
+// RayServiceCompletionFunc Returns a completion function that completes RayService resource names.
+func RayServiceCompletionFunc(f cmdutil.Factory) func(*cobra.Command, []string, string) ([]string, cobra.ShellCompDirective) {
+	return completion.ResourceNameCompletionFunc(f, string(util.RayService))
+}
+
+// RayClusterResourceNameCompletionFunc Returns completions of:
+// 1- RayCluster names that match the toComplete prefix
+// 2- Ray resource types which match the toComplete prefix
+func RayClusterResourceNameCompletionFunc(f cmdutil.Factory) func(*cobra.Command, []string, string) ([]string, cobra.ShellCompDirective) {
+	return func(_ *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
+		var comps []string
+		directive := cobra.ShellCompDirectiveNoFileComp
+		if len(args) == 0 {
+			comps, directive = doRayClusterCompletion(f, toComplete)
+		}
+		return comps, directive
+	}
+}
+
+func getAllRayResourceType() []string {
+	return []string{
+		string(util.RayCluster),
+		string(util.RayJob),
+		string(util.RayService),
+	}
+}
+
+// doRayClusterCompletion Returns completions of:
+// 1- RayCluster names that match the toComplete prefix
+// 2- Ray resource types which match the toComplete prefix
+// Ref: https://github.com/kubernetes/kubectl/blob/262825a8a665c7cae467dfaa42b63be5a5b8e5a2/pkg/util/completion/completion.go#L434
+func doRayClusterCompletion(f cmdutil.Factory, toComplete string) ([]string, cobra.ShellCompDirective) {
+	var comps []string
+	directive := cobra.ShellCompDirectiveNoFileComp
+	slashIdx := strings.Index(toComplete, "/")
+	if slashIdx == -1 {
+		// Standard case, complete RayCluster names
+		comps = completion.CompGetResource(f, string(util.RayCluster), toComplete)
+
+		// Also include resource choices for the <type>/<name> form
+		resourceTypes := getAllRayResourceType()
+
+		if len(comps) == 0 {
+			// If there are no RayCluster to complete, we will only be completing
+			// <type>/.  We should disable adding a space after the /.
+			directive |= cobra.ShellCompDirectiveNoSpace
+		}
+
+		for _, resource := range resourceTypes {
+			if strings.HasPrefix(resource, toComplete) {
+				comps = append(comps, fmt.Sprintf("%s/", resource))
+			}
+		}
+	} else {
+		// Dealing with the <type>/<name> form, use the specified resource type
+		resourceType := toComplete[:slashIdx]
+		toComplete = toComplete[slashIdx+1:]
+		nameComps := completion.CompGetResource(f, resourceType, toComplete)
+		for _, c := range nameComps {
+			comps = append(comps, fmt.Sprintf("%s/%s", resourceType, c))
+		}
+	}
+	return comps, directive
+}

kubectl-plugin/pkg/util/completion/completion_test.go

@@ -0,0 +1,34 @@
+package completion
+
+import (
+	"sort"
+	"testing"
+
+	"github.com/spf13/cobra"
+)
+
+func TestRayResourceTypeCompletionFunc(t *testing.T) {
+	compFunc := RayResourceTypeCompletionFunc()
+	comps, directive := compFunc(nil, []string{}, "")
+	checkCompletion(t, comps, []string{"raycluster", "rayjob", "rayservice"}, directive, cobra.ShellCompDirectiveNoFileComp)
+}
+
+func checkCompletion(t *testing.T, comps, expectedComps []string, directive, expectedDirective cobra.ShellCompDirective) {
+	if e, d := expectedDirective, directive; e != d {
+		t.Errorf("expected directive\n%v\nbut got\n%v", e, d)
+	}
+
+	sort.Strings(comps)
+	sort.Strings(expectedComps)
+
+	if len(expectedComps) != len(comps) {
+		t.Fatalf("expected completions\n%v\nbut got\n%v", expectedComps, comps)
+	}
+
+	for i := range comps {
+		if expectedComps[i] != comps[i] {
+			t.Errorf("expected completions\n%v\nbut got\n%v", expectedComps, comps)
+			break
+		}
+	}
+}