k3k/cli/cmds/completion.go

package cmds

import (
	"context"

	"github.com/sirupsen/logrus"
	"github.com/spf13/cobra"
	"github.com/spf13/pflag"
	"k8s.io/apimachinery/pkg/util/sets"
	"sigs.k8s.io/controller-runtime/pkg/client"

	corev1 "k8s.io/api/core/v1"

	"github.com/rancher/k3k/pkg/apis/k3k.io/v1beta1"
)

var completeClusterMode = cobra.FixedCompletions(
	[]string{
		string(v1beta1.SharedClusterMode),
		string(v1beta1.VirtualClusterMode),
	},
	cobra.ShellCompDirectiveNoFileComp,
)

var completePersistenceMode = cobra.FixedCompletions(
	[]string{
		string(v1beta1.DynamicPersistenceMode),
		string(v1beta1.EphemeralPersistenceMode),
	},
	cobra.ShellCompDirectiveNoFileComp,
)

// completeNamespaces is a cobra.CompletionFunc that completes with every namespace in the host cluster.
func completeNamespaces(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
	client, err := completionClient(cmd)
	if err != nil {
		return nil, cobra.ShellCompDirectiveError
	}

	return namespaceCompletions(cmd, client)
}

// completeClusterNamespaces is a cobra.CompletionFunc that completes with the namespaces with a k3k virtual cluster in it.
func completeClusterNamespaces(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
	client, err := completionClient(cmd)
	if err != nil {
		return nil, cobra.ShellCompDirectiveError
	}

	return clusterNamespaceCompletions(cmd.Context(), client)
}

// namespaceCompletions lists every namespace in the host cluster, excluding any
// already provided to the command's "namespace" flag.
func namespaceCompletions(cmd *cobra.Command, cl client.Client) ([]string, cobra.ShellCompDirective) {
	var namespaces corev1.NamespaceList
	if err := cl.List(cmd.Context(), &namespaces); err != nil {
		return nil, cobra.ShellCompDirectiveError
	}

	allNames := sets.New[string]()
	for _, ns := range namespaces.Items {
		allNames.Insert(ns.Name)
	}

	// Exclude already selected namespaces
	selected := sets.New[string]()
	if vals, err := cmd.Flags().GetStringSlice("namespace"); err == nil {
		selected.Insert(vals...)
	}

	names := allNames.Difference(selected)

	return sets.List(names), cobra.ShellCompDirectiveNoFileComp
}

// clusterNamespaceCompletions lists the unique namespaces that contain at least one k3k Cluster.
func clusterNamespaceCompletions(ctx context.Context, client client.Client) ([]string, cobra.ShellCompDirective) {
	var clusters v1beta1.ClusterList
	if err := client.List(ctx, &clusters); err != nil {
		return nil, cobra.ShellCompDirectiveError
	}

	names := sets.New[string]()
	for _, cluster := range clusters.Items {
		names.Insert(cluster.Namespace)
	}

	return sets.List(names), cobra.ShellCompDirectiveNoFileComp
}

// mustRegisterFlagCompletion registers a completion function for a flag and
// aborts if the flag does not exist. This only fails on programmer error, so
// there is no reason to bubble it up to the caller.
func mustRegisterFlagCompletion(cmd *cobra.Command, flagName string, f cobra.CompletionFunc) {
	if err := cmd.RegisterFlagCompletionFunc(flagName, f); err != nil {
		logrus.Fatal(err)
	}
}

// disableFileCompletion walks the command tree and turns off cobra's default
// filename completion for both positional arguments and flag values, leaving
// anything that is explicitly configured (enum completers, MarkFlagFilename,
// MarkFlagDirname, ValidArgs, bool flags) untouched.
func disableFileCompletion(cmd *cobra.Command) {
	if cmd.ValidArgsFunction == nil && len(cmd.ValidArgs) == 0 {
		cmd.ValidArgsFunction = cobra.NoFileCompletions
	}

	cmd.Flags().VisitAll(func(f *pflag.Flag) {
		if f.Value.Type() == "bool" {
			return
		}

		if _, ok := cmd.GetFlagCompletionFunc(f.Name); ok {
			return
		}

		if _, ok := f.Annotations[cobra.BashCompFilenameExt]; ok {
			return
		}

		if _, ok := f.Annotations[cobra.BashCompSubdirsInDir]; ok {
			return
		}

		_ = cmd.RegisterFlagCompletionFunc(f.Name, cobra.NoFileCompletions)
	})

	for _, sub := range cmd.Commands() {
		disableFileCompletion(sub)
	}
}

// completionClient builds a Kubernetes client for use inside completion functions.
// It checks if the kubeconfig flag is set, to point to the right cluster.
func completionClient(cmd *cobra.Command) (client.Client, error) {
	kubeconfig, err := cmd.Flags().GetString("kubeconfig")
	if err != nil {
		return nil, err
	}

	restConfig, err := loadRESTConfig(kubeconfig, completionRequestTimeout)
	if err != nil {
		return nil, err
	}

	client, err := buildClient(restConfig)
	if err != nil {
		return nil, err
	}

	return client, nil
}