github.com/1aal/kubeblocks@v0.0.0-20231107070852-e1c03e598921/pkg/cli/util/completion.go (about)

     1  /*
     2  Copyright (C) 2022-2023 ApeCloud Co., Ltd
     3  
     4  This file is part of KubeBlocks project
     5  
     6  This program is free software: you can redistribute it and/or modify
     7  it under the terms of the GNU Affero General Public License as published by
     8  the Free Software Foundation, either version 3 of the License, or
     9  (at your option) any later version.
    10  
    11  This program is distributed in the hope that it will be useful
    12  but WITHOUT ANY WARRANTY; without even the implied warranty of
    13  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    14  GNU Affero General Public License for more details.
    15  
    16  You should have received a copy of the GNU Affero General Public License
    17  along with this program.  If not, see <http://www.gnu.org/licenses/>.
    18  */
    19  
    20  package util
    21  
    22  import (
    23  	"bytes"
    24  	"io/ioutil"
    25  	"os"
    26  	"strings"
    27  
    28  	"github.com/spf13/cobra"
    29  	"k8s.io/apimachinery/pkg/api/meta"
    30  	"k8s.io/apimachinery/pkg/runtime/schema"
    31  	"k8s.io/cli-runtime/pkg/genericiooptions"
    32  	"k8s.io/cli-runtime/pkg/printers"
    33  	"k8s.io/kubectl/pkg/cmd/get"
    34  	cmdutil "k8s.io/kubectl/pkg/cmd/util"
    35  	utilcomp "k8s.io/kubectl/pkg/util/completion"
    36  
    37  	"github.com/1aal/kubeblocks/pkg/cli/types"
    38  )
    39  
    40  func ResourceNameCompletionFunc(f cmdutil.Factory, gvr schema.GroupVersionResource) func(*cobra.Command, []string, string) ([]string, cobra.ShellCompDirective) {
    41  	return func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
    42  		comps := utilcomp.CompGetResource(f, GVRToString(gvr), toComplete)
    43  		seen := make(map[string]bool)
    44  
    45  		var availableComps []string
    46  		for _, arg := range args {
    47  			seen[arg] = true
    48  		}
    49  		for _, comp := range comps {
    50  			if !seen[comp] {
    51  				availableComps = append(availableComps, comp)
    52  			}
    53  		}
    54  		return availableComps, cobra.ShellCompDirectiveNoFileComp
    55  	}
    56  }
    57  
    58  // CompGetResourceWithLabels gets the list of the resource specified which begin with `toComplete` and have the specified labels.
    59  // example: CompGetResourceWithLabels(f, cmd, "pods", []string{"app=nginx"}, toComplete)
    60  // gets the name of the pods which have the label `app=nginx` and begin with `toComplete`
    61  func CompGetResourceWithLabels(f cmdutil.Factory, cmd *cobra.Command, resourceName string, labels []string, toComplete string) []string {
    62  	template := "{{ range .items  }}{{ .metadata.name }} {{ end }}"
    63  	return CompGetFromTemplateWithLabels(&template, f, "", cmd, []string{resourceName}, labels, toComplete)
    64  }
    65  
    66  // CompGetFromTemplateWithLabels executes a Get operation using the specified template and args and returns the results
    67  // which begin with `toComplete` and have the specified labels.
    68  // example: CompGetFromTemplateWithLabels(&template, f, "", cmd, []string{"pods"}, []string{"app=nginx"}, toComplete)
    69  // will get the output of `kubectl get pods --template=template -l app=nginx`, and split the output by space and return
    70  func CompGetFromTemplateWithLabels(template *string, f cmdutil.Factory, namespace string, cmd *cobra.Command, args []string, labels []string, toComplete string) []string {
    71  	buf := new(bytes.Buffer)
    72  	streams := genericiooptions.IOStreams{In: os.Stdin, Out: buf, ErrOut: ioutil.Discard}
    73  	o := get.NewGetOptions("kubectl", streams)
    74  
    75  	// Get the list of names of the specified resource
    76  	o.PrintFlags.TemplateFlags.GoTemplatePrintFlags.TemplateArgument = template
    77  	format := "go-template"
    78  	o.PrintFlags.OutputFormat = &format
    79  
    80  	// Do the steps Complete() would have done.
    81  	// We cannot actually call Complete() or Validate() as these function check for
    82  	// the presence of flags, which, in our case won't be there
    83  	if namespace != "" {
    84  		o.Namespace = namespace
    85  		o.ExplicitNamespace = true
    86  	} else {
    87  		var err error
    88  		o.Namespace, o.ExplicitNamespace, err = f.ToRawKubeConfigLoader().Namespace()
    89  		if err != nil {
    90  			return nil
    91  		}
    92  	}
    93  
    94  	o.ToPrinter = func(mapping *meta.RESTMapping, outputObjects *bool, withNamespace bool, withKind bool) (printers.ResourcePrinterFunc, error) {
    95  		printer, err := o.PrintFlags.ToPrinter()
    96  		if err != nil {
    97  			return nil, err
    98  		}
    99  		return printer.PrintObj, nil
   100  	}
   101  
   102  	if len(labels) > 0 {
   103  		o.LabelSelector = strings.Join(labels, ",")
   104  	}
   105  
   106  	_ = o.Run(f, args)
   107  
   108  	var comps []string
   109  	resources := strings.Split(buf.String(), " ")
   110  	for _, res := range resources {
   111  		if res != "" && strings.HasPrefix(res, toComplete) {
   112  			comps = append(comps, res)
   113  		}
   114  	}
   115  	return comps
   116  }
   117  
   118  func RegisterClusterCompletionFunc(cmd *cobra.Command, f cmdutil.Factory) {
   119  	if cmd.Flags().Lookup("cluster") == nil {
   120  		return
   121  	}
   122  
   123  	cmdutil.CheckErr(cmd.RegisterFlagCompletionFunc(
   124  		"cluster",
   125  		func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
   126  			return utilcomp.CompGetResource(f, GVRToString(types.ClusterGVR()), toComplete), cobra.ShellCompDirectiveNoFileComp
   127  		},
   128  	))
   129  }