br: modify collate.newCollationEnabled according to the config of the…

… cluster (#39173)

close #39150

br/pkg/gluetidb/glue.go

@@ -303,7 +303,8 @@ func (gs *tidbSession) showCreatePlacementPolicy(policy *model.PolicyInfo) strin
 
 // mockSession is used for test.
 type mockSession struct {
-	se session.Session
+	se         session.Session
+	globalVars map[string]string
 }
 
 // GetSessionCtx implements glue.Glue
@@ -368,12 +369,16 @@ func (s *mockSession) Close() {
 
 // GetGlobalVariables implements glue.Session.
 func (s *mockSession) GetGlobalVariable(name string) (string, error) {
-	return "true", nil
+	if ret, ok := s.globalVars[name]; ok {
+		return ret, nil
+	}
+	return "True", nil
 }
 
 // MockGlue only used for test
 type MockGlue struct {
-	se session.Session
+	se         session.Session
+	GlobalVars map[string]string
 }
 
 func (m *MockGlue) SetSession(se session.Session) {
@@ -388,7 +393,8 @@ func (*MockGlue) GetDomain(store kv.Storage) (*domain.Domain, error) {
 // CreateSession implements glue.Glue.
 func (m *MockGlue) CreateSession(store kv.Storage) (glue.Session, error) {
 	glueSession := &mockSession{
-		se: m.se,
+		se:         m.se,
+		globalVars: m.GlobalVars,
 	}
 	return glueSession, nil
 }

br/pkg/restore/BUILD.bazel

@@ -58,6 +58,7 @@ go_library(
         "//tablecodec",
         "//util",
         "//util/codec",
+        "//util/collate",
         "//util/hack",
         "//util/mathutil",
         "//util/table-filter",

br/pkg/restore/client.go

@@ -51,6 +51,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"
@@ -2644,3 +2645,47 @@ func TidyOldSchemas(sr *stream.SchemasReplace) *backup.Schemas {
 	}
 	return schemas
 }
+
+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

@@ -1488,3 +1488,75 @@ func TestApplyKVFilesWithBatchMethod4(t *testing.T) {
 		},
 	)
 }
+
+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

@@ -308,12 +308,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

@@ -96,8 +96,6 @@ const (
 	crypterAES192KeyLen = 24
 	crypterAES256KeyLen = 32
 
-	tidbNewCollationEnabled = "new_collation_enabled"
-
 	flagFullBackupType = "type"
 )
 

br/pkg/task/restore.go

@@ -26,7 +26,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"
 	"github.com/pingcap/tidb/util/mathutil"
 	"github.com/spf13/cobra"
@@ -441,42 +440,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")
-		}
-		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
 }
@@ -548,7 +511,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

@@ -15,6 +15,10 @@ import (
 	"go.uber.org/zap"
 )
 
+const (
+	tidbNewCollationEnabled = "new_collation_enabled"
+)
+
 var (
 	// check sql.DB and sql.Conn implement QueryExecutor and DBExecutor
 	_ DBExecutor = &sql.DB{}
@@ -117,3 +121,7 @@ func CheckLogBackupTaskExist() bool {
 func IsLogBackupInUse(ctx sessionctx.Context) bool {
 	return CheckLogBackupEnabled(ctx) && CheckLogBackupTaskExist()
 }
+
+func GetTidbNewCollationEnabled() string {
+	return tidbNewCollationEnabled
+}