support session parameters (pingcap#167)

* support session parameters * delete tidbmemquota for non-TiDB databases
tisonkun · Oct 20, 2020 · 5f3b14a · 5f3b14a
1 parent 5306add
commit 5f3b14a
Show file tree

Hide file tree

Showing 4 changed files with 29 additions and 16 deletions.
diff --git a/dumpling/cmd/dumpling/main.go b/dumpling/cmd/dumpling/main.go
@@ -68,6 +68,7 @@ var (
 	keyPath                 string
 	csvSeparator            string
 	csvDelimiter            string
+	sessionParams           map[string]string
 
 	completeInsert       bool
 	dumpEmptyDatabase    bool
@@ -127,6 +128,7 @@ func main() {
 	pflag.StringVar(&csvDelimiter, "csv-delimiter", "\"", "The delimiter for values in csv files, default '\"'")
 	pflag.StringVar(&outputFilenameFormat, "output-filename-template", "", "The output filename template (without file extension)")
 	pflag.BoolVar(&completeInsert, "complete-insert", false, "Use complete INSERT statements that include column names")
+	pflag.StringToStringVar(&sessionParams, "params", nil, `Extra session variables used while dumping, accepted format: --params "character_set_client=latin1,character_set_connection=latin1"`)
 
 	storage.DefineFlags(pflag.CommandLine)
 
@@ -177,6 +179,9 @@ func main() {
 	}
 
 	conf := export.DefaultConfig()
+	for k, v := range sessionParams {
+		conf.SessionParams[k] = v
+	}
 	conf.Databases = databases
 	conf.Host = host
 	conf.User = user
@@ -208,7 +213,7 @@ func main() {
 	conf.Security.CAPath = caPath
 	conf.Security.CertPath = certPath
 	conf.Security.KeyPath = keyPath
-	conf.SessionParams["tidb_mem_quota_query"] = tidbMemQuotaQuery
+	conf.SessionParams[export.TiDBMemQuotaQueryName] = tidbMemQuotaQuery
 	conf.CsvSeparator = csvSeparator
 	conf.CsvDelimiter = csvDelimiter
 	conf.OutputFileTemplate = tmpl

diff --git a/dumpling/v4/export/config.go b/dumpling/v4/export/config.go
@@ -137,9 +137,11 @@ func (config *Config) createExternalStorage(ctx context.Context) (storage.Extern
 }
 
 const (
-	UnspecifiedSize            = 0
-	DefaultTiDBMemQuotaQuery   = 32 * (1 << 30)
-	DefaultStatementSize       = 1000000
+	UnspecifiedSize          = 0
+	DefaultTiDBMemQuotaQuery = 32 * (1 << 30)
+	DefaultStatementSize     = 1000000
+	TiDBMemQuotaQueryName    = "tidb_mem_quota_query"
+
 	defaultDumpThreads         = 128
 	defaultDumpGCSafePointTTL  = 5 * 60
 	dumplingServiceSafePointID = "dumpling"

diff --git a/dumpling/v4/export/dump.go b/dumpling/v4/export/dump.go
@@ -47,23 +47,27 @@ func Dump(pCtx context.Context, conf *Config) (err error) {
 
 	var doPdGC bool
 	var pdClient pd.Client
-	if conf.ServerInfo.ServerType == ServerTypeTiDB && conf.ServerInfo.ServerVersion.Compare(*gcSafePointVersion) >= 0 {
-		pdAddrs, err := GetPdAddrs(pool)
-		if err != nil {
-			return err
-		}
-		if len(pdAddrs) > 0 {
-			doPdGC, err = checkSameCluster(ctx, pool, pdAddrs)
+	if conf.ServerInfo.ServerType == ServerTypeTiDB {
+		if conf.ServerInfo.ServerVersion.Compare(*gcSafePointVersion) >= 0 {
+			pdAddrs, err := GetPdAddrs(pool)
 			if err != nil {
-				log.Warn("meet error while check whether fetched pd addr and TiDB belongs to one cluster", zap.Error(err), zap.Strings("pdAddrs", pdAddrs))
-			} else if doPdGC {
-				pdClient, err = pd.NewClientWithContext(ctx, pdAddrs, pd.SecurityOption{})
+				return err
+			}
+			if len(pdAddrs) > 0 {
+				doPdGC, err = checkSameCluster(ctx, pool, pdAddrs)
 				if err != nil {
-					log.Warn("create pd client to control GC failed", zap.Error(err), zap.Strings("pdAddrs", pdAddrs))
-					doPdGC = false
+					log.Warn("meet error while check whether fetched pd addr and TiDB belongs to one cluster", zap.Error(err), zap.Strings("pdAddrs", pdAddrs))
+				} else if doPdGC {
+					pdClient, err = pd.NewClientWithContext(ctx, pdAddrs, pd.SecurityOption{})
+					if err != nil {
+						log.Warn("create pd client to control GC failed", zap.Error(err), zap.Strings("pdAddrs", pdAddrs))
+						doPdGC = false
+					}
 				}
 			}
 		}
+	} else {
+		delete(conf.SessionParams, TiDBMemQuotaQueryName)
 	}
 
 	if conf.Snapshot == "" && (doPdGC || conf.Consistency == "snapshot") {

diff --git a/dumpling/v4/export/sql.go b/dumpling/v4/export/sql.go
@@ -482,6 +482,8 @@ func resetDBWithSessionParams(db *sql.DB, dsn string, params map[string]interfac
 
 	for k, v := range support {
 		var s string
+		// Wrap string with quote to handle string with space. For example, '2020-10-20 13:41:40'
+		// For --params argument, quote doesn't matter because it doesn't affect the actual value
 		if str, ok := v.(string); ok {
 			s = wrapStringWith(str, "'")
 		} else {