From 2174d77dc44384193ba612f5588c681a585131a2 Mon Sep 17 00:00:00 2001
From: Manabu McCloskey <manabu.mccloskey@gmail.com>
Date: Thu, 21 Nov 2024 20:57:21 +0000
Subject: [PATCH] pass command context

Signed-off-by: Manabu McCloskey <manabu.mccloskey@gmail.com>
---
 main.go                | 26 ++++++++++++++++++++++++--
 pkg/build/build.go     | 20 ++++++++------------
 pkg/cmd/create/root.go | 12 ++++++++----
 pkg/cmd/get/secrets.go |  3 +--
 pkg/cmd/root.go        |  5 +++--
 5 files changed, 44 insertions(+), 22 deletions(-)

diff --git a/main.go b/main.go
index 4f5ffcf9..71391e3f 100644
--- a/main.go
+++ b/main.go
@@ -1,7 +1,29 @@
 package main
 
-import "github.com/cnoe-io/idpbuilder/pkg/cmd"
+import (
+	"context"
+	"fmt"
+	"os"
+	"os/signal"
+	"syscall"
+
+	"github.com/cnoe-io/idpbuilder/pkg/cmd"
+)
 
 func main() {
-	cmd.Execute()
+	interrupted := make(chan os.Signal, 1)
+	defer close(interrupted)
+	signal.Notify(interrupted, os.Interrupt, syscall.SIGTERM)
+	defer signal.Stop(interrupted)
+
+	ctx, cancel := context.WithCancelCause(context.Background())
+
+	go func() {
+		select {
+		case <-interrupted:
+			cancel(fmt.Errorf("command interrupted"))
+		}
+	}()
+
+	cmd.Execute(ctx)
 }
diff --git a/pkg/build/build.go b/pkg/build/build.go
index f325123c..50d1aaf7 100644
--- a/pkg/build/build.go
+++ b/pkg/build/build.go
@@ -4,8 +4,6 @@ import (
 	"context"
 	"fmt"
 	"os"
-	"os/signal"
-	"syscall"
 	"time"
 
 	"github.com/cnoe-io/idpbuilder/api/v1alpha1"
@@ -159,8 +157,6 @@ func (b *Build) isCompatible(ctx context.Context, kubeClient client.Client) (boo
 }
 
 func (b *Build) Run(ctx context.Context, recreateCluster bool) error {
-	managerExit := make(chan error)
-
 	setupLog.Info("Creating kind cluster")
 	if err := b.ReconcileKindCluster(ctx, recreateCluster); err != nil {
 		return err
@@ -227,6 +223,8 @@ func (b *Build) Run(ctx context.Context, recreateCluster bool) error {
 		return err
 	}
 
+	managerExit := make(chan error)
+
 	setupLog.V(1).Info("Running controllers")
 	if err := b.RunControllers(ctx, mgr, managerExit, dir); err != nil {
 		setupLog.Error(err, "Error running controllers")
@@ -268,17 +266,15 @@ func (b *Build) Run(ctx context.Context, recreateCluster bool) error {
 		return fmt.Errorf("creating localbuild resource: %w", err)
 	}
 
-	interrupted := make(chan os.Signal, 1)
-	defer close(interrupted)
-	signal.Notify(interrupted, os.Interrupt, syscall.SIGTERM)
-
 	select {
 	case mgrErr := <-managerExit:
-		return mgrErr
-	case <-interrupted:
-		b.CancelFunc()
-		return fmt.Errorf("command interrupted")
+		if mgrErr != nil {
+			return mgrErr
+		}
+	case <-ctx.Done():
+		return nil
 	}
+	return nil
 }
 
 func isBuildCustomizationSpecEqual(s1, s2 v1alpha1.BuildCustomizationSpec) bool {
diff --git a/pkg/cmd/create/root.go b/pkg/cmd/create/root.go
index 1e2ff4b9..6d22d7ba 100644
--- a/pkg/cmd/create/root.go
+++ b/pkg/cmd/create/root.go
@@ -15,7 +15,6 @@ import (
 	"github.com/cnoe-io/idpbuilder/pkg/k8s"
 	"github.com/spf13/cobra"
 	"k8s.io/client-go/util/homedir"
-	ctrl "sigs.k8s.io/controller-runtime"
 )
 
 const (
@@ -89,8 +88,9 @@ func preCreateE(cmd *cobra.Command, args []string) error {
 }
 
 func create(cmd *cobra.Command, args []string) error {
-	ctx, ctxCancel := context.WithCancel(ctrl.SetupSignalHandler())
-	defer ctxCancel()
+
+	ctx, cancel := context.WithCancel(cmd.Context())
+	defer cancel()
 
 	kubeConfigPath := filepath.Join(homedir.HomeDir(), ".kube", "config")
 
@@ -152,7 +152,7 @@ func create(cmd *cobra.Command, args []string) error {
 		PackageCustomization: o,
 
 		Scheme:     k8s.GetScheme(),
-		CancelFunc: ctxCancel,
+		CancelFunc: cancel,
 	}
 
 	b := build.NewBuild(opts)
@@ -161,6 +161,10 @@ func create(cmd *cobra.Command, args []string) error {
 		return err
 	}
 
+	if cmd.Context().Err() != nil {
+		return context.Cause(cmd.Context())
+	}
+
 	printSuccessMsg()
 	return nil
 }
diff --git a/pkg/cmd/get/secrets.go b/pkg/cmd/get/secrets.go
index 5cad39c3..11b3f52d 100644
--- a/pkg/cmd/get/secrets.go
+++ b/pkg/cmd/get/secrets.go
@@ -19,7 +19,6 @@ import (
 	"k8s.io/apimachinery/pkg/labels"
 	"k8s.io/apimachinery/pkg/selection"
 	"k8s.io/client-go/util/homedir"
-	ctrl "sigs.k8s.io/controller-runtime"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 	"sigs.k8s.io/yaml"
 )
@@ -54,7 +53,7 @@ type TemplateData struct {
 }
 
 func getSecretsE(cmd *cobra.Command, args []string) error {
-	ctx, ctxCancel := context.WithCancel(ctrl.SetupSignalHandler())
+	ctx, ctxCancel := context.WithCancel(cmd.Context())
 	defer ctxCancel()
 	kubeConfigPath := filepath.Join(homedir.HomeDir(), ".kube", "config")
 
diff --git a/pkg/cmd/root.go b/pkg/cmd/root.go
index 32d97ab4..345aca66 100644
--- a/pkg/cmd/root.go
+++ b/pkg/cmd/root.go
@@ -1,6 +1,7 @@
 package cmd
 
 import (
+	"context"
 	"fmt"
 	"os"
 
@@ -27,8 +28,8 @@ func init() {
 	rootCmd.AddCommand(version.VersionCmd)
 }
 
-func Execute() {
-	if err := rootCmd.Execute(); err != nil {
+func Execute(ctx context.Context) {
+	if err := rootCmd.ExecuteContext(ctx); err != nil {
 		fmt.Fprintln(os.Stderr, err)
 		os.Exit(1)
 	}