fix panic when cant openDB

pingcap · Sep 17, 2021 · 5662d19 · 5662d19
1 parent 2e089a2
commit 5662d19
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 5 deletions.
diff --git a/components/bench/main.go b/components/bench/main.go
@@ -51,30 +51,32 @@ 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)
 	dsn := ds + dbName
 	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()
 		if strings.Contains(errString, unknownDB) {
 			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
@@ -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
@@ -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