Refactor context handling

Stop using context.Background() and use a root context that's got
signal.NotifyContext() on it. This allows control-C to interrupt
everything.

Use the main context for the MCP handler as well

Author: graveland

cmd/tiger/main.go

@@ -1,7 +1,64 @@
 package main
 
-import "github.com/timescale/tiger-cli/internal/tiger/cmd"
+import (
+	"context"
+	"errors"
+	"fmt"
+	"os"
+	"os/signal"
+	"syscall"
+
+	"go.uber.org/zap"
+
+	"github.com/timescale/tiger-cli/internal/tiger/cmd"
+	"github.com/timescale/tiger-cli/internal/tiger/logging"
+)
 
 func main() {
-	cmd.Execute()
+	if err := run(); err != nil {
+		// Check if it's a custom exit code error
+		if exitErr, ok := err.(interface{ ExitCode() int }); ok {
+			os.Exit(exitErr.ExitCode())
+		}
+		os.Exit(1)
+	}
+	os.Exit(0)
+}
+
+func run() (err error) {
+	ctx, cancel := notifyContext(context.Background())
+	defer func() {
+		cancel()
+		if r := recover(); r != nil {
+			err = errors.Join(err, fmt.Errorf("panic: %v", r))
+			_, _ = fmt.Fprintln(os.Stderr, err.Error())
+		}
+	}()
+	err = cmd.Execute(ctx)
+	return
+}
+
+// noifyContext sets up graceful shutdown handling and returns a context and
+// cleanup function. This is nearly identical to [signal.NotifyContext], except
+// that it logs a message when a signal is received and also restores the default
+// signal handling behavior.
+func notifyContext(parent context.Context) (context.Context, context.CancelFunc) {
+	ctx, cancel := context.WithCancel(parent)
+
+	sigChan := make(chan os.Signal, 1)
+	signal.Notify(sigChan, os.Interrupt, syscall.SIGINT, syscall.SIGTERM)
+	go func() {
+		select {
+		case sig := <-sigChan:
+			logging.Info("Received interrupt signal", zap.Stringer("signal", sig))
+			signal.Stop(sigChan) // Restore default signal handling behavior
+			cancel()
+		case <-ctx.Done():
+		}
+	}()
+
+	return ctx, func() {
+		cancel()
+		signal.Stop(sigChan)
+	}
 }

internal/tiger/cmd/auth_test.go

@@ -2,6 +2,7 @@ package cmd
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -70,24 +71,27 @@ func setupAuthTest(t *testing.T) string {
 	return tmpDir
 }
 
