dumpling: fix default collation with upstream when dump database and table

dumpling/export/dump.go

@@ -8,14 +8,15 @@ import (
 	"database/sql"
 	"encoding/hex"
 	"fmt"
-	"github.com/go-sql-driver/mysql"
 	"math/big"
 	"sort"
 	"strconv"
 	"strings"
 	"sync/atomic"
 	"time"
 
+	"github.com/go-sql-driver/mysql"
+
 	// import mysql driver
 	_ "github.com/go-sql-driver/mysql"
 	"github.com/pingcap/errors"
@@ -31,6 +32,9 @@ import (
 	"github.com/pingcap/tidb/dumpling/cli"
 	tcontext "github.com/pingcap/tidb/dumpling/context"
 	"github.com/pingcap/tidb/dumpling/log"
+	"github.com/pingcap/tidb/parser"
+	"github.com/pingcap/tidb/parser/ast"
+	"github.com/pingcap/tidb/parser/format"
 	"github.com/pingcap/tidb/store/helper"
 	"github.com/pingcap/tidb/tablecodec"
 	"github.com/pingcap/tidb/util/codec"
@@ -331,12 +335,19 @@ func (d *Dumper) dumpDatabases(tctx *tcontext.Context, metaConn *sql.Conn, taskC
 		}
 	}
 
+	parser1 := parser.New()
 	for dbName, tables := range allTables {
+		var dbCollation string
 		if !conf.NoSchemas {
 			createDatabaseSQL, err := ShowCreateDatabase(metaConn, dbName)
 			if err != nil {
 				return err
 			}
+			// adjust db collation
+			createDatabaseSQL, dbCollation, err = adjustDatabaseCollation(metaConn, parser1, createDatabaseSQL, dbName)
+			if err != nil {
+				return err
+			}
 			task := NewTaskDatabaseMeta(dbName, createDatabaseSQL)
 			ctxDone := d.sendTaskToChan(tctx, task, taskChan)
 			if ctxDone {
@@ -351,7 +362,17 @@ func (d *Dumper) dumpDatabases(tctx *tcontext.Context, metaConn *sql.Conn, taskC
 			if err != nil {
 				return err
 			}
-
+			// adjust table collation
+			if table.Type == TableTypeBase {
+				newCreateSQL, defaultCollation, err := adjustTableCollation(metaConn, parser1, meta.ShowCreateTable(), dbCollation, dbName, table.Name)
+				if err != nil {
+					return err
+				}
+				if dbCollation == "" {
+					dbCollation = defaultCollation
+				}
+				meta.(*tableMeta).showCreateTable = newCreateSQL
+			}
 			if !conf.NoSchemas {
 				if table.Type == TableTypeView {
 					task := NewTaskViewMeta(dbName, table.Name, meta.ShowCreateTable(), meta.ShowCreateView())
@@ -379,6 +400,83 @@ func (d *Dumper) dumpDatabases(tctx *tcontext.Context, metaConn *sql.Conn, taskC
 	return nil
 }
 
+// adjustDatabaseCollation adjusts db collation and return new create sql and collation
+func adjustDatabaseCollation(db *sql.Conn, parser1 *parser.Parser, originSQL string, dbName string) (string, string, error) {
+	stmt, err := parser1.ParseOneStmt(originSQL, "", "")
+	if err != nil {
+		return "", "", err
+	}
+	createStmt, ok := stmt.(*ast.CreateDatabaseStmt)
+	if !ok {
+		return originSQL, "", nil
+	}
+	for _, createOption := range createStmt.Options {
+		// already have 'Collation'
+		if createOption.Tp == ast.DatabaseOptionCollate {
+			return originSQL, createOption.Value, nil
+		}
+	}
+	// get db collation
+	collation, err := GetDBCollation(db, dbName)
+	if err != nil {
+		return "", "", err
+	}
+	// add collation
+	createStmt.Options = append(createStmt.Options, &ast.DatabaseOption{Tp: ast.DatabaseOptionCollate, Value: collation})
+	// rewrite sql
+	var b []byte
+	bf := bytes.NewBuffer(b)
+	err = createStmt.Restore(&format.RestoreCtx{
+		Flags: format.DefaultRestoreFlags,
+		In:    bf,
+	})
+	if err != nil {
+		return "", "", err
+	}
+	return bf.String(), collation, nil
+}
+
+// adjustTableCollation adjusts table collation and return new create sql and default db collation
+func adjustTableCollation(db *sql.Conn, parser1 *parser.Parser, originSQL string, dbCollation string, dbName string, tableName string) (string, string, error) {
+	stmt, err := parser1.ParseOneStmt(originSQL, "", "")
+	if err != nil {
+		return "", "", err
+	}
+	createStmt, ok := stmt.(*ast.CreateTableStmt)
+	if !ok {
+		return originSQL, "", nil
+	}
+	for _, createOption := range createStmt.Options {
+		// already have 'Collation'
+		if createOption.Tp == ast.TableOptionCollate {
+			return originSQL, dbCollation, nil
+		}
+	}
+
+	collation := dbCollation
+	if collation == "" {
+		// get db collation
+		collation, err = GetDBCollation(db, dbName)
+		if err != nil {
+			return "", "", err
+		}
+	}
+
+	// add collation
+	createStmt.Options = append(createStmt.Options, &ast.TableOption{Tp: ast.TableOptionCollate, StrValue: collation})
+	// rewrite sql
+	var b []byte
+	bf := bytes.NewBuffer(b)
+	err = createStmt.Restore(&format.RestoreCtx{
+		Flags: format.DefaultRestoreFlags,
+		In:    bf,
+	})
+	if err != nil {
+		return "", "", err
+	}
+	return bf.String(), collation, nil
+}
+
 func (d *Dumper) dumpTableData(tctx *tcontext.Context, conn *sql.Conn, meta TableMeta, taskChan chan<- Task) error {
 	conf := d.conf
 	if conf.NoData {

dumpling/export/dump_test.go

@@ -15,6 +15,7 @@ import (
 
 	"github.com/pingcap/tidb/br/pkg/version"
 	tcontext "github.com/pingcap/tidb/dumpling/context"
+	"github.com/pingcap/tidb/parser"
 )
 
 func TestDumpBlock(t *testing.T) {
@@ -29,6 +30,9 @@ func TestDumpBlock(t *testing.T) {
 	mock.ExpectQuery(fmt.Sprintf("SHOW CREATE DATABASE `%s`", escapeString(database))).
 		WillReturnRows(sqlmock.NewRows([]string{"Database", "Create Database"}).
 			AddRow("test", "CREATE DATABASE `test` /*!40100 DEFAULT CHARACTER SET utf8mb4 */"))
+	mock.ExpectQuery(fmt.Sprintf("SELECT DEFAULT_COLLATION_NAME FROM INFORMATION_SCHEMA.SCHEMATA WHERE SCHEMA_NAME = '%s'", escapeString(database))).
+		WillReturnRows(sqlmock.NewRows([]string{"DEFAULT_COLLATION_NAME"}).
+			AddRow("utf8mb4_bin"))
 
 	tctx, cancel := tcontext.Background().WithLogger(appLogger).WithCancel()
 	defer cancel()
@@ -128,3 +132,101 @@ func TestGetListTableTypeByConf(t *testing.T) {
 		require.Equalf(t, x.expected, getListTableTypeByConf(conf), "server info: %s, consistency: %s", x.serverInfo, x.consistency)
 	}
 }
+
+func TestAdjustDatabaseCollation(t *testing.T) {
+	t.Parallel()
+
+	db, mock, err := sqlmock.New()
+	require.NoError(t, err)
+	defer func() {
+		require.NoError(t, db.Close())
+	}()
+
+	tctx, cancel := tcontext.Background().WithLogger(appLogger).WithCancel()
+	defer cancel()
+	conn, err := db.Conn(tctx)
+	require.NoError(t, err)
+
+	parser1 := parser.New()
+
+	originSQLs := []string{
+		"create database `test` CHARACTER SET=utf8mb4 COLLATE=utf8mb4_bin",
+		"create database `test` COLLATE=utf8mb4_bin",
+		"create database `test` CHARACTER SET=utf8mb4",
+		"create database if not exists `test`",
+	}
+
+	expectedSQLs := []string{
+		"create database `test` CHARACTER SET=utf8mb4 COLLATE=utf8mb4_bin",
+		"create database `test` COLLATE=utf8mb4_bin",
+		"CREATE DATABASE `test` CHARACTER SET = utf8mb4 COLLATE = utf8mb4_bin",
+		"CREATE DATABASE IF NOT EXISTS `test` COLLATE = utf8mb4_bin",
+	}
+
+	for i, originSQL := range originSQLs {
+		if i > 1 {
+			mock.ExpectQuery("SELECT DEFAULT_COLLATION_NAME FROM INFORMATION_SCHEMA.SCHEMATA WHERE SCHEMA_NAME = 'test'").
+				WillReturnRows(sqlmock.NewRows([]string{"DEFAULT_COLLATION_NAME"}).
+					AddRow("utf8mb4_bin"))
+		}
+		newSQL, collation, err := adjustDatabaseCollation(conn, parser1, originSQL, "test")
+		require.NoError(t, err)
+		require.Equal(t, expectedSQLs[i], newSQL)
+		require.Equal(t, "utf8mb4_bin", collation)
+	}
+}
+
+func TestAdjustTableCollation(t *testing.T) {
+	t.Parallel()
+
+	db, mock, err := sqlmock.New()
+	require.NoError(t, err)
+	defer func() {
+		require.NoError(t, db.Close())
+	}()
+
+	tctx, cancel := tcontext.Background().WithLogger(appLogger).WithCancel()
+	defer cancel()
+	conn, err := db.Conn(tctx)
+	require.NoError(t, err)
+
+	parser1 := parser.New()
+
+	originSQLs := []string{
+		"create table `test`.`t1` (id int) CHARSET=utf8mb4 COLLATE=utf8mb4_bin",
+		"create table `test`.`t1` (id int) COLLATE=utf8mb4_bin",
+		"create table `test`.`t1` (id int) CHARSET=utf8mb4",
+		"create table `test`.`t1` (id int)",
+	}
+
+	expectedSQLs := []string{
+		"create table `test`.`t1` (id int) CHARSET=utf8mb4 COLLATE=utf8mb4_bin",
+		"create table `test`.`t1` (id int) COLLATE=utf8mb4_bin",
+		"CREATE TABLE `test`.`t1` (`id` INT) DEFAULT CHARACTER SET = UTF8MB4 DEFAULT COLLATE = UTF8MB4_BIN",
+		"CREATE TABLE `test`.`t1` (`id` INT) DEFAULT COLLATE = UTF8MB4_BIN",
+	}
+
+	for i, originSQL := range originSQLs {
+		newSQL, collation, err := adjustTableCollation(conn, parser1, originSQL, "utf8mb4_bin", "test", "t1")
+		require.NoError(t, err)
+		require.Equal(t, expectedSQLs[i], newSQL)
+		require.Equal(t, "utf8mb4_bin", collation)
+	}
+
+	for i, originSQL := range originSQLs {
+		if i > 1 {
+			mock.ExpectQuery("SELECT DEFAULT_COLLATION_NAME FROM INFORMATION_SCHEMA.SCHEMATA WHERE SCHEMA_NAME = 'test'").
+				WillReturnRows(sqlmock.NewRows([]string{"DEFAULT_COLLATION_NAME"}).
+					AddRow("utf8mb4_bin"))
+		}
+		newSQL, collation, err := adjustTableCollation(conn, parser1, originSQL, "", "test", "t1")
+		require.NoError(t, err)
+		require.Equal(t, expectedSQLs[i], newSQL)
+		if i > 1 {
+			require.Equal(t, "utf8mb4_bin", collation)
+		} else {
+			require.Equal(t, "", collation)
+		}
+
+	}
+}

dumpling/export/sql.go

@@ -1387,3 +1387,13 @@ func GetRegionInfos(db *sql.Conn) (*helper.RegionsInfo, error) {
 	})
 	return regionsInfo, err
 }
+
+// GetDBCollation gets db collation from INFORMATION_SCHEMA.SCHEMATA
+func GetDBCollation(db *sql.Conn, database string) (string, error) {
+	query := fmt.Sprintf("SELECT DEFAULT_COLLATION_NAME FROM INFORMATION_SCHEMA.SCHEMATA WHERE SCHEMA_NAME = '%s'", database)
+	var collation string
+	err := simpleQuery(db, query, func(rows *sql.Rows) error {
+		return rows.Scan(&collation)
+	})
+	return collation, err
+}