refactor: Using cobra contexte

Signed-off-by: Vincent Boutour <[email protected]>
ViBiOh · Jan 8, 2023 · 46608f3 · 46608f3
1 parent b9efc07
commit 46608f3
Show file tree

Hide file tree

Showing 8 changed files with 27 additions and 24 deletions.
diff --git a/cmd/choose.go b/cmd/choose.go
@@ -21,16 +21,16 @@ func getCommonNamespace(kube client.Kube, namespace string) string {
 	return ""
 }
 
-func getCommonObjects(namespace string, lister resource.Lister) []string {
+func getCommonObjects(ctx context.Context, namespace string, lister resource.Lister) []string {
 	output := make(chan string, len(clients))
 	successChan := make(chan bool, len(clients))
 
 	go func() {
 		defer close(output)
 		defer close(successChan)
 
-		clients.Execute(func(kube client.Kube) error {
-			items, err := lister(context.Background(), kube, getCommonNamespace(kube, namespace), metav1.ListOptions{})
+		clients.Execute(ctx, func(ctx context.Context, kube client.Kube) error {
+			items, err := lister(ctx, kube, getCommonNamespace(kube, namespace), metav1.ListOptions{})
 			if err != nil {
 				return err
 			}

diff --git a/cmd/image.go b/cmd/image.go
@@ -36,7 +36,7 @@ var imageCmd = &cobra.Command{
 				return nil, cobra.ShellCompDirectiveError
 			}
 
-			return getCommonObjects(viper.GetString("namespace"), lister), cobra.ShellCompDirectiveNoFileComp
+			return getCommonObjects(cmd.Context(), viper.GetString("namespace"), lister), cobra.ShellCompDirectiveNoFileComp
 		}
 
 		return nil, cobra.ShellCompDirectiveNoFileComp
@@ -46,15 +46,15 @@ var imageCmd = &cobra.Command{
 		resourceType := args[0]
 		resourceName := args[1]
 
-		ctx, cancel := context.WithCancel(context.Background())
+		ctx, cancel := context.WithCancel(cmd.Context())
 		defer cancel()
 
 		go func() {
 			waitForEnd(syscall.SIGINT, syscall.SIGTERM)
 			cancel()
 		}()
 
-		clients.Execute(func(kube client.Kube) error {
+		clients.Execute(ctx, func(ctx context.Context, kube client.Kube) error {
 			podTemplate, err := resource.PodTemplateGetter(ctx, kube, resourceType, resourceName)
 			if err != nil {
 				return err

diff --git a/cmd/log.go b/cmd/log.go
@@ -68,7 +68,7 @@ var logCmd = &cobra.Command{
 				return nil, cobra.ShellCompDirectiveError
 			}
 
-			return getCommonObjects(viper.GetString("namespace"), lister), cobra.ShellCompDirectiveNoFileComp
+			return getCommonObjects(cmd.Context(), viper.GetString("namespace"), lister), cobra.ShellCompDirectiveNoFileComp
 		}
 
 		return nil, cobra.ShellCompDirectiveNoFileComp
@@ -78,7 +78,7 @@ var logCmd = &cobra.Command{
 			return errors.New("either labels or `TYPE NAME` args must be specified")
 		}
 
-		ctx, cancel := context.WithCancel(context.Background())
+		ctx, cancel := context.WithCancel(cmd.Context())
 		defer cancel()
 
 		go func() {
@@ -120,7 +120,7 @@ var logCmd = &cobra.Command{
 			resourceName = args[1]
 		}
 
-		clients.Execute(func(kube client.Kube) error {
+		clients.Execute(ctx, func(ctx context.Context, kube client.Kube) error {
 			podWatcher, err := resource.WatchPods(ctx, kube, resourceType, resourceName, labelsSelector, dryRun)
 			if err != nil {
 				return fmt.Errorf("watch pods: %w", err)

diff --git a/cmd/port_forward.go b/cmd/port_forward.go
@@ -51,7 +51,7 @@ var portForwardCmd = &cobra.Command{
 				return nil, cobra.ShellCompDirectiveError
 			}
 
-			return getCommonObjects(viper.GetString("namespace"), lister), cobra.ShellCompDirectiveNoFileComp
+			return getCommonObjects(cmd.Context(), viper.GetString("namespace"), lister), cobra.ShellCompDirectiveNoFileComp
 		}
 
 		return nil, cobra.ShellCompDirectiveNoFileComp
@@ -76,7 +76,7 @@ var portForwardCmd = &cobra.Command{
 			remotePort = ports[0]
 		}
 
-		ctx, cancel := context.WithCancel(context.Background())
+		ctx, cancel := context.WithCancel(cmd.Context())
 		defer cancel()
 
 		output.Std("", "Listening tcp on %d", localPort)
@@ -92,7 +92,7 @@ var portForwardCmd = &cobra.Command{
 			cancel()
 		}()
 
-		clients.Execute(func(kube client.Kube) error {
+		clients.Execute(ctx, func(ctx context.Context, kube client.Kube) error {
 			remotePort := remotePort
 
 			if resource.IsService(resourceType) {

diff --git a/cmd/restart.go b/cmd/restart.go
@@ -50,7 +50,7 @@ var restartCmd = &cobra.Command{
 				return nil, cobra.ShellCompDirectiveError
 			}
 
-			return getCommonObjects(viper.GetString("namespace"), lister), cobra.ShellCompDirectiveNoFileComp
+			return getCommonObjects(cmd.Context(), viper.GetString("namespace"), lister), cobra.ShellCompDirectiveNoFileComp
 		}
 
 		return nil, cobra.ShellCompDirectiveNoFileComp
@@ -60,7 +60,7 @@ var restartCmd = &cobra.Command{
 		resourceType := args[0]
 		resourceName := args[1]
 
-		ctx, cancel := context.WithCancel(context.Background())
+		ctx, cancel := context.WithCancel(cmd.Context())
 		defer cancel()
 
 		var patch restartPatch
@@ -74,7 +74,7 @@ var restartCmd = &cobra.Command{
 			return
 		}
 
-		clients.Execute(func(kube client.Kube) error {
+		clients.Execute(ctx, func(ctx context.Context, kube client.Kube) error {
 			switch resourceType {
 			case "ds", "daemonset", "daemonsets":
 				_, err := kube.AppsV1().DaemonSets(kube.Namespace).Patch(ctx, resourceName, types.MergePatchType, payload, v1.PatchOptions{})

diff --git a/cmd/root.go b/cmd/root.go
@@ -1,6 +1,7 @@
 package cmd
 
 import (
+	"context"
 	"fmt"
 	"path/filepath"
 	"strings"
@@ -37,7 +38,7 @@ var rootCmd = &cobra.Command{
 		<-output.Done()
 	},
 	Run: func(cmd *cobra.Command, args []string) {
-		clients.Execute(func(kube client.Kube) error {
+		clients.Execute(cmd.Context(), func(ctx context.Context, kube client.Kube) error {
 			info, err := kube.Discovery().ServerVersion()
 			if err != nil {
 				return fmt.Errorf("get server version: %w", err)
@@ -172,7 +173,7 @@ func completeContext(_ *cobra.Command, _ []string, _ string) ([]string, cobra.Sh
 	return completeContexts, cobra.ShellCompDirectiveNoFileComp
 }
 
-func completeNamespace(_ *cobra.Command, _ []string, _ string) ([]string, cobra.ShellCompDirective) {
+func completeNamespace(cmd *cobra.Command, _ []string, _ string) ([]string, cobra.ShellCompDirective) {
 	lister, err := resource.ListerFor("namespace")
 	if err != nil {
 		return nil, cobra.ShellCompDirectiveError
@@ -183,7 +184,7 @@ func completeNamespace(_ *cobra.Command, _ []string, _ string) ([]string, cobra.
 		return nil, cobra.ShellCompDirectiveError
 	}
 
-	return getCommonObjects("", lister), cobra.ShellCompDirectiveDefault
+	return getCommonObjects(cmd.Context(), "", lister), cobra.ShellCompDirectiveDefault
 }
 
 func contains(arr []string, value string) bool {

diff --git a/cmd/watch.go b/cmd/watch.go
@@ -39,7 +39,7 @@ var watchCmd = &cobra.Command{
 	Use:   "watch",
 	Short: "Get all pods in the namespace",
 	Run: func(cmd *cobra.Command, args []string) {
-		ctx, cancel := context.WithCancel(context.Background())
+		ctx, cancel := context.WithCancel(cmd.Context())
 		defer cancel()
 
 		go func() {
@@ -50,7 +50,7 @@ var watchCmd = &cobra.Command{
 		watchTable := initWatchTable()
 		initialsPodsHash := displayInitialPods(ctx, watchTable)
 
-		clients.Execute(func(kube client.Kube) error {
+		clients.Execute(ctx, func(ctx context.Context, kube client.Kube) error {
 			watcher, err := resource.WatchPods(ctx, kube, "namespace", kube.Namespace, labelsSelector, false)
 			if err != nil {
 				return fmt.Errorf("watch pods: %w", err)
@@ -128,7 +128,7 @@ func displayInitialPods(ctx context.Context, watchTable *table.Table) map[string
 		}
 	}()
 
-	clients.Execute(func(kube client.Kube) error {
+	clients.Execute(ctx, func(ctx context.Context, kube client.Kube) error {
 		watcher, err := resource.WatchPods(ctx, kube, "namespace", kube.Namespace, labelsSelector, true)
 		if err != nil {
 			return fmt.Errorf("watch pods: %w", err)

diff --git a/pkg/client/client.go b/pkg/client/client.go
@@ -1,6 +1,8 @@
 package client
 
 import (
+	"context"
+
 	"github.com/ViBiOh/kmux/pkg/concurrent"
 	"github.com/ViBiOh/kmux/pkg/output"
 	"k8s.io/client-go/kubernetes"
@@ -25,18 +27,18 @@ func New(name, namespace string, config *rest.Config, clientset *kubernetes.Clie
 	}
 }
 
-type Action func(Kube) error
+type Action func(context.Context, Kube) error
 
 type Array []Kube
 
-func (a Array) Execute(action Action) {
+func (a Array) Execute(ctx context.Context, action Action) {
 	parallel := concurrent.NewSimple()
 
 	for _, client := range a {
 		client := client
 
 		parallel.Go(func() {
-			if err := action(client); err != nil {
+			if err := action(ctx, client); err != nil {
 				client.Err("%s", err)
 			}
 		})