From 20455dbd7b45e5772c7633c9bab9e5c09c70ac92 Mon Sep 17 00:00:00 2001
From: nexustar <inexustar@gmail.com>
Date: Sat, 18 Sep 2021 11:40:42 +0800
Subject: [PATCH] fix panic when cant openDB on tiup-bench (#1557)

---
 components/bench/main.go | 8 +++++---
 components/bench/tpcc.go | 6 +++++-
 components/bench/tpch.go | 8 +++++++-
 3 files changed, 17 insertions(+), 5 deletions(-)

diff --git a/components/bench/main.go b/components/bench/main.go
index 65ab85b105..379e2ab08f 100644
--- a/components/bench/main.go
+++ b/components/bench/main.go
@@ -51,7 +51,7 @@ func closeDB() {
 	globalDB = nil
 }
 
-func openDB() {
+func openDB() error {
 	// TODO: support other drivers
 	var tmpDB *sql.DB
 	ds := fmt.Sprintf("%s:%s@tcp(%s:%d)/", user, password, host, port)
@@ -59,7 +59,7 @@ func openDB() {
 	var err error
 	globalDB, err = sql.Open(mysqlDriver, dsn)
 	if err != nil {
-		panic(err)
+		return err
 	}
 	if err := globalDB.Ping(); err != nil {
 		errString := err.Error()
@@ -67,14 +67,16 @@ func openDB() {
 			tmpDB, _ = sql.Open(mysqlDriver, ds)
 			defer tmpDB.Close()
 			if _, err := tmpDB.Exec(createDBDDL + dbName); err != nil {
-				panic(fmt.Errorf("failed to create database, err %v", err))
+				return fmt.Errorf("failed to create database, err %v", err)
 			}
 		} else {
 			globalDB = nil
+			return err
 		}
 	} else {
 		globalDB.SetMaxIdleConns(threads + 1)
 	}
+	return nil
 }
 
 func main() {
diff --git a/components/bench/tpcc.go b/components/bench/tpcc.go
index efb6293717..6360be49e0 100644
--- a/components/bench/tpcc.go
+++ b/components/bench/tpcc.go
@@ -26,7 +26,11 @@ func executeTpcc(action string) {
 	}
 	runtime.GOMAXPROCS(maxProcs)
 
-	openDB()
+	if err := openDB(); err != nil {
+		fmt.Println(err)
+		fmt.Println("Cannot open database, pleae check it (ip/port/username/password)")
+		os.Exit(1)
+	}
 
 	tpccConfig.DBName = dbName
 	tpccConfig.Threads = threads
diff --git a/components/bench/tpch.go b/components/bench/tpch.go
index aae37fe0be..8eb1e2eed9 100644
--- a/components/bench/tpch.go
+++ b/components/bench/tpch.go
@@ -2,6 +2,8 @@ package main
 
 import (
 	"context"
+	"fmt"
+	"os"
 	"strings"
 
 	"github.com/pingcap/go-tpc/tpch"
@@ -11,7 +13,11 @@ import (
 var tpchConfig tpch.Config
 
 func executeTpch(action string, _ []string) {
-	openDB()
+	if err := openDB(); err != nil {
+		fmt.Println(err)
+		fmt.Println("Cannot open database, pleae check it (ip/port/username/password)")
+		os.Exit(1)
+	}
 	defer closeDB()
 
 	tpchConfig.DBName = dbName