br: reset cloud_admin and root after ebs restoration (#39986)

br/pkg/restore/BUILD.bazel

@@ -115,6 +115,7 @@ go_test(
         "search_test.go",
         "split_test.go",
         "stream_metas_test.go",
+        "systable_restore_test.go",
         "util_test.go",
     ],
     embed = [":restore"],

br/pkg/restore/systable_restore.go

@@ -19,7 +19,11 @@ import (
 	"go.uber.org/zap"
 )
 
-const sysUserTableName = "user"
+const (
+	rootUser         = "root"
+	sysUserTableName = "user"
+	cloudAdminUser   = "cloud_admin"
+)
 
 var statsTables = map[string]struct{}{
 	"stats_buckets":    {},
@@ -51,14 +55,14 @@ var unRecoverableTable = map[string]struct{}{
 // skip clearing or restoring 'cloud_admin'@'%' which is a special
 // user on TiDB Cloud
 var sysPrivilegeTableMap = map[string]string{
-	"user":          "not (user = 'cloud_admin' and host = '%')",       // since v1.0.0
-	"db":            "not (user = 'cloud_admin' and host = '%')",       // since v1.0.0
-	"tables_priv":   "not (user = 'cloud_admin' and host = '%')",       // since v1.0.0
-	"columns_priv":  "not (user = 'cloud_admin' and host = '%')",       // since v1.0.0
-	"default_roles": "not (user = 'cloud_admin' and host = '%')",       // since v3.0.0
-	"role_edges":    "not (to_user = 'cloud_admin' and to_host = '%')", // since v3.0.0
-	"global_priv":   "not (user = 'cloud_admin' and host = '%')",       // since v3.0.8
-	"global_grants": "not (user = 'cloud_admin' and host = '%')",       // since v5.0.3
+	"user":          "(user = '%s' and host = '%%')",       // since v1.0.0
+	"db":            "(user = '%s' and host = '%%')",       // since v1.0.0
+	"tables_priv":   "(user = '%s' and host = '%%')",       // since v1.0.0
+	"columns_priv":  "(user = '%s' and host = '%%')",       // since v1.0.0
+	"default_roles": "(user = '%s' and host = '%%')",       // since v3.0.0
+	"role_edges":    "(to_user = '%s' and to_host = '%%')", // since v3.0.0
+	"global_priv":   "(user = '%s' and host = '%%')",       // since v3.0.8
+	"global_grants": "(user = '%s' and host = '%%')",       // since v5.0.3
 }
 
 func isUnrecoverableTable(tableName string) bool {
@@ -71,6 +75,78 @@ func isStatsTable(tableName string) bool {
 	return ok
 }
 
+func generateResetSQLs(db *database, resetUsers []string) []string {
+	if db.Name.L != mysql.SystemDB {
+		return nil
+	}
+	sqls := make([]string, 0, 10)
+	// we only need reset root password once
+	rootReset := false
+	for tableName := range db.ExistingTables {
+		if sysPrivilegeTableMap[tableName] != "" {
+			for _, name := range resetUsers {
+				if strings.ToLower(name) == rootUser {
+					if !rootReset {
+						updateSQL := fmt.Sprintf("UPDATE %s.%s SET authentication_string='',"+
+							" Shutdown_priv='Y',"+
+							" Config_priv='Y'"+
+							" WHERE USER='root' AND Host='%%';",
+							db.Name.L, sysUserTableName)
+						sqls = append(sqls, updateSQL)
+						rootReset = true
+					} else {
+						continue
+					}
+				} else {
+					/* #nosec G202: SQL string concatenation */
+					whereClause := fmt.Sprintf("WHERE "+sysPrivilegeTableMap[tableName], name)
+					deleteSQL := fmt.Sprintf("DELETE FROM %s %s;",
+						utils.EncloseDBAndTable(db.Name.L, tableName), whereClause)
+					sqls = append(sqls, deleteSQL)
+				}
+			}
+		}
+	}
+	return sqls
+}
+
+// ClearSystemUsers is used for volume-snapshot restoration.
+// because we can not support restore user in some scenarios, for example in cloud.
+// we'd better use this function to drop cloud_admin user after volume-snapshot restore.
+func (rc *Client) ClearSystemUsers(ctx context.Context, resetUsers []string) error {
+	sysDB := mysql.SystemDB
+	db, ok := rc.getDatabaseByName(sysDB)
+	if !ok {
+		log.Warn("target database not exist, aborting", zap.String("database", sysDB))
+		return nil
+	}
+	execSQL := func(sql string) error {
+		// SQLs here only contain table name and database name, seems it is no need to redact them.
+		if err := rc.db.se.Execute(ctx, sql); err != nil {
+			log.Warn("failed to clear system users",
+				zap.Stringer("database", db.Name),
+				zap.String("sql", sql),
+				zap.Error(err),
+			)
+			return berrors.ErrUnknown.Wrap(err).GenWithStack("failed to execute %s", sql)
+		}
+		log.Info("successfully clear system users after restoration",
+			zap.Stringer("database", db.Name),
+			zap.String("sql", sql),
+		)
+		return nil
+	}
+
+	sqls := generateResetSQLs(db, resetUsers)
+	for _, sql := range sqls {
+		log.Info("reset system user for cloud", zap.String("sql", sql))
+		if err := execSQL(sql); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
 // RestoreSystemSchemas restores the system schema(i.e. the `mysql` schema).
 // Detail see https://github.com/pingcap/br/issues/679#issuecomment-762592254.
 func (rc *Client) RestoreSystemSchemas(ctx context.Context, f filter.Filter) {
@@ -203,14 +279,15 @@ func (rc *Client) replaceTemporaryTableToSystable(ctx context.Context, ti *model
 	}
 
 	if db.ExistingTables[tableName] != nil {
-		whereClause := ""
+		whereNotClause := ""
 		if rc.fullClusterRestore && sysPrivilegeTableMap[tableName] != "" {
 			// cloud_admin is a special user on tidb cloud, need to skip it.
-			whereClause = fmt.Sprintf("WHERE %s", sysPrivilegeTableMap[tableName])
+			/* #nosec G202: SQL string concatenation */
+			whereNotClause = fmt.Sprintf("WHERE NOT "+sysPrivilegeTableMap[tableName], cloudAdminUser)
 			log.Info("full cluster restore, delete existing data",
 				zap.String("table", tableName), zap.Stringer("schema", db.Name))
 			deleteSQL := fmt.Sprintf("DELETE FROM %s %s;",
-				utils.EncloseDBAndTable(db.Name.L, tableName), whereClause)
+				utils.EncloseDBAndTable(db.Name.L, tableName), whereNotClause)
 			if err := execSQL(deleteSQL); err != nil {
 				return err
 			}
@@ -228,7 +305,7 @@ func (rc *Client) replaceTemporaryTableToSystable(ctx context.Context, ti *model
 			utils.EncloseDBAndTable(db.Name.L, tableName),
 			colListStr, colListStr,
 			utils.EncloseDBAndTable(db.TemporaryName.L, tableName),
-			whereClause)
+			whereNotClause)
 		return execSQL(replaceIntoSQL)
 	}
 

br/pkg/restore/systable_restore_test.go

@@ -0,0 +1,72 @@
+// Copyright 2020 PingCAP, Inc. Licensed under Apache-2.0.
+
+package restore
+
+import (
+	"regexp"
+	"testing"
+
+	"github.com/pingcap/tidb/br/pkg/utils"
+	"github.com/pingcap/tidb/parser/model"
+	"github.com/stretchr/testify/require"
+)
+
+func testTableInfo(name string) *model.TableInfo {
+	return &model.TableInfo{
+		Name: model.NewCIStr(name),
+	}
+}
+
+func TestGenerateResetSQL(t *testing.T) {
+	// case #1: ignore non-mysql databases
+	mockDB := &database{
+		ExistingTables: map[string]*model.TableInfo{},
+		Name:           model.NewCIStr("non-mysql"),
+		TemporaryName:  utils.TemporaryDBName("non-mysql"),
+	}
+	for name := range sysPrivilegeTableMap {
+		mockDB.ExistingTables[name] = testTableInfo(name)
+	}
+	resetUsers := []string{"cloud_admin", "root"}
+	require.Equal(t, 0, len(generateResetSQLs(mockDB, resetUsers)))
+
+	// case #2: ignore non expected table
+	mockDB = &database{
+		ExistingTables: map[string]*model.TableInfo{},
+		Name:           model.NewCIStr("mysql"),
+		TemporaryName:  utils.TemporaryDBName("mysql"),
+	}
+	for name := range sysPrivilegeTableMap {
+		name += "non_available"
+		mockDB.ExistingTables[name] = testTableInfo(name)
+	}
+	resetUsers = []string{"cloud_admin", "root"}
+	require.Equal(t, 0, len(generateResetSQLs(mockDB, resetUsers)))
+
+	// case #3: only reset cloud admin account
+	for name := range sysPrivilegeTableMap {
+		mockDB.ExistingTables[name] = testTableInfo(name)
+	}
+	resetUsers = []string{"cloud_admin"}
+	sqls := generateResetSQLs(mockDB, resetUsers)
+	require.Equal(t, 8, len(sqls))
+	for _, sql := range sqls {
+		// for cloud_admin we only generate DELETE sql
+		require.Regexp(t, regexp.MustCompile("DELETE*"), sql)
+	}
+
+	// case #4: reset cloud admin/other account
+	resetUsers = []string{"cloud_admin", "cloud_other"}
+	sqls = generateResetSQLs(mockDB, resetUsers)
+	require.Equal(t, 16, len(sqls))
+	for _, sql := range sqls {
+		// for cloud_admin/cloud_other we only generate DELETE sql
+		require.Regexp(t, regexp.MustCompile("DELETE*"), sql)
+	}
+
+	// case #5: reset cloud admin && root account
+	resetUsers = []string{"cloud_admin", "root"}
+	sqls = generateResetSQLs(mockDB, resetUsers)
+	// 8 DELETE sqls for cloud admin and 1 UPDATE sql for root
+	require.Equal(t, 9, len(sqls))
+}

br/pkg/task/restore.go

@@ -62,6 +62,8 @@ const (
 	FlagPiTRBatchSize   = "pitr-batch-size"
 	FlagPiTRConcurrency = "pitr-concurrency"
 
+	FlagResetSysUsers = "reset-sys-users"
+
 	defaultPiTRBatchCount     = 8
 	defaultPiTRBatchSize      = 16 * 1024 * 1024
 	defaultRestoreConcurrency = 128
@@ -93,6 +95,8 @@ type RestoreCommonConfig struct {
 
 	// determines whether enable restore sys table on default, see fullClusterRestore in restore/client.go
 	WithSysTable bool `json:"with-sys-table" toml:"with-sys-table"`
+
+	ResetSysUsers []string `json:"reset-sys-users" toml:"reset-sys-users"`
 }
 
 // adjust adjusts the abnormal config value in the current config.
@@ -118,10 +122,12 @@ func DefineRestoreCommonFlags(flags *pflag.FlagSet) {
 	flags.Uint(FlagPDConcurrency, defaultPDConcurrency,
 		"concurrency pd-relative operations like split & scatter.")
 	flags.Duration(FlagBatchFlushInterval, defaultBatchFlushInterval,
-		"after how long a restore batch would be auto sended.")
+		"after how long a restore batch would be auto sent.")
 	flags.Uint(FlagDdlBatchSize, defaultFlagDdlBatchSize,
-		"batch size for ddl to create a batch of tabes once.")
+		"batch size for ddl to create a batch of tables once.")
 	flags.Bool(flagWithSysTable, false, "whether restore system privilege tables on default setting")
+	flags.StringArrayP(FlagResetSysUsers, "", []string{"cloud_admin", "root"}, "whether reset these users after restoration")
+	_ = flags.MarkHidden(FlagResetSysUsers)
 	_ = flags.MarkHidden(FlagMergeRegionSizeBytes)
 	_ = flags.MarkHidden(FlagMergeRegionKeyCount)
 	_ = flags.MarkHidden(FlagPDConcurrency)
@@ -150,6 +156,10 @@ func (cfg *RestoreCommonConfig) ParseFromFlags(flags *pflag.FlagSet) error {
 			return errors.Trace(err)
 		}
 	}
+	cfg.ResetSysUsers, err = flags.GetStringArray(FlagResetSysUsers)
+	if err != nil {
+		return errors.Trace(err)
+	}
 	return errors.Trace(err)
 }
 

br/pkg/task/restore_data.go

@@ -19,6 +19,7 @@ import (
 	"github.com/pingcap/tidb/br/pkg/storage"
 	"github.com/pingcap/tidb/br/pkg/summary"
 	"github.com/pingcap/tidb/br/pkg/utils"
+	tidbconfig "github.com/pingcap/tidb/config"
 	"go.uber.org/zap"
 )
 
@@ -71,6 +72,11 @@ func RunResolveKvData(c context.Context, g glue.Glue, cmdName string, cfg *Resto
 	defer mgr.Close()
 
 	keepaliveCfg.PermitWithoutStream = true
+	tc := tidbconfig.GetGlobalConfig()
+	tc.SkipRegisterToDashboard = true
+	tc.EnableGlobalKill = false
+	tidbconfig.StoreGlobalConfig(tc)
+
 	client := restore.NewRestoreClient(mgr.GetPDClient(), mgr.GetTLSConfig(), keepaliveCfg, false)
 
 	restoreTS, err := client.GetTS(ctx)
@@ -153,6 +159,18 @@ func RunResolveKvData(c context.Context, g glue.Glue, cmdName string, cfg *Resto
 
 	//TODO: restore volume type into origin type
 	//ModifyVolume(*ec2.ModifyVolumeInput) (*ec2.ModifyVolumeOutput, error) by backupmeta
+	// this is used for cloud restoration
+	err = client.Init(g, mgr.GetStorage())
+	if err != nil {
+		return errors.Trace(err)
+	}
+	defer client.Close()
+	log.Info("start to clear system user for cloud")
+	err = client.ClearSystemUsers(ctx, cfg.ResetSysUsers)
+
+	if err != nil {
+		return errors.Trace(err)
+	}
 
 	// since we cannot reset tiflash automaticlly. so we should start it manually
 	if err = client.ResetTiFlashReplicas(ctx, g, mgr.GetStorage()); err != nil {