br: reset cloud_admin and root after ebs restoration

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,72 @@ func isStatsTable(tableName string) bool {
 	return ok
 }
 
+// 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, filterUsers []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 execute SQL restore system database",
+				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 restore system database",
+			zap.Stringer("database", db.Name),
+			zap.String("sql", sql),
+		)
+		return nil
+	}
+
+	for _, name := range filterUsers {
+		if strings.ToLower(name) == rootUser {
+			// we cannot directly remove root account. instead, we reset its password. and let cloud_admin control it lately.
+			updateSQL := fmt.Sprintf("UPDATE %s.%s SET authentication_string='',"+
+				" Shutdown_priv='Y',"+
+				" Config_priv='Y'"+
+				" WHERE USER='root' AND Host='%%';",
+				sysDB, sysUserTableName)
+			log.Info("clear root user for cloud", zap.String("sql", updateSQL))
+			err := execSQL(updateSQL)
+			if err != nil {
+				return err
+			}
+			// we need treat the root account specially.
+			break
+		}
+	}
+
+	for tableName := range db.ExistingTables {
+		if sysPrivilegeTableMap[tableName] != "" {
+			for _, name := range filterUsers {
+				// we already handle root account before.
+				if strings.ToLower(name) != rootUser {
+					/* #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)
+					log.Info("clear system user for cloud", zap.String("sql", deleteSQL))
+					if err := execSQL(deleteSQL); 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 +273,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 +299,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/task/restore.go

@@ -62,6 +62,8 @@ const (
 	FlagPiTRBatchSize   = "pitr-batch-size"
 	FlagPiTRConcurrency = "pitr-concurrency"
 
+	FlagFilterSysUsers = "filter-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"`
+
+	FilterSysUsers []string `json:"filter-sys-users" toml:"filter-sys-users"`
 }
 
 // adjust adjusts the abnormal config value in the current config.
@@ -122,6 +126,8 @@ func DefineRestoreCommonFlags(flags *pflag.FlagSet) {
 	flags.Uint(FlagDdlBatchSize, defaultFlagDdlBatchSize,
 		"batch size for ddl to create a batch of tabes once.")
 	flags.Bool(flagWithSysTable, false, "whether restore system privilege tables on default setting")
+	flags.StringArrayP(FlagFilterSysUsers, "", []string{"cloud_admin", "root"}, "whether reset these users after restoration")
+	_ = flags.MarkHidden(FlagFilterSysUsers)
 	_ = 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.FilterSysUsers, err = flags.GetStringArray(FlagFilterSysUsers)
+	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.FilterSysUsers)
+
+	if err != nil {
+		return errors.Trace(err)
+	}
 
 	progress.Close()
 	summary.CollectDuration("restore duration", time.Since(startAll))