Repository: pingcap/tidb-tools (go)

Issue:
Add max_execution_time parameter for sync diff inspector
## Feature Request

**Is your feature request related to a problem? Please describe:**
* We enforce cluster max_execution_time as a safety mechanism at 30s and find that sync diff inspector needs longer runtimes to avoid experiencing errors due to canceled requests.
* This patch allows for setting a different session level max_execution_time value to override the cluster global value.
* It maintains the default behavior (no session level limit) unless an additional new argument is provided to sync-diff-inspector

**Describe the feature you'd like:**
* Session level override for max_execution_time

**Describe alternatives you've considered:**
* No good ones came to mind and would require relaxing our reliability stance.

**Teachability, Documentation, Adoption, Migration Strategy:**
* N/A, upgrade and use argument

Filing follow up PR using the patch we run in production.

Pull request description:
sync-diff-inspector: add session configuration in toml
<!--
Thank you for contributing to TiDB! Please read TiDB's [CONTRIBUTING](https://github.com/pingcap/tidb/blob/master/CONTRIBUTING.md) document **BEFORE** filing this PR.
-->

### What problem does this PR solve?
<!--
Please create an issue first to describe the problem.

There MUST be one line starting with "Issue Number:  " and 
linking the relevant issues via the "close" or "ref".

For more info, check https://pingcap.github.io/tidb-dev-guide/contribute-to-tidb/contribute-code.html#referring-to-an-issue.
-->

Issue Number: ref #836, close #844

### What is changed and how it works?

Add session variable configuration for each data source.

### Check List <!--REMOVE the items that are not applicable-->

Tests <!-- At least one of them must be included. -->

 - Unit test
 - Integration test
 - Manual test (add detailed scripts or steps below)
 - No code

Code changes

 - Has exported function/method change
 - Has exported variable/fields change
 - Has interface methods change
 - Has persistent data change

Side effects

 - Possible performance regression
 - Increased code complexity
 - Breaking backward compatibility

Related changes

 - Need to cherry-pick to the release branch
 - Need to update the documentation
 - Need to be included in the release note


Patch:
diff --git a/OWNERS b/OWNERS
index 8582ca6..b9fbc15 100644
--- a/OWNERS
+++ b/OWNERS
@@ -99,6 +99,7 @@ reviewers:
   - ChenPeng2013
   - dhysum
   - fzzf678
+  - gengliqi
   - ghazalfamilyusa
   - JmPotato
   - js00070
diff --git a/sync_diff_inspector/config/config.go b/sync_diff_inspector/config/config.go
index 7fd99cc..b0f311e 100644
--- a/sync_diff_inspector/config/config.go
+++ b/sync_diff_inspector/config/config.go
@@ -121,7 +121,8 @@ type DataSource struct {
 	Router         *router.Table
 	RouteTargetSet map[string]struct{} `json:"-"`
 
-	Conn *sql.DB
+	Conn          *sql.DB
+	SessionConfig SessionConfig `toml:"session" json:"session"`
 }
 
 // IsAutoSnapshot returns true if the tidb_snapshot is expected to automatically
@@ -192,6 +193,14 @@ func (d *DataSource) ToDriverConfig() *mysql.Config {
 		cfg.TLSConfig = d.Security.TLSName
 	}
 
+	for param, value := range d.SessionConfig {
+		switch v := value.(type) {
+		case string:
+			cfg.Params[param] = "'" + strings.ReplaceAll(v, "'", "''") + "'"
+		default:
+			cfg.Params[param] = fmt.Sprint(v)
+		}
+	}
 	return cfg
 }
 
@@ -357,6 +366,9 @@ func (t *TaskConfig) ComputeConfigHash() (string, error) {
 	return fmt.Sprintf("%x", sha256.Sum256(hash)), nil
 }
 
