feat: add talosctl completions to copy, usage, logs, restart and service

talosctl copy: path completion of src node path, dst client path was done by default talosctl usage: directory completion of node path talosctl logs: service completion talosctl restart: running container/pod completion talosctl service: service and action completion Signed-off-by: Nico Berlee <[email protected]> Signed-off-by: Andrey Smirnov <[email protected]>
siderolabs · Mar 16, 2022 · 6bec084 · 6bec084
1 parent 355b1a4
commit 6bec084
Show file tree

Hide file tree

Showing 6 changed files with 146 additions and 4 deletions.
diff --git a/cmd/talosctl/cmd/talos/copy.go b/cmd/talosctl/cmd/talos/copy.go
@@ -32,6 +32,16 @@ talosctl creates a directory if <local-path> doesn't exist. Command doesn't pres
 ownership and access mode for the files in extract mode, while  streamed .tar archive
 captures ownership and permission bits.`,
 	Args: cobra.ExactArgs(2),
+	ValidArgsFunction: func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
+		switch len(args) {
+		case 0:
+			return completePathFromNode(toComplete), cobra.ShellCompDirectiveNoFileComp
+		case 1:
+			return nil, cobra.ShellCompDirectiveDefault
+		}
+
+		return nil, cobra.ShellCompDirectiveError | cobra.ShellCompDirectiveNoFileComp
+	},
 	RunE: func(cmd *cobra.Command, args []string) error {
 		return WithClient(func(ctx context.Context, c *client.Client) error {
 			if err := helpers.FailIfMultiNodes(ctx, "copy"); err != nil {

diff --git a/cmd/talosctl/cmd/talos/diskusage.go b/cmd/talosctl/cmd/talos/diskusage.go
@@ -30,6 +30,21 @@ var duCmd = &cobra.Command{
 	Aliases: []string{"du"},
 	Short:   "Retrieve a disk usage",
 	Long:    ``,
+	ValidArgsFunction: func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
+		if len(args) != 0 {
+			return nil, cobra.ShellCompDirectiveError | cobra.ShellCompDirectiveNoFileComp
+		}
+
+		var completeOnlyPaths []string
+
+		for _, path := range completePathFromNode(toComplete) {
+			if path[len(path)-1:] == "/" {
+				completeOnlyPaths = append(completeOnlyPaths, path)
+			}
+		}
+
+		return completeOnlyPaths, cobra.ShellCompDirectiveNoFileComp
+	},
 	RunE: func(cmd *cobra.Command, args []string) error {
 		return WithClient(func(ctx context.Context, c *client.Client) error {
 			var paths []string

diff --git a/cmd/talosctl/cmd/talos/logs.go b/cmd/talosctl/cmd/talos/logs.go
@@ -34,6 +34,17 @@ var logsCmd = &cobra.Command{
 	Short: "Retrieve logs for a service",
 	Long:  ``,
 	Args:  cobra.ExactArgs(1),
+	ValidArgsFunction: func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
+		if len(args) != 0 {
+			return nil, cobra.ShellCompDirectiveError | cobra.ShellCompDirectiveNoFileComp
+		}
+
+		if kubernetes {
+			return getContainersFromNode(kubernetes), cobra.ShellCompDirectiveNoFileComp
+		}
+
+		return mergeSuggestions(getServiceFromNode(), getContainersFromNode(kubernetes)), cobra.ShellCompDirectiveNoFileComp
+	},
 	RunE: func(cmd *cobra.Command, args []string) error {
 		return WithClient(func(ctx context.Context, c *client.Client) error {
 			var (

diff --git a/cmd/talosctl/cmd/talos/restart.go b/cmd/talosctl/cmd/talos/restart.go
@@ -22,6 +22,13 @@ var restartCmd = &cobra.Command{
 	Short: "Restart a process",
 	Long:  ``,
 	Args:  cobra.ExactArgs(1),
+	ValidArgsFunction: func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
+		if len(args) != 0 {
+			return nil, cobra.ShellCompDirectiveError | cobra.ShellCompDirectiveNoFileComp
+		}
+
+		return getContainersFromNode(kubernetes), cobra.ShellCompDirectiveNoFileComp
+	},
 	RunE: func(cmd *cobra.Command, args []string) error {
 		return WithClient(func(ctx context.Context, c *client.Client) error {
 			var (

diff --git a/cmd/talosctl/cmd/talos/root.go b/cmd/talosctl/cmd/talos/root.go
@@ -9,17 +9,23 @@ import (
 	"crypto/tls"
 	"fmt"
 	"io"
+	"sort"
 	"strings"
 
+	criconstants "github.com/containerd/cri/pkg/constants"
 	"github.com/spf13/cobra"
 	"github.com/talos-systems/crypto/x509"
+	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/peer"
 
 	"github.com/talos-systems/talos/pkg/cli"
 	_ "github.com/talos-systems/talos/pkg/grpc/codec" //nolint:gci // register codec
+	"github.com/talos-systems/talos/pkg/machinery/api/common"
 	machineapi "github.com/talos-systems/talos/pkg/machinery/api/machine"
 	"github.com/talos-systems/talos/pkg/machinery/client"
 	clientconfig "github.com/talos-systems/talos/pkg/machinery/client/config"
+	"github.com/talos-systems/talos/pkg/machinery/constants"
 )
 
 var kubernetes bool
@@ -164,7 +170,7 @@ func completePathFromNode(inputPath string) []string {
 func getPathFromNode(path, filter string) map[string]struct{} {
 	paths := make(map[string]struct{})
 
-	if WithClient(func(ctx context.Context, c *client.Client) error {
+	WithClient(func(ctx context.Context, c *client.Client) error { //nolint:errcheck
 		ctx, cancel := context.WithCancel(ctx)
 		defer cancel()
 
@@ -218,13 +224,96 @@ func getPathFromNode(path, filter string) map[string]struct{} {
 				}
 			}
 		}
-	}) != nil {
-		return paths
-	}
+	})
 
 	return paths
 }
 
+func getServiceFromNode() []string {
+	var svcIds []string
+
+	WithClient(func(ctx context.Context, c *client.Client) error { //nolint:errcheck
+		var remotePeer peer.Peer
+
+		resp, err := c.ServiceList(ctx, grpc.Peer(&remotePeer))
+		if err != nil {
+			return err
+		}
+
+		for _, msg := range resp.Messages {
+			for _, s := range msg.Services {
+				svc := cli.ServiceInfoWrapper{ServiceInfo: s}
+				svcIds = append(svcIds, svc.Id)
+			}
+		}
+
+		return nil
+	})
+
+	return svcIds
+}
+
+func getContainersFromNode(kubernetes bool) []string {
+	var containerIds []string
+
+	WithClient(func(ctx context.Context, c *client.Client) error { //nolint:errcheck
+		var (
+			namespace string
+			driver    common.ContainerDriver
+		)
+
+		if kubernetes {
+			namespace = criconstants.K8sContainerdNamespace
+			driver = common.ContainerDriver_CRI
+		} else {
+			namespace = constants.SystemContainerdNamespace
+			driver = common.ContainerDriver_CONTAINERD
+		}
+
+		resp, err := c.Containers(ctx, namespace, driver)
+		if err != nil {
+			return err
+		}
+
+		for _, msg := range resp.Messages {
+			for _, p := range msg.Containers {
+				if p.Pid == 0 {
+					continue
+				}
+
+				if kubernetes && p.Id == p.PodId {
+					continue
+				}
+
+				containerIds = append(containerIds, p.Id)
+			}
+		}
+
+		return nil
+	})
+
+	return containerIds
+}
+
+func mergeSuggestions(a, b []string) []string {
+	merged := append(append([]string(nil), a...), b...)
+
+	sort.Strings(merged)
+
+	n := 1
+
+	for i := 1; i < len(merged); i++ {
+		if merged[i] != merged[i-1] {
+			merged[n] = merged[i]
+			n++
+		}
+	}
+
+	merged = merged[:n]
+
+	return merged
+}
+
 func relativeTo(fullPath string, filter string) bool {
 	return strings.HasPrefix(fullPath, filter)
 }
diff --git a/cmd/talosctl/cmd/talos/service.go b/cmd/talosctl/cmd/talos/service.go
@@ -27,6 +27,16 @@ var serviceCmd = &cobra.Command{
 If service ID is specified, default action 'status' is executed which shows status of a single list service.
 With actions 'start', 'stop', 'restart', service state is updated respectively.`,
 	Args: cobra.MaximumNArgs(2),
+	ValidArgsFunction: func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
+		switch len(args) {
+		case 0:
+			return getServiceFromNode(), cobra.ShellCompDirectiveNoFileComp
+		case 1:
+			return []string{"start", "stop", "restart", "status"}, cobra.ShellCompDirectiveNoFileComp
+		}
+
+		return nil, cobra.ShellCompDirectiveError | cobra.ShellCompDirectiveNoFileComp
+	},
 	RunE: func(cmd *cobra.Command, args []string) error {
 		action := "status"
 		serviceID := ""