Refactor context handling

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, press control-C again to exit", 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/api/client_util.go

@@ -66,18 +66,18 @@ func NewTigerClient(apiKey string) (*ClientWithResponses, error) {
 }
 
 // ValidateAPIKey validates the API key by making a test API call
-func ValidateAPIKey(apiKey string, projectID string) error {
+func ValidateAPIKey(ctx context.Context, apiKey string, projectID string) error {
 	client, err := NewTigerClient(apiKey)
 	if err != nil {
 		return fmt.Errorf("failed to create client: %w", err)
 	}
 
-	return ValidateAPIKeyWithClient(client, projectID)
+	return ValidateAPIKeyWithClient(ctx, client, projectID)
 }
 
 // ValidateAPIKeyWithClient validates the API key using the provided client interface
-func ValidateAPIKeyWithClient(client ClientWithResponsesInterface, projectID string) error {
-	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+func ValidateAPIKeyWithClient(ctx context.Context, client ClientWithResponsesInterface, projectID string) error {
+	ctx, cancel := context.WithTimeout(ctx, 10*time.Second)
 	defer cancel()
 
 	// Use provided project ID if available, otherwise use a dummy one

internal/tiger/api/client_util_test.go

@@ -95,7 +95,7 @@ func TestValidateAPIKeyWithClient(t *testing.T) {
 			mockClient := mocks.NewMockClientWithResponsesInterface(ctrl)
 			tt.setupMock(mockClient)
 
-			err := api.ValidateAPIKeyWithClient(mockClient, "")
+			err := api.ValidateAPIKeyWithClient(context.Background(), mockClient, "")
 
 			if tt.expectedError == "" {
 				if err != nil {

internal/tiger/cmd/auth.go

@@ -2,9 +2,9 @@ package cmd
 
 import (
 	"bufio"
+	"context"
 	"fmt"
 	"os"
-	"strings"
 
 	"github.com/spf13/cobra"
 	"golang.org/x/term"
@@ -97,13 +97,13 @@ Examples:
 					out: cmd.OutOrStdout(),
 				}
 
-				creds, err = l.loginWithOAuth()
+				creds, err = l.loginWithOAuth(cmd.Context())
 				if err != nil {
 					return err
 				}
 			} else if creds.publicKey == "" || creds.secretKey == "" || creds.projectID == "" {
 				// If some credentials were provided, prompt for missing ones
-				creds, err = promptForCredentials(cfg.ConsoleURL, creds)
+				creds, err = promptForCredentials(cmd.Context(), cfg.ConsoleURL, creds)
 				if err != nil {
 					return fmt.Errorf("failed to get credentials: %w", err)
 				}
@@ -122,7 +122,7 @@ Examples:
 
 			// Validate the API key by making a test API call
 			fmt.Fprintln(cmd.OutOrStdout(), "Validating API key...")
-			if err := validateAPIKeyForLogin(apiKey, creds.projectID); err != nil {
+			if err := validateAPIKeyForLogin(cmd.Context(), apiKey, creds.projectID); err != nil {
 				return fmt.Errorf("API key validation failed: %w", err)
 			}
 
@@ -209,7 +209,7 @@ func flagOrEnvVar(flagVal, envVarName string) string {
 }
 
 // promptForCredentials prompts the user to enter any missing credentials
-func promptForCredentials(consoleURL string, creds credentials) (credentials, error) {
+func promptForCredentials(ctx context.Context, consoleURL string, creds credentials) (credentials, error) {
 	// Check if we're in a terminal for interactive input
 	if !util.IsTerminal(os.Stdin) {
 		return credentials{}, fmt.Errorf("TTY not detected - credentials required. Use flags (--public-key, --secret-key, --project-id) or environment variables (TIGER_PUBLIC_KEY, TIGER_SECRET_KEY, TIGER_PROJECT_ID)")
@@ -222,32 +222,35 @@ func promptForCredentials(consoleURL string, creds credentials) (credentials, er
 	// Prompt for public key if missing
 	if creds.publicKey == "" {
 		fmt.Print("Enter your public key: ")
-		publicKey, err := reader.ReadString('\n')
+		publicKey, err := readString(ctx, func() (string, error) { return reader.ReadString('\n') })
 		if err != nil {
 			return credentials{}, err
 		}
-		creds.publicKey = strings.TrimSpace(publicKey)
+		creds.publicKey = publicKey
 	}
 
 	// Prompt for secret key if missing
 	if creds.secretKey == "" {
 		fmt.Print("Enter your secret key: ")
-		bytePassword, err := term.ReadPassword(int(os.Stdin.Fd()))
+		password, err := readString(ctx, func() (string, error) {
+			val, err := term.ReadPassword(int(os.Stdin.Fd()))
+			return string(val), err
+		})
 		if err != nil {
 			return credentials{}, err
 		}
 		fmt.Println() // Print newline after hidden input
-		creds.secretKey = strings.TrimSpace(string(bytePassword))
+		creds.secretKey = password
 	}
 
 	// Prompt for project ID if missing
 	if creds.projectID == "" {
 		fmt.Print("Enter your project ID: ")
-		projectID, err := reader.ReadString('\n')
+		projectID, err := readString(ctx, func() (string, error) { return reader.ReadString('\n') })
 		if err != nil {
 			return credentials{}, err
 		}
-		creds.projectID = strings.TrimSpace(projectID)
+		creds.projectID = projectID
 	}
 
 	return creds, nil

internal/tiger/cmd/auth_test.go

@@ -2,6 +2,7 @@ package cmd
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -27,7 +28,7 @@ func setupAuthTest(t *testing.T) string {
 
 	// Mock the API key validation for testing
 	originalValidator := validateAPIKeyForLogin
-	validateAPIKeyForLogin = func(apiKey, projectID string) error {
+	validateAPIKeyForLogin = func(ctx context.Context, apiKey, projectID string) error {
 		// Always return success for testing
 		return nil
 	}
@@ -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

@@ -1,6 +1,7 @@
 package cmd
 
 import (
+	"context"
 	"errors"
 	"os"
 	"strings"
@@ -23,7 +24,7 @@ func TestAuthLogin_APIKeyValidationFailure(t *testing.T) {
 	originalValidator := validateAPIKeyForLogin
 
 	// Mock the validator to return an error
-	validateAPIKeyForLogin = func(apiKey, projectID string) error {
+	validateAPIKeyForLogin = func(ctx context.Context, apiKey, projectID string) error {
 		return errors.New("invalid API key: authentication failed")
 	}
 
@@ -42,7 +43,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")
 	}
@@ -77,7 +78,7 @@ func TestAuthLogin_APIKeyValidationSuccess(t *testing.T) {
 	originalValidator := validateAPIKeyForLogin
 
 	// Mock the validator to return success
-	validateAPIKeyForLogin = func(apiKey, projectID string) error {
+	validateAPIKeyForLogin = func(ctx context.Context, apiKey, projectID string) error {
 		return nil // Success
 	}
 
@@ -96,7 +97,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)
 	}