From df81cf7a74e6f78c27055b82c20375f53976cea8 Mon Sep 17 00:00:00 2001
From: Kristin Laemmert <mildwonkey@users.noreply.github.com>
Date: Thu, 3 Oct 2024 12:58:27 -0400
Subject: [PATCH] chore: add global command context for program cancelation and
 everything else (#696)

This implements cancelation with CTRL-C and adds the context with cancelation to the cobra root command. The context is not used anywhere else yet, but sub-commands can access the ctx in their Run methods w/ cmd.Context().
---
 go.sum          |  6 ------
 src/cmd/root.go | 23 +++++++++++++++++++++--
 2 files changed, 21 insertions(+), 8 deletions(-)

diff --git a/go.sum b/go.sum
index c08549bd..3f2ee55e 100644
--- a/go.sum
+++ b/go.sum
@@ -58,16 +58,10 @@ github.com/charmbracelet/bubbletea v1.1.1 h1:KJ2/DnmpfqFtDNVTvYZ6zpPFL9iRCRr0qqK
 github.com/charmbracelet/bubbletea v1.1.1/go.mod h1:9Ogk0HrdbHolIKHdjfFpyXJmiCzGwy+FesYkZr7hYU4=
 github.com/charmbracelet/lipgloss v0.13.0 h1:4X3PPeoWEDCMvzDvGmTajSyYPcZM4+y8sCA/SsA3cjw=
 github.com/charmbracelet/lipgloss v0.13.0/go.mod h1:nw4zy0SBX/F/eAO1cWdcvy6qnkDUxr8Lw7dvFrAIbbY=
-github.com/charmbracelet/x/ansi v0.2.3 h1:VfFN0NUpcjBRd4DnKfRaIRo53KRgey/nhOoEqosGDEY=
-github.com/charmbracelet/x/ansi v0.2.3/go.mod h1:dk73KoMTT5AX5BsX0KrqhsTqAnhZZoCBjs7dGWp4Ktw=
 github.com/charmbracelet/x/ansi v0.3.2 h1:wsEwgAN+C9U06l9dCVMX0/L3x7ptvY1qmjMwyfE6USY=
 github.com/charmbracelet/x/ansi v0.3.2/go.mod h1:dk73KoMTT5AX5BsX0KrqhsTqAnhZZoCBjs7dGWp4Ktw=
-github.com/charmbracelet/x/exp/golden v0.0.0-20240815200342-61de596daa2b h1:MnAMdlwSltxJyULnrYbkZpp4k58Co7Tah3ciKhSNo0Q=
-github.com/charmbracelet/x/exp/golden v0.0.0-20240815200342-61de596daa2b/go.mod h1:wDlXFlCrmJ8J+swcL/MnGUuYnqgQdW9rhSD61oNMb6U=
 github.com/charmbracelet/x/exp/golden v0.0.0-20240919170804-a4978c8e603a h1:IUy+N6nKpGfijckOe8KGnAQwBUT6xz63n3tbb0Gy8aY=
 github.com/charmbracelet/x/exp/golden v0.0.0-20240919170804-a4978c8e603a/go.mod h1:wDlXFlCrmJ8J+swcL/MnGUuYnqgQdW9rhSD61oNMb6U=
-github.com/charmbracelet/x/exp/teatest v0.0.0-20240918160051-227168dc0568 h1:MWPyZsxZMe/oKhkt5j34zAba/2nXOxWuf3CvQqO8SDA=
-github.com/charmbracelet/x/exp/teatest v0.0.0-20240918160051-227168dc0568/go.mod h1:NDRRSMP6bZbCs4jyc4i1/4UG4M+0PEiQdpivQgD0Mio=
 github.com/charmbracelet/x/exp/teatest v0.0.0-20240919170804-a4978c8e603a h1:sS42HbmCab8rCehUwNO/bQEZQoJ6GavhZyO+245mBwA=
 github.com/charmbracelet/x/exp/teatest v0.0.0-20240919170804-a4978c8e603a/go.mod h1:NDRRSMP6bZbCs4jyc4i1/4UG4M+0PEiQdpivQgD0Mio=
 github.com/charmbracelet/x/term v0.2.0 h1:cNB9Ot9q8I711MyZ7myUR5HFWL/lc3OpU8jZ4hwm0x0=
diff --git a/src/cmd/root.go b/src/cmd/root.go
index 2f13ebd9..e73d1dd5 100644
--- a/src/cmd/root.go
+++ b/src/cmd/root.go
@@ -1,7 +1,11 @@
 package cmd
 
 import (
-	"github.com/spf13/cobra"
+	"context"
+	"fmt"
+	"os"
+	"os/signal"
+	"syscall"
 
 	"github.com/defenseunicorns/lula/src/cmd/common"
 	"github.com/defenseunicorns/lula/src/cmd/console"
@@ -11,6 +15,7 @@ import (
 	"github.com/defenseunicorns/lula/src/cmd/tools"
 	"github.com/defenseunicorns/lula/src/cmd/validate"
 	"github.com/defenseunicorns/lula/src/cmd/version"
+	"github.com/spf13/cobra"
 )
 
 var LogLevelCLI string
@@ -32,8 +37,22 @@ func RootCommand() *cobra.Command {
 }
 
 func Execute() {
+	ctx, cancel := context.WithCancel(context.Background())
+	c := make(chan os.Signal, 1)
+	signal.Notify(c, os.Interrupt, syscall.SIGHUP, syscall.SIGTERM)
+
+	go func() {
+		select {
+		case <-c:
+			fmt.Println("Got signal, shutting down...")
+			cancel()
+			os.Exit(2)
+		case <-ctx.Done():
+			return
+		}
+	}()
 
-	cobra.CheckErr(rootCmd.Execute())
+	cobra.CheckErr(rootCmd.ExecuteContext(ctx))
 }
 
 func init() {