+// SessionConfig the the session level configuration for data source.
+type SessionConfig map[string]any
+
 // Config is the configuration.
 type Config struct {
 	*flag.FlagSet `json:"-"`
diff --git a/sync_diff_inspector/config/config.toml b/sync_diff_inspector/config/config.toml
index b555721..3c53cc0 100644
--- a/sync_diff_inspector/config/config.toml
+++ b/sync_diff_inspector/config/config.toml
@@ -20,6 +20,8 @@ check-struct-only = false
     port = 3306
     user = "root"
     password = ""
+    session.max_execution_time = 3600
+    session.net_read_timeout = 60
 # MySQL doesn't has snapshot config
 
 [data-sources.tidb0]
@@ -27,6 +29,8 @@ check-struct-only = false
     port = 4000
     user = "root"
     password = ""
+    session.tidb_opt_prefer_range_scan = "ON"
+    session.max_execution_time = 86400
 
 # Support tls connection
     # security.ca-path = "..."
diff --git a/sync_diff_inspector/config/config_sharding.toml b/sync_diff_inspector/config/config_sharding.toml
index 59a70c5..53bc6e6 100644
--- a/sync_diff_inspector/config/config_sharding.toml
+++ b/sync_diff_inspector/config/config_sharding.toml
@@ -49,6 +49,8 @@ check-struct-only = false
     port = 4000
     user = "root"
     password = ""
+    session.tidb_opt_prefer_range_scan = "ON"
+    session.max_execution_time = 86400
     # remove comment if use tidb's snapshot data
     # snapshot = "2016-10-08 16:45:26"
 
diff --git a/sync_diff_inspector/source/common/conn.go b/sync_diff_inspector/source/common/conn.go
index 286d1f5..3edefed 100755
--- a/sync_diff_inspector/source/common/conn.go
+++ b/sync_diff_inspector/source/common/conn.go
@@ -16,11 +16,15 @@ package common
 import (
 	"database/sql"
 	"encoding/base64"
+	"fmt"
 
 	"github.com/go-sql-driver/mysql"
 	"github.com/pingcap/errors"
 	"github.com/pingcap/failpoint"
+	"github.com/pingcap/log"
+	"github.com/pingcap/tidb-tools/sync_diff_inspector/config"
 	tmysql "github.com/pingcap/tidb/pkg/parser/mysql"
+	"go.uber.org/zap"
 )
 
 func tryConnectMySQL(cfg *mysql.Config) (*sql.DB, error) {
@@ -43,8 +47,39 @@ func tryConnectMySQL(cfg *mysql.Config) (*sql.DB, error) {
 	return db, nil
 }
 
+func verifyParams(db *sql.DB, sessionCfg *config.SessionConfig) error {
+	if sessionCfg == nil {
+		return nil
+	}
+	for param, value := range *sessionCfg {
+		res, err := db.Query("show session variables like ?", param)
+		if err != nil {
+			return err
+		}
+		//nolint: errcheck
+		defer res.Close()
+		if res.Next() {
+			var paramName, actual string
+			if err := res.Scan(&paramName, &actual); err != nil {
+				return err
+			}
+			expected := fmt.Sprint(value)
+			if actual != expected {
+				log.Warn("The session variable was set, but the database returned a different value",
+					zap.String("variable", param),
+					zap.String("configured_value", expected),
+					zap.String("effective_value", actual),
+				)
+			}
+		} else {
+			return fmt.Errorf("parameter %s not found", param)
+		}
+	}
+	return nil
+}
+
 // ConnectMySQL creates sql.DB used for select data
-func ConnectMySQL(cfg *mysql.Config, num int) (db *sql.DB, err error) {
+func ConnectMySQL(sessionCfg *config.SessionConfig, cfg *mysql.Config, num int) (db *sql.DB, err error) {
 	defer func() {
 		if err == nil && db != nil {
 			// SetMaxOpenConns and SetMaxIdleConns for connection to avoid error like
@@ -55,7 +90,7 @@ func ConnectMySQL(cfg *mysql.Config, num int) (db *sql.DB, err error) {
 	}()
 	// Try plain password first.
 	db, firstErr := tryConnectMySQL(cfg)
-	if firstErr == nil {
+	if firstErr == nil && verifyParams(db, sessionCfg) == nil {
 		return db, nil
 	}
 	// If access is denied and password is encoded by base64, try the decoded string as well.
@@ -64,7 +99,7 @@ func ConnectMySQL(cfg *mysql.Config, num int) (db *sql.DB, err error) {
 		if password, decodeErr := base64.StdEncoding.DecodeString(cfg.Passwd); decodeErr == nil && string(password) != cfg.Passwd {
 			cfg.Passwd = string(password)
 			db2, err := tryConnectMySQL(cfg)
-			if err == nil {
+			if err == nil && verifyParams(db2, sessionCfg) == nil {
 				return db2, nil
 			}
 		}
diff --git a/sync_diff_inspector/source/source.go b/sync_diff_inspector/source/source.go
index c0c8271..9fd7106 100644
--- a/sync_diff_inspector/source/source.go
+++ b/sync_diff_inspector/source/source.go
@@ -237,7 +237,7 @@ func buildSourceFromCfg(ctx context.Context, tableDiffs []*common.TableDiff, con
 }
 
 func getAutoSnapshotPosition(cfg *mysql.Config) (string, string, error) {
-	tmpConn, err := common.ConnectMySQL(cfg, 2)
+	tmpConn, err := common.ConnectMySQL(nil, cfg, 2)
 	if err != nil {
 		return "", "", errors.Annotatef(err, "connecting to auto-position tidb_snapshot failed")
 	}
@@ -269,22 +269,32 @@ func initDBConn(ctx context.Context, cfg *config.Config) error {
 	}
 	// we had `cfg.SplitThreadCount` producers and `cfg.CheckThreadCount` consumer to use db connections maybe and `cfg.CheckThreadCount` splitter to split buckets.
 	// so the connection count need to be cfg.SplitThreadCount + cfg.CheckThreadCount + cfg.CheckThreadCount.
-	targetConn, err := common.ConnectMySQL(cfg.Task.TargetInstance.ToDriverConfig(), cfg.SplitThreadCount+2*cfg.CheckThreadCount)
+	targetConn, err := common.ConnectMySQL(
+		&cfg.Task.TargetInstance.SessionConfig,
+		cfg.Task.TargetInstance.ToDriverConfig(),
+		cfg.SplitThreadCount+2*cfg.CheckThreadCount,
+	)
 	if err != nil {
+		log.Error("failed to configure session", zap.String("data-source", cfg.Task.Target), zap.Error(err))
 		return errors.Trace(err)
 	}
 
 	cfg.Task.TargetInstance.Conn = targetConn
 
-	for _, source := range cfg.Task.SourceInstances {
+	for sourceIdx, source := range cfg.Task.SourceInstances {
 		// If it is still set to AUTO it means it was not set on the target.
 		// We require it to be set to AUTO on both.
 		if source.IsAutoSnapshot() {
 			return errors.Errorf("'auto' snapshot should be set on both target and source")
 		}
 		// connect source db with target db time_zone
-		conn, err := common.ConnectMySQL(source.ToDriverConfig(), cfg.SplitThreadCount+2*cfg.CheckThreadCount)
+		conn, err := common.ConnectMySQL(
+			&source.SessionConfig,
+			source.ToDriverConfig(),
+			cfg.SplitThreadCount+2*cfg.CheckThreadCount,
+		)
 		if err != nil {
+			log.Error("failed to configure session", zap.String("data-source", cfg.Task.Source[sourceIdx]), zap.Error(err))
 			return errors.Trace(err)
 		}
 		source.Conn = conn


Test patch:
diff --git a/sync_diff_inspector/config/config_test.go b/sync_diff_inspector/config/config_test.go
index 7c12c26..d07286f 100644
--- a/sync_diff_inspector/config/config_test.go
+++ b/sync_diff_inspector/config/config_test.go
@@ -50,10 +50,10 @@ func TestParseConfig(t *testing.T) {
 
 	// we might not use the same config to run this test. e.g. MYSQL_PORT can be 4000
 	require.JSONEq(t, cfg.String(),
-		"{\"check-thread-count\":4,\"split-thread-count\":5,\"export-fix-sql\":true,\"check-struct-only\":false,\"dm-addr\":\"\",\"dm-task\":\"\",\"data-sources\":{\"mysql1\":{\"host\":\"127.0.0.1\",\"port\":3306,\"user\":\"root\",\"password\":\"******\",\"sql-mode\":\"\",\"snapshot\":\"\",\"security\":null,\"route-rules\":[\"rule1\",\"rule2\"],\"Router\":{\"Selector\":{}},\"Conn\":null},\"mysql2\":{\"host\":\"127.0.0.1\",\"port\":3306,\"user\":\"root\",\"password\":\"******\",\"sql-mode\":\"\",\"snapshot\":\"\",\"security\":null,\"route-rules\":[\"rule1\",\"rule2\"],\"Router\":{\"Selector\":{}},\"Conn\":null},\"mysql3\":{\"host\":\"127.0.0.1\",\"port\":3306,\"user\":\"root\",\"password\":\"******\",\"sql-mode\":\"\",\"snapshot\":\"\",\"security\":null,\"route-rules\":[\"rule1\",\"rule3\"],\"Router\":{\"Selector\":{}},\"Conn\":null},\"tidb0\":{\"host\":\"127.0.0.1\",\"port\":4000,\"user\":\"root\",\"password\":\"******\",\"sql-mode\":\"\",\"snapshot\":\"\",\"security\":null,\"route-rules\":null,\"Router\":{\"Selector\":{}},\"Conn\":null}},\"routes\":{\"rule1\":{\"schema-pattern\":\"test_*\",\"table-pattern\":\"t_*\",\"target-schema\":\"test\",\"target-table\":\"t\"},\"rule2\":{\"schema-pattern\":\"test2_*\",\"table-pattern\":\"t2_*\",\"target-schema\":\"test2\",\"target-table\":\"t2\"},\"rule3\":{\"schema-pattern\":\"test2_*\",\"table-pattern\":\"t2_*\",\"target-schema\":\"test\",\"target-table\":\"t\"}},\"table-configs\":{\"config1\":{\"target-tables\":[\"schema*.table*\",\"test2.t2\"],\"Schema\":\"\",\"Table\":\"\",\"ConfigIndex\":0,\"HasMatched\":false,\"IgnoreColumns\":[\"\",\"\"],\"Fields\":[\"\"],\"Range\":\"age \\u003e 10 AND age \\u003c 20\",\"TargetTableInfo\":null,\"Collation\":\"\",\"chunk-size\":0}},\"task\":{\"source-instances\":[\"mysql1\",\"mysql2\",\"mysql3\"],\"source-routes\":null,\"target-instance\":\"tidb0\",\"target-check-tables\":[\"schema*.table*\",\"!c.*\",\"test2.t2\"],\"target-configs\":[\"config1\"],\"output-dir\":\"/tmp/output/config\",\"SourceInstances\":[{\"host\":\"127.0.0.1\",\"port\":3306,\"user\":\"root\",\"password\":\"******\",\"sql-mode\":\"\",\"snapshot\":\"\",\"security\":null,\"route-rules\":[\"rule1\",\"rule2\"],\"Router\":{\"Selector\":{}},\"Conn\":null},{\"host\":\"127.0.0.1\",\"port\":3306,\"user\":\"root\",\"password\":\"******\",\"sql-mode\":\"\",\"snapshot\":\"\",\"security\":null,\"route-rules\":[\"rule1\",\"rule2\"],\"Router\":{\"Selector\":{}},\"Conn\":null},{\"host\":\"127.0.0.1\",\"port\":3306,\"user\":\"root\",\"password\":\"******\",\"sql-mode\":\"\",\"snapshot\":\"\",\"security\":null,\"route-rules\":[\"rule1\",\"rule3\"],\"Router\":{\"Selector\":{}},\"Conn\":null}],\"TargetInstance\":{\"host\":\"127.0.0.1\",\"port\":4000,\"user\":\"root\",\"password\":\"******\",\"sql-mode\":\"\",\"snapshot\":\"\",\"security\":null,\"route-rules\":null,\"Router\":{\"Selector\":{}},\"Conn\":null},\"TargetTableConfigs\":[{\"target-tables\":[\"schema*.table*\",\"test2.t2\"],\"Schema\":\"\",\"Table\":\"\",\"ConfigIndex\":0,\"HasMatched\":false,\"IgnoreColumns\":[\"\",\"\"],\"Fields\":[\"\"],\"Range\":\"age \\u003e 10 AND age \\u003c 20\",\"TargetTableInfo\":null,\"Collation\":\"\",\"chunk-size\":0}],\"TargetCheckTables\":[{},{},{}],\"FixDir\":\"/tmp/output/config/fix-on-tidb0\",\"CheckpointDir\":\"/tmp/output/config/checkpoint\",\"HashFile\":\"\"},\"ConfigFile\":\"config_sharding.toml\",\"PrintVersion\":false}")
+		"{\"check-thread-count\":4,\"split-thread-count\":5,\"export-fix-sql\":true,\"check-struct-only\":false,\"dm-addr\":\"\",\"dm-task\":\"\",\"data-sources\":{\"mysql1\":{\"host\":\"127.0.0.1\",\"port\":3306,\"user\":\"root\",\"password\":\"******\",\"sql-mode\":\"\",\"snapshot\":\"\",\"security\":null,\"route-rules\":[\"rule1\",\"rule2\"],\"Router\":{\"Selector\":{}},\"Conn\":null,\"session\":null},\"mysql2\":{\"host\":\"127.0.0.1\",\"port\":3306,\"user\":\"root\",\"password\":\"******\",\"sql-mode\":\"\",\"snapshot\":\"\",\"security\":null,\"route-rules\":[\"rule1\",\"rule2\"],\"Router\":{\"Selector\":{}},\"Conn\":null,\"session\":null},\"mysql3\":{\"host\":\"127.0.0.1\",\"port\":3306,\"user\":\"root\",\"password\":\"******\",\"sql-mode\":\"\",\"snapshot\":\"\",\"security\":null,\"route-rules\":[\"rule1\",\"rule3\"],\"Router\":{\"Selector\":{}},\"Conn\":null,\"session\":null},\"tidb0\":{\"host\":\"127.0.0.1\",\"port\":4000,\"user\":\"root\",\"password\":\"******\",\"sql-mode\":\"\",\"snapshot\":\"\",\"security\":null,\"route-rules\":null,\"Router\":{\"Selector\":{}},\"Conn\":null,\"session\":{\"max_execution_time\":86400,\"tidb_opt_prefer_range_scan\":\"ON\"}}},\"routes\":{\"rule1\":{\"schema-pattern\":\"test_*\",\"table-pattern\":\"t_*\",\"target-schema\":\"test\",\"target-table\":\"t\"},\"rule2\":{\"schema-pattern\":\"test2_*\",\"table-pattern\":\"t2_*\",\"target-schema\":\"test2\",\"target-table\":\"t2\"},\"rule3\":{\"schema-pattern\":\"test2_*\",\"table-pattern\":\"t2_*\",\"target-schema\":\"test\",\"target-table\":\"t\"}},\"table-configs\":{\"config1\":{\"target-tables\":[\"schema*.table*\",\"test2.t2\"],\"Schema\":\"\",\"Table\":\"\",\"ConfigIndex\":0,\"HasMatched\":false,\"IgnoreColumns\":[\"\",\"\"],\"Fields\":[\"\"],\"Range\":\"age \\u003e 10 AND age \\u003c 20\",\"TargetTableInfo\":null,\"Collation\":\"\",\"chunk-size\":0}},\"task\":{\"source-instances\":[\"mysql1\",\"mysql2\",\"mysql3\"],\"source-routes\":null,\"target-instance\":\"tidb0\",\"target-check-tables\":[\"schema*.table*\",\"!c.*\",\"test2.t2\"],\"target-configs\":[\"config1\"],\"output-dir\":\"/tmp/output/config\",\"SourceInstances\":[{\"host\":\"127.0.0.1\",\"port\":3306,\"user\":\"root\",\"password\":\"******\",\"sql-mode\":\"\",\"snapshot\":\"\",\"security\":null,\"route-rules\":[\"rule1\",\"rule2\"],\"Router\":{\"Selector\":{}},\"Conn\":null,\"session\":null},{\"host\":\"127.0.0.1\",\"port\":3306,\"user\":\"root\",\"password\":\"******\",\"sql-mode\":\"\",\"snapshot\":\"\",\"security\":null,\"route-rules\":[\"rule1\",\"rule2\"],\"Router\":{\"Selector\":{}},\"Conn\":null,\"session\":null},{\"host\":\"127.0.0.1\",\"port\":3306,\"user\":\"root\",\"password\":\"******\",\"sql-mode\":\"\",\"snapshot\":\"\",\"security\":null,\"route-rules\":[\"rule1\",\"rule3\"],\"Router\":{\"Selector\":{}},\"Conn\":null,\"session\":null}],\"TargetInstance\":{\"host\":\"127.0.0.1\",\"port\":4000,\"user\":\"root\",\"password\":\"******\",\"sql-mode\":\"\",\"snapshot\":\"\",\"security\":null,\"route-rules\":null,\"Router\":{\"Selector\":{}},\"Conn\":null,\"session\":{\"max_execution_time\":86400,\"tidb_opt_prefer_range_scan\":\"ON\"}},\"TargetTableConfigs\":[{\"target-tables\":[\"schema*.table*\",\"test2.t2\"],\"Schema\":\"\",\"Table\":\"\",\"ConfigIndex\":0,\"HasMatched\":false,\"IgnoreColumns\":[\"\",\"\"],\"Fields\":[\"\"],\"Range\":\"age \\u003e 10 AND age \\u003c 20\",\"TargetTableInfo\":null,\"Collation\":\"\",\"chunk-size\":0}],\"TargetCheckTables\":[{},{},{}],\"FixDir\":\"/tmp/output/config/fix-on-tidb0\",\"CheckpointDir\":\"/tmp/output/config/checkpoint\",\"HashFile\":\"\"},\"ConfigFile\":\"config_sharding.toml\",\"PrintVersion\":false}")
 	hash, err := cfg.Task.ComputeConfigHash()
 	require.NoError(t, err)
-	require.Equal(t, hash, "c080f9894ec24aadb4aaec1109cd1951454f09a1233f2034bc3b06e0903cb289")
+	require.Equal(t, hash, "553515158a9da2724195f6f71eb8daeebb605bc7de8f3dd9b7447468d6f385b9")
 
 	require.True(t, cfg.TableConfigs["config1"].Valid())
 
diff --git a/sync_diff_inspector/source/common/conn_test.go b/sync_diff_inspector/source/common/conn_test.go
index 143a2e5..f77d451 100644
--- a/sync_diff_inspector/source/common/conn_test.go
+++ b/sync_diff_inspector/source/common/conn_test.go
@@ -40,9 +40,9 @@ func TestConnect(t *testing.T) {
 		User:     "root",
 		Password: utils.SecretString(plainPsw),
 	}
-	_, err := ConnectMySQL(dataSource.ToDriverConfig(), 2)
+	_, err := ConnectMySQL(nil, dataSource.ToDriverConfig(), 2)
 	require.NoError(t, err)
 	dataSource.Password = utils.SecretString(base64.StdEncoding.EncodeToString([]byte(plainPsw)))
-	_, err = ConnectMySQL(dataSource.ToDriverConfig(), 2)
+	_, err = ConnectMySQL(nil, dataSource.ToDriverConfig(), 2)
 	require.NoError(t, err)
 }
diff --git a/tests/sync_diff_inspector/table_config/config.toml b/tests/sync_diff_inspector/table_config/config.toml
index 5f5bc45..75d9c80 100644
--- a/tests/sync_diff_inspector/table_config/config.toml
+++ b/tests/sync_diff_inspector/table_config/config.toml
@@ -28,6 +28,8 @@ check-struct-only = false
     port = 4000
     user = "root"
     password = ""
+    session.tidb_opt_prefer_range_scan = "ON"
+    session.max_execution_time = 86400
     # remove comment if use tidb's snapshot data
     # snapshot = "2016-10-08 16:45:26"