Repository: pingcap/tidb-lightning
Pull request #280 (2020-03-10 05:41:38)

Make sure the session param takes effect

<!--
Thank you for contributing to TiDB-Lightning! Please read the [CONTRIBUTING](https://github.com/pingcap/tidb-lightning/blob/master/CONTRIBUTING.md) document **BEFORE** filing this PR.
-->

### What problem does this PR solve? <!--add issue link with summary if exists-->
Make sure the session param takes effect.

### What is changed and how it works?
Set session var for every new conn. For every new connection created by the driver, it will set the params first.

Note the previous `setSessionConcurrencyVars` will set one connection
from the db connection pool, so we can't make sure the session will take
effect later using the `sql.DB`.

### Check List <!--REMOVE the items that are not applicable-->

Tests <!-- At least one of them must be included. -->

 - Unit test
 - Manual test (add detailed scripts or steps below)
check tidb log for the session params setted using the dsn.

Side effects



Related changes

 - Need to cherry-pick to the release branch

 - Need to be included in the release note


Patch:
diff --git a/lightning/common/util.go b/lightning/common/util.go
index 581573a..c650c1c 100644
--- a/lightning/common/util.go
+++ b/lightning/common/util.go
@@ -23,6 +23,7 @@ import (
 	"io/ioutil"
 	"net"
 	"net/http"
+	"net/url"
 	"os"
 	"reflect"
 	"regexp"
@@ -55,12 +56,19 @@ type MySQLConnectParam struct {
 	SQLMode          string
 	MaxAllowedPacket uint64
 	TLS              string
+	Vars             map[string]string
 }
 
 func (param *MySQLConnectParam) ToDSN() string {
-	return fmt.Sprintf("%s:%s@tcp(%s:%d)/?charset=utf8mb4&sql_mode='%s'&maxAllowedPacket=%d&tls=%s",
+	dsn := fmt.Sprintf("%s:%s@tcp(%s:%d)/?charset=utf8mb4&sql_mode='%s'&maxAllowedPacket=%d&tls=%s",
 		param.User, param.Password, param.Host, param.Port,
 		param.SQLMode, param.MaxAllowedPacket, param.TLS)
+
+	for k, v := range param.Vars {
+		dsn += fmt.Sprintf("&%s=%s", k, url.QueryEscape(v))
+	}
+
+	return dsn
 }
 
 func (param *MySQLConnectParam) Connect() (*sql.DB, error) {
diff --git a/lightning/restore/restore.go b/lightning/restore/restore.go
index 8c5a749..8dd46ee 100644
--- a/lightning/restore/restore.go
+++ b/lightning/restore/restore.go
@@ -1002,8 +1002,6 @@ func (t *TableRestore) postProcess(ctx context.Context, rc *RestoreController, c
 		return nil
 	}
 
-	setSessionConcurrencyVars(ctx, rc.tidbMgr.db, rc.cfg.TiDB)
-
 	// 3. alter table set auto_increment
 	if cp.Status < CheckpointStatusAlteredAutoInc {
 		rc.alterTableLock.Lock()
@@ -1466,15 +1464,6 @@ type RemoteChecksum struct {
 	TotalBytes uint64
 }
 
-func setSessionConcurrencyVars(ctx context.Context, db *sql.DB, dsn config.DBStore) {
-	common.SQLWithRetry{DB: db, Logger: log.L()}.Exec(ctx, "set session concurrency variables", `SET
-		SESSION tidb_build_stats_concurrency = ?,
-		SESSION tidb_distsql_scan_concurrency = ?,
-		SESSION tidb_index_serial_scan_concurrency = ?,
-		SESSION tidb_checksum_table_concurrency = ?;
-	`, dsn.BuildStatsConcurrency, dsn.DistSQLScanConcurrency, dsn.IndexSerialScanConcurrency, dsn.ChecksumTableConcurrency)
-}
-
 // DoChecksum do checksum for tables.
 // table should be in <db>.<table>, format.  e.g. foo.bar
 func DoChecksum(ctx context.Context, db *sql.DB, table string) (*RemoteChecksum, error) {
diff --git a/lightning/restore/tidb.go b/lightning/restore/tidb.go
index 1305bcc..df47cfc 100644
--- a/lightning/restore/tidb.go
+++ b/lightning/restore/tidb.go
@@ -17,6 +17,7 @@ import (
 	"context"
 	"database/sql"
 	"fmt"
+	"strconv"
 	"strings"
 
 	"github.com/pingcap/errors"
@@ -49,6 +50,12 @@ func NewTiDBManager(dsn config.DBStore, tls *common.TLS) (*TiDBManager, error) {
 		SQLMode:          dsn.StrSQLMode,
 		MaxAllowedPacket: dsn.MaxAllowedPacket,
 		TLS:              dsn.TLS,
+		Vars: map[string]string{
+			"tidb_build_stats_concurrency":       strconv.Itoa(dsn.BuildStatsConcurrency),
+			"tidb_distsql_scan_concurrency":      strconv.Itoa(dsn.DistSQLScanConcurrency),
+			"tidb_index_serial_scan_concurrency": strconv.Itoa(dsn.IndexSerialScanConcurrency),
+			"tidb_checksum_table_concurrency":    strconv.Itoa(dsn.ChecksumTableConcurrency),
+		},
 	}
 	db, err := param.Connect()
 	if err != nil {


Test patch:
diff --git a/lightning/common/util_test.go b/lightning/common/util_test.go
index 9ac69c8..ff8893d 100644
--- a/lightning/common/util_test.go
+++ b/lightning/common/util_test.go
@@ -129,8 +129,11 @@ func (s *utilSuite) TestToDSN(c *C) {
 		SQLMode:          "strict",
 		MaxAllowedPacket: 1234,
 		TLS:              "cluster",
+		Vars: map[string]string{
+			"tidb_distsql_scan_concurrency": "1",
+		},
 	}
-	c.Assert(param.ToDSN(), Equals, "root:123456@tcp(127.0.0.1:4000)/?charset=utf8mb4&sql_mode='strict'&maxAllowedPacket=1234&tls=cluster")
+	c.Assert(param.ToDSN(), Equals, "root:123456@tcp(127.0.0.1:4000)/?charset=utf8mb4&sql_mode='strict'&maxAllowedPacket=1234&tls=cluster&tidb_distsql_scan_concurrency=1")
 }
 
 func (s *utilSuite) TestIsContextCanceledError(c *C) {
diff --git a/lightning/restore/restore_test.go b/lightning/restore/restore_test.go
index 10f0fb2..2b56d61 100644
--- a/lightning/restore/restore_test.go
+++ b/lightning/restore/restore_test.go
@@ -264,32 +264,6 @@ func (s *restoreSuite) TestDoChecksumWithErrorAndLongOriginalLifetime(c *C) {
 	c.Assert(mock.ExpectationsWereMet(), IsNil)
 }
 
-func (s *restoreSuite) TestSetSessionConcurrencyVars(c *C) {
-	db, mock, err := sqlmock.New()
-	c.Assert(err, IsNil)
-
-	mock.ExpectExec(
-		`SET\s+`+
-			`SESSION tidb_build_stats_concurrency = \?,\s+`+
-			`SESSION tidb_distsql_scan_concurrency = \?,\s+`+
-			`SESSION tidb_index_serial_scan_concurrency = \?,\s+`+
-			`SESSION tidb_checksum_table_concurrency = \?`).
-		WithArgs(123, 456, 789, 543).
-		WillReturnResult(sqlmock.NewResult(1, 4))
-	mock.ExpectClose()
-
-	ctx := context.Background()
-	setSessionConcurrencyVars(ctx, db, config.DBStore{
-		BuildStatsConcurrency:      123,
-		DistSQLScanConcurrency:     456,
-		IndexSerialScanConcurrency: 789,
-		ChecksumTableConcurrency:   543,
-	})
-
-	c.Assert(db.Close(), IsNil)
-	c.Assert(mock.ExpectationsWereMet(), IsNil)
-}
-
 var _ = Suite(&tableRestoreSuite{})
 
 type tableRestoreSuite struct {