br: modify collate.newCollationEnabled according to the config of the cluster (#39173)

br/pkg/gluetidb/glue.go

@@ -298,3 +298,116 @@ func (gs *tidbSession) showCreateDatabase(db *model.DBInfo) (string, error) {
 func (gs *tidbSession) showCreatePlacementPolicy(policy *model.PolicyInfo) string {
 	return executor.ConstructResultOfShowCreatePlacementPolicy(policy)
 }
+
+// mockSession is used for test.
+type mockSession struct {
+	se         session.Session
+	globalVars map[string]string
+}
+
+// GetSessionCtx implements glue.Glue
+func (s *mockSession) GetSessionCtx() sessionctx.Context {
+	return s.se
+}
+
+// Execute implements glue.Session.
+func (s *mockSession) Execute(ctx context.Context, sql string) error {
+	return s.ExecuteInternal(ctx, sql)
+}
+
+func (s *mockSession) ExecuteInternal(ctx context.Context, sql string, args ...interface{}) error {
+	return nil
+}
+
+// CreateDatabase implements glue.Session.
+func (s *mockSession) CreateDatabase(ctx context.Context, schema *model.DBInfo) error {
+	log.Fatal("unimplemented CreateDatabase for mock session")
+	return nil
+}
+
+// CreatePlacementPolicy implements glue.Session.
+func (s *mockSession) CreatePlacementPolicy(ctx context.Context, policy *model.PolicyInfo) error {
+	log.Fatal("unimplemented CreateDatabase for mock session")
+	return nil
+}
+
+// CreateTables implements glue.BatchCreateTableSession.
+func (s *mockSession) CreateTables(ctx context.Context, tables map[string][]*model.TableInfo) error {
+	log.Fatal("unimplemented CreateDatabase for mock session")
+	return nil
+}
+
+// CreateTable implements glue.Session.
+func (s *mockSession) CreateTable(ctx context.Context, dbName model.CIStr, table *model.TableInfo) error {
+	log.Fatal("unimplemented CreateDatabase for mock session")
+	return nil
+}
+
+// Close implements glue.Session.
+func (s *mockSession) Close() {
+	s.se.Close()
+}
+
+// GetGlobalVariables implements glue.Session.
+func (s *mockSession) GetGlobalVariable(name string) (string, error) {
+	if ret, ok := s.globalVars[name]; ok {
+		return ret, nil
+	}
+	return "True", nil
+}
+
+// MockGlue only used for test
+type MockGlue struct {
+	se         session.Session
+	GlobalVars map[string]string
+}
+
+func (m *MockGlue) SetSession(se session.Session) {
+	m.se = se
+}
+
+// GetDomain implements glue.Glue.
+func (*MockGlue) GetDomain(store kv.Storage) (*domain.Domain, error) {
+	return nil, nil
+}
+
+// CreateSession implements glue.Glue.
+func (m *MockGlue) CreateSession(store kv.Storage) (glue.Session, error) {
+	glueSession := &mockSession{
+		se:         m.se,
+		globalVars: m.GlobalVars,
+	}
+	return glueSession, nil
+}
+
+// Open implements glue.Glue.
+func (*MockGlue) Open(path string, option pd.SecurityOption) (kv.Storage, error) {
+	return nil, nil
+}
+
+// OwnsStorage implements glue.Glue.
+func (*MockGlue) OwnsStorage() bool {
+	return true
+}
+
+// StartProgress implements glue.Glue.
+func (*MockGlue) StartProgress(ctx context.Context, cmdName string, total int64, redirectLog bool) glue.Progress {
+	return nil
+}
+
+// Record implements glue.Glue.
+func (*MockGlue) Record(name string, value uint64) {
+}
+
+// GetVersion implements glue.Glue.
+func (*MockGlue) GetVersion() string {
+	return "mock glue"
+}
+
+// UseOneShotSession implements glue.Glue.
+func (m *MockGlue) UseOneShotSession(store kv.Storage, closeDomain bool, fn func(glue.Session) error) error {
+	glueSession := &mockSession{
+		se: m.se,
+	}
+	return fn(glueSession)
+}

br/pkg/restore/client.go

@@ -46,6 +46,7 @@ import (
 	"github.com/pingcap/tidb/store/pdtypes"
 	"github.com/pingcap/tidb/tablecodec"
 	"github.com/pingcap/tidb/util/codec"
+	"github.com/pingcap/tidb/util/collate"
 	"github.com/pingcap/tidb/util/mathutil"
 	filter "github.com/pingcap/tidb/util/table-filter"
 	"github.com/tikv/client-go/v2/oracle"
@@ -2016,3 +2017,47 @@ func (rc *Client) SaveSchemas(
 	}
 	return nil
 }
+
+func CheckNewCollationEnable(
+	backupNewCollationEnable string,
+	g glue.Glue,
+	storage kv.Storage,
+	CheckRequirements bool,
+) error {
+	if backupNewCollationEnable == "" {
+		if CheckRequirements {
+			return errors.Annotatef(berrors.ErrUnknown,
+				"the config 'new_collations_enabled_on_first_bootstrap' not found in backupmeta. "+
+					"you can use \"show config WHERE name='new_collations_enabled_on_first_bootstrap';\" to manually check the config. "+
+					"if you ensure the config 'new_collations_enabled_on_first_bootstrap' in backup cluster is as same as restore cluster, "+
+					"use --check-requirements=false to skip this check")
+		}
+		log.Warn("the config 'new_collations_enabled_on_first_bootstrap' is not in backupmeta")
+		return nil
+	}
+
+	se, err := g.CreateSession(storage)
+	if err != nil {
+		return errors.Trace(err)
+	}
+
+	newCollationEnable, err := se.GetGlobalVariable(utils.GetTidbNewCollationEnabled())
+	if err != nil {
+		return errors.Trace(err)
+	}
+
+	if !strings.EqualFold(backupNewCollationEnable, newCollationEnable) {
+		return errors.Annotatef(berrors.ErrUnknown,
+			"the config 'new_collations_enabled_on_first_bootstrap' not match, upstream:%v, downstream: %v",
+			backupNewCollationEnable, newCollationEnable)
+	}
+
+	// collate.newCollationEnabled is set to 1 when the collate package is initialized,
+	// so we need to modify this value according to the config of the cluster
+	// before using the collate package.
+	enabled := newCollationEnable == "True"
+	// modify collate.newCollationEnabled according to the config of the cluster
+	collate.SetNewCollationEnabledForTest(enabled)
+	log.Info("set new_collation_enabled", zap.Bool("new_collation_enabled", enabled))
+	return nil
+}

br/pkg/restore/client_test.go

@@ -9,6 +9,7 @@ import (
 	"testing"
 	"time"
 
+	backuppb "github.com/pingcap/kvproto/pkg/brpb"
 	"github.com/pingcap/kvproto/pkg/metapb"
 	"github.com/pingcap/tidb/br/pkg/gluetidb"
 	"github.com/pingcap/tidb/br/pkg/metautil"
@@ -240,3 +241,75 @@ func TestPreCheckTableTiFlashReplicas(t *testing.T) {
 		require.Nil(t, tables[i].Info.TiFlashReplica)
 	}
 }
+
+func TestCheckNewCollationEnable(t *testing.T) {
+	caseList := []struct {
+		backupMeta                  *backuppb.BackupMeta
+		newCollationEnableInCluster string
+		CheckRequirements           bool
+		isErr                       bool
+	}{
+		{
+			backupMeta:                  &backuppb.BackupMeta{NewCollationsEnabled: "True"},
+			newCollationEnableInCluster: "True",
+			CheckRequirements:           true,
+			isErr:                       false,
+		},
+		{
+			backupMeta:                  &backuppb.BackupMeta{NewCollationsEnabled: "True"},
+			newCollationEnableInCluster: "False",
+			CheckRequirements:           true,
+			isErr:                       true,
+		},
+		{
+			backupMeta:                  &backuppb.BackupMeta{NewCollationsEnabled: "False"},
+			newCollationEnableInCluster: "True",
+			CheckRequirements:           true,
+			isErr:                       true,
+		},
+		{
+			backupMeta:                  &backuppb.BackupMeta{NewCollationsEnabled: "False"},
+			newCollationEnableInCluster: "false",
+			CheckRequirements:           true,
+			isErr:                       false,
+		},
+		{
+			backupMeta:                  &backuppb.BackupMeta{NewCollationsEnabled: "False"},
+			newCollationEnableInCluster: "True",
+			CheckRequirements:           false,
+			isErr:                       true,
+		},
+		{
+			backupMeta:                  &backuppb.BackupMeta{NewCollationsEnabled: "True"},
+			newCollationEnableInCluster: "False",
+			CheckRequirements:           false,
+			isErr:                       true,
+		},
+		{
+			backupMeta:                  &backuppb.BackupMeta{NewCollationsEnabled: ""},
+			newCollationEnableInCluster: "True",
+			CheckRequirements:           false,
+			isErr:                       false,
+		},
+		{
+			backupMeta:                  &backuppb.BackupMeta{NewCollationsEnabled: ""},
+			newCollationEnableInCluster: "True",
+			CheckRequirements:           true,
+			isErr:                       true,
+		},
+	}
+
+	for i, ca := range caseList {
+		g := &gluetidb.MockGlue{
+			GlobalVars: map[string]string{"new_collation_enabled": ca.newCollationEnableInCluster},
+		}
+		err := restore.CheckNewCollationEnable(ca.backupMeta.GetNewCollationsEnabled(), g, nil, ca.CheckRequirements)
+
+		t.Logf("[%d] Got Error: %v\n", i, err)
+		if ca.isErr {
+			require.Error(t, err)
+		} else {
+			require.NoError(t, err)
+		}
+	}
+}

br/pkg/task/backup.go

@@ -264,12 +264,12 @@ func RunBackup(c context.Context, g glue.Glue, cmdName string, cfg *BackupConfig
 
 	var newCollationEnable string
 	err = g.UseOneShotSession(mgr.GetStorage(), !needDomain, func(se glue.Session) error {
-		newCollationEnable, err = se.GetGlobalVariable(tidbNewCollationEnabled)
+		newCollationEnable, err = se.GetGlobalVariable(utils.GetTidbNewCollationEnabled())
 		if err != nil {
 			return errors.Trace(err)
 		}
 		log.Info("get new_collations_enabled_on_first_bootstrap config from system table",
-			zap.String(tidbNewCollationEnabled, newCollationEnable))
+			zap.String(utils.GetTidbNewCollationEnabled(), newCollationEnable))
 		return nil
 	})
 	if err != nil {

br/pkg/task/common.go

@@ -87,8 +87,6 @@ const (
 	crypterAES128KeyLen = 16
 	crypterAES192KeyLen = 24
 	crypterAES256KeyLen = 32
-
-	tidbNewCollationEnabled = "new_collation_enabled"
 )
 
 // TLSConfig is the common configuration for TLS connection.

br/pkg/task/restore.go

@@ -4,7 +4,6 @@ package task
 
 import (
 	"context"
-	"strings"
 	"time"
 
 	"github.com/opentracing/opentracing-go"
@@ -23,7 +22,6 @@ import (
 	"github.com/pingcap/tidb/br/pkg/utils"
 	"github.com/pingcap/tidb/br/pkg/version"
 	"github.com/pingcap/tidb/config"
-	"github.com/pingcap/tidb/kv"
 	"github.com/pingcap/tidb/util/mathutil"
 	"github.com/spf13/cobra"
 	"github.com/spf13/pflag"
@@ -330,43 +328,6 @@ func CheckRestoreDBAndTable(client *restore.Client, cfg *RestoreConfig) error {
 	return nil
 }
 
-func CheckNewCollationEnable(
-	backupNewCollationEnable string,
-	g glue.Glue,
-	storage kv.Storage,
-	CheckRequirements bool,
-) error {
-	if backupNewCollationEnable == "" {
-		if CheckRequirements {
-			return errors.Annotatef(berrors.ErrUnknown,
-				"the config 'new_collations_enabled_on_first_bootstrap' not found in backupmeta. "+
-					"you can use \"show config WHERE name='new_collations_enabled_on_first_bootstrap';\" to manually check the config. "+
-					"if you ensure the config 'new_collations_enabled_on_first_bootstrap' in backup cluster is as same as restore cluster, "+
-					"use --check-requirements=false to skip this check")
-		} else {
-			log.Warn("the config 'new_collations_enabled_on_first_bootstrap' is not in backupmeta")
-			return nil
-		}
-	}
-
-	se, err := g.CreateSession(storage)
-	if err != nil {
-		return errors.Trace(err)
-	}
-
-	newCollationEnable, err := se.GetGlobalVariable(tidbNewCollationEnabled)
-	if err != nil {
-		return errors.Trace(err)
-	}
-
-	if !strings.EqualFold(backupNewCollationEnable, newCollationEnable) {
-		return errors.Annotatef(berrors.ErrUnknown,
-			"the config 'new_collations_enabled_on_first_bootstrap' not match, upstream:%v, downstream: %v",
-			backupNewCollationEnable, newCollationEnable)
-	}
-	return nil
-}
-
 func isFullRestore(cmdName string) bool {
 	return cmdName == FullRestoreCmd
 }
@@ -439,7 +400,7 @@ func RunRestore(c context.Context, g glue.Glue, cmdName string, cfg *RestoreConf
 			return errors.Trace(versionErr)
 		}
 	}
-	if err = CheckNewCollationEnable(backupMeta.GetNewCollationsEnabled(), g, mgr.GetStorage(), cfg.CheckRequirements); err != nil {
+	if err = restore.CheckNewCollationEnable(backupMeta.GetNewCollationsEnabled(), g, mgr.GetStorage(), cfg.CheckRequirements); err != nil {
 		return errors.Trace(err)
 	}
 

br/pkg/utils/db.go

@@ -7,6 +7,10 @@ import (
 	"database/sql"
 )
 
+const (
+	tidbNewCollationEnabled = "new_collation_enabled"
+)
+
 var (
 	// check sql.DB and sql.Conn implement QueryExecutor and DBExecutor
 	_ DBExecutor = &sql.DB{}
@@ -30,3 +34,7 @@ type DBExecutor interface {
 	StmtExecutor
 	BeginTx(ctx context.Context, opts *sql.TxOptions) (*sql.Tx, error)
 }
+
+func GetTidbNewCollationEnabled() string {
+	return tidbNewCollationEnabled
+}