-func executeAuthCommand(args ...string) (string, error) {
+func executeAuthCommand(ctx context.Context, args ...string) (string, error) {
 	// Use buildRootCmd() to get a complete root command with all flags and subcommands
-	testRoot := buildRootCmd()
+	testRoot, err := buildRootCmd(ctx)
+	if err != nil {
+		return "", err
+	}
 
 	buf := new(bytes.Buffer)
 	testRoot.SetOut(buf)
 	testRoot.SetErr(buf)
 	testRoot.SetArgs(args)
 
-	err := testRoot.Execute()
+	err = testRoot.Execute()
 	return buf.String(), err
 }
 
 func TestAuthLogin_KeyAndProjectIDFlags(t *testing.T) {
 	setupAuthTest(t)
 
 	// Execute login command with public and secret key flags and project ID
-	output, err := executeAuthCommand("auth", "login", "--public-key", "test-public-key", "--secret-key", "test-secret-key", "--project-id", "test-project-123")
+	output, err := executeAuthCommand(t.Context(), "auth", "login", "--public-key", "test-public-key", "--secret-key", "test-secret-key", "--project-id", "test-project-123")
 	if err != nil {
 		t.Fatalf("Login failed: %v", err)
 	}
@@ -120,7 +124,7 @@ func TestAuthLogin_KeyFlags_NoProjectID(t *testing.T) {
 
 	// Execute login command with only public and secret key flags (no project ID)
 	// This should fail since project ID is now required
-	_, err := executeAuthCommand("auth", "login", "--public-key", "test-public-key", "--secret-key", "test-secret-key")
+	_, err := executeAuthCommand(t.Context(), "auth", "login", "--public-key", "test-public-key", "--secret-key", "test-secret-key")
 	if err == nil {
 		t.Fatal("Expected login to fail without project ID, but it succeeded")
 	}
@@ -149,7 +153,7 @@ func TestAuthLogin_KeyAndProjectIDEnvironmentVariables(t *testing.T) {
 	defer os.Unsetenv("TIGER_PROJECT_ID")
 
 	// Execute login command with project ID flag but using env vars for keys
-	output, err := executeAuthCommand("auth", "login")
+	output, err := executeAuthCommand(t.Context(), "auth", "login")
 	if err != nil {
 		t.Fatalf("Login failed: %v", err)
 	}
@@ -184,7 +188,7 @@ func TestAuthLogin_KeyEnvironmentVariables_ProjectIDFlag(t *testing.T) {
 	defer os.Unsetenv("TIGER_SECRET_KEY")
 
 	// Execute login command with project ID flag but using env vars for keys
-	output, err := executeAuthCommand("auth", "login", "--project-id", "test-project-456")
+	output, err := executeAuthCommand(t.Context(), "auth", "login", "--project-id", "test-project-456")
 	if err != nil {
 		t.Fatalf("Login failed: %v", err)
 	}
@@ -434,7 +438,7 @@ func TestAuthLogin_OAuth_SingleProject(t *testing.T) {
 	}, "project-123")
 
 	// Execute login command - the mocked openBrowser will handle the callback automatically
-	output, err := executeAuthCommand("auth", "login")
+	output, err := executeAuthCommand(t.Context(), "auth", "login")
 
 	if err != nil {
 		t.Fatalf("Login failed: %v", err)
@@ -491,7 +495,7 @@ func TestAuthLogin_OAuth_MultipleProjects(t *testing.T) {
 	}
 
 	// Execute login command - both mocked functions will handle OAuth flow and project selection
-	output, err := executeAuthCommand("auth", "login")
+	output, err := executeAuthCommand(t.Context(), "auth", "login")
 
 	if err != nil {
 		t.Fatalf("Login failed: %v", err)
@@ -535,7 +539,7 @@ func TestAuthLogin_KeyringFallback(t *testing.T) {
 	// by ensuring the API key gets stored to file when keyring might not be available
 
 	// Execute login command with public and secret key flags and project ID
-	output, err := executeAuthCommand("auth", "login", "--public-key", "fallback-public", "--secret-key", "fallback-secret", "--project-id", "test-project-fallback")
+	output, err := executeAuthCommand(t.Context(), "auth", "login", "--public-key", "fallback-public", "--secret-key", "fallback-secret", "--project-id", "test-project-fallback")
 	if err != nil {
 		t.Fatalf("Login failed: %v", err)
 	}
@@ -571,7 +575,7 @@ func TestAuthLogin_KeyringFallback(t *testing.T) {
 	}
 
 	// Test status with file-only storage
-	output, err = executeAuthCommand("auth", "status")
+	output, err = executeAuthCommand(t.Context(), "auth", "status")
 	if err != nil {
 		t.Fatalf("Status failed with file storage: %v", err)
 	}
@@ -580,7 +584,7 @@ func TestAuthLogin_KeyringFallback(t *testing.T) {
 	}
 
 	// Test logout with file-only storage
-	output, err = executeAuthCommand("auth", "logout")
+	output, err = executeAuthCommand(t.Context(), "auth", "logout")
 	if err != nil {
 		t.Fatalf("Logout failed with file storage: %v", err)
 	}
@@ -607,7 +611,7 @@ func TestAuthLogin_EnvironmentVariable_FileOnly(t *testing.T) {
 	defer os.Unsetenv("TIGER_PROJECT_ID")
 
 	// Execute login command without any flags (all from env vars)
-	output, err := executeAuthCommand("auth", "login")
+	output, err := executeAuthCommand(t.Context(), "auth", "login")
 	if err != nil {
 		t.Fatalf("Login failed: %v", err)
 	}
@@ -652,7 +656,7 @@ func TestAuthStatus_LoggedIn(t *testing.T) {
 	}
 
 	// Execute status command
-	output, err := executeAuthCommand("auth", "status")
+	output, err := executeAuthCommand(t.Context(), "auth", "status")
 	if err != nil {
 		t.Fatalf("Status failed: %v", err)
 	}
@@ -666,7 +670,7 @@ func TestAuthStatus_NotLoggedIn(t *testing.T) {
 	setupAuthTest(t)
 
 	// Execute status command without being logged in
-	_, err := executeAuthCommand("auth", "status")
+	_, err := executeAuthCommand(t.Context(), "auth", "status")
 	if err == nil {
 		t.Fatal("Expected status to fail when not logged in")
 	}
@@ -693,7 +697,7 @@ func TestAuthLogout_Success(t *testing.T) {
 	}
 
 	// Execute logout command
-	output, err := executeAuthCommand("auth", "logout")
+	output, err := executeAuthCommand(t.Context(), "auth", "logout")
 	if err != nil {
 		t.Fatalf("Logout failed: %v", err)
 	}

internal/tiger/cmd/auth_validation_test.go

@@ -42,7 +42,7 @@ func TestAuthLogin_APIKeyValidationFailure(t *testing.T) {
 	defer config.RemoveCredentials()
 
 	// Execute login command with public and secret key flags - should fail validation
-	output, err := executeAuthCommand("auth", "login", "--public-key", "invalid-public", "--secret-key", "invalid-secret", "--project-id", "test-project-invalid")
+	output, err := executeAuthCommand(t.Context(), "auth", "login", "--public-key", "invalid-public", "--secret-key", "invalid-secret", "--project-id", "test-project-invalid")
 	if err == nil {
 		t.Fatal("Expected login to fail with invalid keys, but it succeeded")
 	}
@@ -96,7 +96,7 @@ func TestAuthLogin_APIKeyValidationSuccess(t *testing.T) {
 	defer config.RemoveCredentials()
 
 	// Execute login command with public and secret key flags - should succeed
-	output, err := executeAuthCommand("auth", "login", "--public-key", "valid-public", "--secret-key", "valid-secret", "--project-id", "test-project-valid")
+	output, err := executeAuthCommand(t.Context(), "auth", "login", "--public-key", "valid-public", "--secret-key", "valid-secret", "--project-id", "test-project-valid")
 	if err != nil {
 		t.Fatalf("Expected login to succeed with valid keys, got error: %v", err)
 	}