basedb: close db when fail ping (#942)

pkg/conn/basedb.go

@@ -16,12 +16,16 @@ package conn
 import (
 	"context"
 	"database/sql"
+	"errors"
 	"fmt"
 	"net/url"
 	"strconv"
 	"sync"
 	"sync/atomic"
 
+	"github.com/DATA-DOG/go-sqlmock"
+	"github.com/pingcap/failpoint"
+
 	"github.com/pingcap/dm/dm/config"
 	"github.com/pingcap/dm/pkg/retry"
 	"github.com/pingcap/dm/pkg/terror"
@@ -48,6 +52,9 @@ func init() {
 	DefaultDBProvider = &DefaultDBProviderImpl{}
 }
 
+// mock is used in unit test
+var mock sqlmock.Sqlmock
+
 // Apply will build BaseDB with DBConfig
 func (d *DefaultDBProviderImpl) Apply(config config.DBConfig) (*BaseDB, error) {
 	dsn := fmt.Sprintf("%s:%s@tcp(%s:%d)/?charset=utf8mb4&interpolateParams=true&maxAllowedPacket=%d",
@@ -96,8 +103,20 @@ func (d *DefaultDBProviderImpl) Apply(config config.DBConfig) (*BaseDB, error) {
 	if err != nil {
 		return nil, terror.DBErrorAdapt(err, terror.ErrDBDriverError)
 	}
-
-	if err = db.Ping(); err != nil {
+	failpoint.Inject("failDBPing", func(_ failpoint.Value) {
+		db.Close()
+		db, mock, _ = sqlmock.New()
+		mock.ExpectPing()
+		mock.ExpectClose()
+	})
+
+	err = db.Ping()
+	failpoint.Inject("failDBPing", func(_ failpoint.Value) {
+		err = errors.New("injected error")
+	})
+	if err != nil {
+		db.Close()
+		doFuncInClose()
 		return nil, terror.DBErrorAdapt(err, terror.ErrDBDriverError)
 	}
 

pkg/conn/basedb_test.go

@@ -15,8 +15,11 @@ package conn
 
 import (
 	. "github.com/pingcap/check"
+	"github.com/pingcap/failpoint"
 
 	sqlmock "github.com/DATA-DOG/go-sqlmock"
+
+	"github.com/pingcap/dm/dm/config"
 	tcontext "github.com/pingcap/dm/pkg/context"
 )
 
@@ -57,3 +60,17 @@ func (t *testBaseDBSuite) TestGetBaseConn(c *C) {
 	c.Assert(err, IsNil)
 	c.Assert(affected, Equals, 1)
 }
+
+func (t *testBaseDBSuite) TestFailDBPing(c *C) {
+	c.Assert(failpoint.Enable("github.com/pingcap/dm/pkg/conn/failDBPing", "return"), IsNil)
+	defer failpoint.Disable("github.com/pingcap/dm/pkg/conn/failDBPing")
+
+	cfg := config.DBConfig{User: "root", Host: "127.0.0.1", Port: 3306}
+	cfg.Adjust()
+	db, err := DefaultDBProvider.Apply(cfg)
+	c.Assert(db, IsNil)
+	c.Assert(err, NotNil)
+
+	err = mock.ExpectationsWereMet()
+	c.Assert(err, IsNil)
+}