From 072e503efc5cc6421f9202e4d421e4ac27f95347 Mon Sep 17 00:00:00 2001
From: Daniel Mitterdorfer <daniel.mitterdorfer@supabase.io>
Date: Fri, 22 May 2026 13:49:17 +0200
Subject: [PATCH] Propagate errors to top-level for TPCC

With this commit we propagate all errors in any of the TPCC subcommands to the
top-level and exit the application with an error message and a non-zero
exit code.

Closes #206
---
 cmd/go-tpc/main.go |  4 +++-
 cmd/go-tpc/misc.go | 18 +++++++++++-------
 cmd/go-tpc/tpcc.go | 27 +++++++++++++++++----------
 3 files changed, 31 insertions(+), 18 deletions(-)

diff --git a/cmd/go-tpc/main.go b/cmd/go-tpc/main.go
index c4125582..79d0794c 100644
--- a/cmd/go-tpc/main.go
+++ b/cmd/go-tpc/main.go
@@ -271,7 +271,9 @@ func main() {
 		}
 	}()
 
-	rootCmd.Execute()
+	if err := rootCmd.Execute(); err != nil {
+		os.Exit(1)
+	}
 
 	cancel()
 }
diff --git a/cmd/go-tpc/misc.go b/cmd/go-tpc/misc.go
index 2e8288c4..400bc2a2 100644
--- a/cmd/go-tpc/misc.go
+++ b/cmd/go-tpc/misc.go
@@ -9,16 +9,18 @@ import (
 	"github.com/pingcap/go-tpc/pkg/workload"
 )
 
-func checkPrepare(ctx context.Context, w workload.Workloader) {
+func checkPrepare(ctx context.Context, w workload.Workloader) error {
 	// skip preparation check in csv case
 	if w.Name() == "tpcc-csv" {
 		fmt.Println("Skip preparing checking. Please load CSV data into database and check later.")
-		return
+		return nil
 	}
 	if w.Name() == "tpcc" && tpccConfig.NoCheck {
-		return
+		return nil
 	}
 
+	errCh := make(chan error, threads)
+	defer close(errCh)
 	var wg sync.WaitGroup
 	wg.Add(threads)
 	for i := 0; i < threads; i++ {
@@ -29,12 +31,12 @@ func checkPrepare(ctx context.Context, w workload.Workloader) {
 			defer w.CleanupThread(ctx, index)
 
 			if err := w.CheckPrepare(ctx, index); err != nil {
-				fmt.Printf("check prepare failed, err %v\n", err)
-				return
+				errCh <- err
 			}
 		}(i)
 	}
 	wg.Wait()
+	return <-errCh
 }
 
 func execute(timeoutCtx context.Context, w workload.Workloader, action string, threads, index int) error {
@@ -101,7 +103,7 @@ func execute(timeoutCtx context.Context, w workload.Workloader, action string, t
 	return nil
 }
 
-func executeWorkload(ctx context.Context, w workload.Workloader, threads int, action string) {
+func executeWorkload(ctx context.Context, w workload.Workloader, threads int, action string) error {
 	var wg sync.WaitGroup
 	wg.Add(threads)
 
@@ -185,11 +187,13 @@ func executeWorkload(ctx context.Context, w workload.Workloader, threads int, ac
 
 	wg.Wait()
 
+	var checkErr error
 	if action == "prepare" {
 		// For prepare, we must check the data consistency after all prepare finished
-		checkPrepare(ctx, w)
+		checkErr = checkPrepare(ctx, w)
 	}
 	outputCancel()
 
 	<-ch
+	return checkErr
 }
diff --git a/cmd/go-tpc/tpcc.go b/cmd/go-tpc/tpcc.go
index 4dd5d80b..966853f1 100644
--- a/cmd/go-tpc/tpcc.go
+++ b/cmd/go-tpc/tpcc.go
@@ -18,7 +18,7 @@ import (
 
 var tpccConfig tpcc.Config
 
-func executeTpcc(action string) {
+func executeTpcc(action string) error {
 	if pprofAddr != "" {
 		go func() {
 			if err := http.ListenAndServe(pprofAddr, http.DefaultServeMux); err != nil {
@@ -81,10 +81,13 @@ func executeTpcc(action string) {
 	timeoutCtx, cancel := context.WithTimeout(globalCtx, totalTime)
 	defer cancel()
 
-	executeWorkload(timeoutCtx, w, threads, action)
+	if err := executeWorkload(timeoutCtx, w, threads, action); err != nil {
+		return err
+	}
 
 	fmt.Println("Finished")
 	w.OutputStats(true)
+	return nil
 }
 
 func registerTpcc(root *cobra.Command) {
@@ -99,8 +102,9 @@ func registerTpcc(root *cobra.Command) {
 	var cmdPrepare = &cobra.Command{
 		Use:   "prepare",
 		Short: "Prepare data for TPCC",
-		Run: func(cmd *cobra.Command, _ []string) {
-			executeTpcc("prepare")
+		RunE: func(cmd *cobra.Command, _ []string) error {
+			cmd.SilenceUsage = true
+			return executeTpcc("prepare")
 		},
 	}
 	cmdPrepare.PersistentFlags().BoolVar(&tpccConfig.NoCheck, "no-check", false, "TPCC prepare check, default false")
@@ -117,8 +121,9 @@ func registerTpcc(root *cobra.Command) {
 	var cmdRun = &cobra.Command{
 		Use:   "run",
 		Short: "Run workload",
-		Run: func(cmd *cobra.Command, _ []string) {
-			executeTpcc("run")
+		RunE: func(cmd *cobra.Command, _ []string) error {
+			cmd.SilenceUsage = true
+			return executeTpcc("run")
 		},
 	}
 	cmdRun.PersistentFlags().BoolVar(&tpccConfig.Wait, "wait", false, "including keying & thinking time described on TPC-C Standard Specification")
@@ -129,16 +134,18 @@ func registerTpcc(root *cobra.Command) {
 	var cmdCleanup = &cobra.Command{
 		Use:   "cleanup",
 		Short: "Cleanup data for the workload",
-		Run: func(cmd *cobra.Command, _ []string) {
-			executeTpcc("cleanup")
+		RunE: func(cmd *cobra.Command, _ []string) error {
+			cmd.SilenceUsage = true
+			return executeTpcc("cleanup")
 		},
 	}
 
 	var cmdCheck = &cobra.Command{
 		Use:   "check",
 		Short: "Check data consistency for the workload",
-		Run: func(cmd *cobra.Command, _ []string) {
-			executeTpcc("check")
+		RunE: func(cmd *cobra.Command, _ []string) error {
+			cmd.SilenceUsage = true
+			return executeTpcc("check")
 		},
 	}