From 5cd055510f1b025ff4d80ad14ad8996a08b84189 Mon Sep 17 00:00:00 2001
From: Herman Schaaf <hermanschaaf@gmail.com>
Date: Wed, 18 Jan 2023 14:00:10 +0000
Subject: [PATCH] Log received signal when shutting down

---
 cli/main.go | 23 ++++++++++++++++++-----
 1 file changed, 18 insertions(+), 5 deletions(-)

diff --git a/cli/main.go b/cli/main.go
index 867ad49a5ac5ee..cbd55bbf8e6fd9 100644
--- a/cli/main.go
+++ b/cli/main.go
@@ -6,6 +6,8 @@ import (
 	"os"
 	"os/signal"
 	"runtime/debug"
+	"sync"
+	"syscall"
 
 	"github.com/cloudquery/cloudquery/cli/cmd"
 	"github.com/getsentry/sentry-go"
@@ -15,23 +17,34 @@ import (
 
 func executeRootCmdWithContext() error {
 	ctx := context.Background()
-	// trap Ctrl+C and call cancel on the context
+	// trap Ctrl+C and other signals, then call cancel on the context
 	ctx, cancel := context.WithCancel(ctx)
 	c := make(chan os.Signal, 1)
-	signal.Notify(c, os.Interrupt)
+	signal.Notify(c, os.Interrupt, syscall.SIGTERM)
 	defer func() {
 		signal.Stop(c)
-		cancel()
 	}()
+	var gotSignal os.Signal
+	var wg sync.WaitGroup
+	wg.Add(1)
 	go func() {
+		defer wg.Done()
 		select {
-		case <-c:
+		case gotSignal = <-c:
 			cancel()
 		case <-ctx.Done():
 		}
 	}()
 
-	return cmd.NewCmdRoot().ExecuteContext(ctx)
+	err := cmd.NewCmdRoot().ExecuteContext(ctx)
+	cancel()
+	wg.Wait()
+	if gotSignal != nil && err != nil {
+		err = fmt.Errorf("received %v signal from OS: %w", gotSignal.String(), err)
+	} else if gotSignal != nil {
+		err = fmt.Errorf("received %v signal from OS", gotSignal.String())
+	}
+	return err
 }
 
 func main() {