sharding restore and inc backup validations

loomts · loomts · commit 8b790de829b6 · 2025-12-18T11:27:09.000+08:00
diff --git a/controllers/apps/cluster/transformer_cluster_restore.go b/controllers/apps/cluster/transformer_cluster_restore.go
@@ -72,43 +72,51 @@ func (c *clusterRestoreTransformer) Transform(ctx graph.TransformContext, dag *g
 		if err != nil {
 			return err
 		}
+		targets := backup.Status.Targets
 		// obtain components that have already been assigned targets.
 		allocateTargetMap := map[string]string{}
 		restoreDoneForShardComponents := true
 		for _, v := range shardComponents {
 			if model.IsObjectDeleting(&v) {
 				continue
 			}
-			if v.Annotations[constant.RestoreDoneAnnotationKey] != "true" {
+
+			compName := v.Labels[constant.KBAppComponentLabelKey]
+			compAnnotations := c.initClusterAnnotations(compName)
+
+			if v.Annotations[constant.BackupSourceTargetAnnotationKey] != "" && v.Annotations[constant.RestoreDoneAnnotationKey] != "true" {
 				restoreDoneForShardComponents = false
 			}
 			if targetName, ok := v.Annotations[constant.BackupSourceTargetAnnotationKey]; ok {
-				compName := v.Labels[constant.KBAppComponentLabelKey]
 				allocateTargetMap[targetName] = compName
-				c.initClusterAnnotations(compName)
-				c.annotations[compName][constant.BackupSourceTargetAnnotationKey] = targetName
+				compAnnotations[constant.BackupSourceTargetAnnotationKey] = targetName
 			}
 		}
-		if len(allocateTargetMap) == len(backup.Status.Targets) {
+		if len(allocateTargetMap) == len(targets) {
 			// check if the restore is completed when all source target have allocated.
 			if err = c.cleanupRestoreAnnotationForSharding(dag, spec.Name, restoreDoneForShardComponents); err != nil {
 				return err
 			}
-			continue
 		}
-		for _, target := range backup.Status.Targets {
+		for _, target := range targets {
 			if _, ok = allocateTargetMap[target.Name]; ok {
 				continue
 			}
 			for _, compSpec := range c.shardingComps[spec.Name] {
-				if _, ok = c.annotations[compSpec.Name][constant.BackupSourceTargetAnnotationKey]; ok {
+				compAnnotations := c.initClusterAnnotations(compSpec.Name)
+				if _, ok = compAnnotations[constant.BackupSourceTargetAnnotationKey]; ok {
 					continue
 				}
-				c.initClusterAnnotations(compSpec.Name)
-				c.annotations[compSpec.Name][constant.BackupSourceTargetAnnotationKey] = target.Name
+				compAnnotations[constant.BackupSourceTargetAnnotationKey] = target.Name
 				break
 			}
 		}
+		for _, compSpec := range c.shardingComps[spec.Name] {
+			compAnnotations := c.initClusterAnnotations(compSpec.Name)
+			if compAnnotations[constant.BackupSourceTargetAnnotationKey] == "" {
+				compAnnotations[constant.SkipRestoreAnnotationKey] = "true"
+			}
+		}
 	}
 	// if component needs to do post ready restore after cluster is running, annotate component
 	if c.Cluster.Status.Phase == appsv1.RunningClusterPhase {
@@ -132,13 +140,14 @@ func (c *clusterRestoreTransformer) Transform(ctx graph.TransformContext, dag *g
 	return nil
 }
 
-func (c *clusterRestoreTransformer) initClusterAnnotations(compName string) {
+func (c *clusterRestoreTransformer) initClusterAnnotations(compName string) map[string]string {
 	if c.annotations == nil {
-		c.annotations = map[string]map[string]string{}
+		c.annotations = make(map[string]map[string]string)
 	}
 	if c.annotations[compName] == nil {
-		c.annotations[compName] = map[string]string{}
+		c.annotations[compName] = make(map[string]string)
 	}
+	return c.annotations[compName]
 }
 
 func (c *clusterRestoreTransformer) cleanupRestoreAnnotationForSharding(dag *graph.DAG,
diff --git a/controllers/dataprotection/backup_controller.go b/controllers/dataprotection/backup_controller.go
@@ -1071,7 +1071,7 @@ func prepare4Incremental(request *dpbackup.Request) (*dpbackup.Request, error) {
 		return nil, fmt.Errorf("backupRepo for incremental backup can't be empty")
 	}
 	// get and validate parent backup
-	parentBackup, err := GetParentBackup(request.Ctx, request.Client, request.Backup, request.BackupMethod, request.BackupRepo.Name)
+	parentBackup, err := GetParentBackup(request.Ctx, request.Client, request.Backup, request.BackupMethod, request.BackupPolicy, request.BackupRepo.Name)
 	if err != nil {
 		return nil, err
 	}
diff --git a/controllers/dataprotection/utils.go b/controllers/dataprotection/utils.go
@@ -590,7 +590,7 @@ func fromFlattenName(flatten string) (name string, namespace string) {
 // then validate and return the parent backup.
 // If parentBackupName is not specified, find the latest valid parent backup.
 func GetParentBackup(ctx context.Context, cli client.Client, backup *dpv1alpha1.Backup,
-	backupMethod *dpv1alpha1.BackupMethod, backupRepoName string) (*dpv1alpha1.Backup, error) {
+	backupMethod *dpv1alpha1.BackupMethod, backupPolicy *dpv1alpha1.BackupPolicy, backupRepoName string) (*dpv1alpha1.Backup, error) {
 	if backup == nil || backupMethod == nil {
 		return nil, fmt.Errorf("backup or backupMethod is nil")
 	}
@@ -616,16 +616,23 @@ func GetParentBackup(ctx context.Context, cli client.Client, backup *dpv1alpha1.
 		}, parentBackup); err != nil {
 			return nil, err
 		}
-		if err := ValidateParentBackup(ctx, cli, backup, parentBackup, backupMethod, backupRepoName); err != nil {
+		if err := ValidateParentBackup(ctx, cli, backup, parentBackup, backupMethod, backupPolicy, backupRepoName); err != nil {
 			return nil, fmt.Errorf("failed to validate parent backup %s: %w", parentBackupName, err)
 		}
 		return parentBackup, nil
 	}
-	parentBackup, err := FindParentBackupIfNotSet(ctx, cli, backup, backupMethod, backupRepoName)
+	parentBackup, err := FindParentBackupIfNotSet(ctx, cli, backup, backupMethod, backupPolicy, backupRepoName)
 	if err != nil {
 		return nil, fmt.Errorf("failed to find parent backup: %w", err)
 	}
 	if parentBackup == nil {
+		// output sharding message
+		expectedTargets := dputils.GetBackupTargets(backupPolicy, backupMethod)
+		if len(expectedTargets) > 0 {
+			return nil, fmt.Errorf("failed to find a valid parent backup for backup %s/%s: "+
+				"current backup expects %d shards, this may be due to shard count mismatch with existing backups",
+				backup.Namespace, backup.Name, len(expectedTargets))
+		}
 		return nil, fmt.Errorf("failed to find a valid parent backup for backup %s/%s", backup.Namespace, backup.Name)
 	}
 	return parentBackup, nil
@@ -637,7 +644,7 @@ func GetParentBackup(ctx context.Context, cli client.Client, backup *dpv1alpha1.
 // b. return the latest incremental backup.
 // c. return the latest full backup if incremental backups are not found.
 func FindParentBackupIfNotSet(ctx context.Context, cli client.Client, backup *dpv1alpha1.Backup,
-	backupMethod *dpv1alpha1.BackupMethod, backupRepoName string) (*dpv1alpha1.Backup, error) {
+	backupMethod *dpv1alpha1.BackupMethod, backupPolicy *dpv1alpha1.BackupPolicy, backupRepoName string) (*dpv1alpha1.Backup, error) {
 	getLatestBackup := func(backupList []*dpv1alpha1.Backup) *dpv1alpha1.Backup {
 		if len(backupList) == 0 {
 			return nil
@@ -655,7 +662,7 @@ func FindParentBackupIfNotSet(ctx context.Context, cli client.Client, backup *dp
 			client.MatchingLabels(labels)); err != nil && !apierrors.IsNotFound(err) {
 			return nil, err
 		}
-		filteredbackupList := FilterParentBackups(ctx, cli, backupList, backup, backupMethod, incremental, backupRepoName)
+		filteredbackupList := FilterParentBackups(ctx, cli, backupList, backup, backupMethod, backupPolicy, incremental, backupRepoName)
 		return getLatestBackup(filteredbackupList), nil
 	}
 
@@ -701,13 +708,13 @@ func FindParentBackupIfNotSet(ctx context.Context, cli client.Client, backup *dp
 
 // FilterParentBackups filters the parent backups by backup phase, backup method and end time.
 func FilterParentBackups(ctx context.Context, cli client.Client, backupList *dpv1alpha1.BackupList, targetBackup *dpv1alpha1.Backup,
-	backupMethod *dpv1alpha1.BackupMethod, incremental bool, backupRepoName string) []*dpv1alpha1.Backup {
+	backupMethod *dpv1alpha1.BackupMethod, backupPolicy *dpv1alpha1.BackupPolicy, incremental bool, backupRepoName string) []*dpv1alpha1.Backup {
 	var res []*dpv1alpha1.Backup
 	if backupList == nil || len(backupList.Items) == 0 {
 		return res
 	}
 	for i, backup := range backupList.Items {
-		if err := ValidateParentBackup(ctx, cli, targetBackup, &backup, backupMethod, backupRepoName); err != nil {
+		if err := ValidateParentBackup(ctx, cli, targetBackup, &backup, backupMethod, backupPolicy, backupRepoName); err != nil {
 			continue
 		}
 		// backups are listed by backup type label, validate if the backup method matches
@@ -728,7 +735,7 @@ func FilterParentBackups(ctx context.Context, cli client.Client, backupList *dpv
 
 // ValidateParentBackup validates the parent backup.
 func ValidateParentBackup(ctx context.Context, cli client.Client, backup *dpv1alpha1.Backup, parentBackup *dpv1alpha1.Backup,
-	backupMethod *dpv1alpha1.BackupMethod, backupRepoName string) error {
+	backupMethod *dpv1alpha1.BackupMethod, backupPolicy *dpv1alpha1.BackupPolicy, backupRepoName string) error {
 	// validate parent backup is completed
 	if parentBackup.Status.Phase != dpv1alpha1.BackupPhaseCompleted {
 		return fmt.Errorf("parent backup %s/%s is not completed", parentBackup.Namespace, parentBackup.Name)
@@ -759,6 +766,13 @@ func ValidateParentBackup(ctx context.Context, cli client.Client, backup *dpv1al
 				parentBackup.Namespace, parentBackup.Name, err)
 		}
 	}
+	// validate shard count consistency for incremental backup
+	expectedTargets := dputils.GetBackupTargets(backupPolicy, backupMethod)
+	if len(expectedTargets) > 0 && len(parentBackup.Status.Targets) != len(expectedTargets) {
+		return fmt.Errorf("parent backup %s/%s has %d shards, but current backup expects %d shards; "+
+			"incremental backup requires consistent shard count",
+			parentBackup.Namespace, parentBackup.Name, len(parentBackup.Status.Targets), len(expectedTargets))
+	}
 	return nil
 }
 
diff --git a/pkg/constant/annotations.go b/pkg/constant/annotations.go
@@ -33,6 +33,7 @@ const (
 	RestoreFromBackupAnnotationKey       = "kubeblocks.io/restore-from-backup"
 	RestoreDoneAnnotationKey             = "kubeblocks.io/restore-done"
 	BackupSourceTargetAnnotationKey      = "kubeblocks.io/backup-source-target" // RestoreFromBackupAnnotationKey specifies the component to recover from the backup.
+	SkipRestoreAnnotationKey             = "kubeblocks.io/skip-restore"         // SkipRestoreAnnotationKey indicates the shard component should skip sharding restore scheduling.
 
 	KBAppClusterUIDKey                   = "apps.kubeblocks.io/cluster-uid"
 	BackupPolicyTemplateAnnotationKey    = "apps.kubeblocks.io/backup-policy-template"
diff --git a/pkg/controller/plan/restore.go b/pkg/controller/plan/restore.go
@@ -86,6 +86,9 @@ func NewRestoreManager(ctx context.Context,
 }
 
 func (r *RestoreManager) DoRestore(comp *component.SynthesizedComponent, compObj *appsv1.Component, postProvisionDone bool) error {
+	if compObj.Annotations[constant.SkipRestoreAnnotationKey] == "true" {
+		return nil
+	}
 	backupObj, err := r.initFromAnnotation(comp, compObj)
 	if err != nil {
 		return err

Original file line number	Diff line number	Diff line change
`@@ -72,43 +72,51 @@ func (c clusterRestoreTransformer) Transform(ctx graph.TransformContext, dag g`
`72`	`72`	`if err != nil {`
`73`	`73`	`return err`
`74`	`74`	`}`
	`75`	`+ targets := backup.Status.Targets`
`75`	`76`	`// obtain components that have already been assigned targets.`
`76`	`77`	`allocateTargetMap := map[string]string{}`
`77`	`78`	`restoreDoneForShardComponents := true`
`78`	`79`	`for _, v := range shardComponents {`
`79`	`80`	`if model.IsObjectDeleting(&v) {`
`80`	`81`	`continue`
`81`	`82`	`}`
`82`		`- if v.Annotations[constant.RestoreDoneAnnotationKey] != "true" {`
	`83`	`+`
	`84`	`+ compName := v.Labels[constant.KBAppComponentLabelKey]`
	`85`	`+ compAnnotations := c.initClusterAnnotations(compName)`
	`86`	`+`
	`87`	`+ if v.Annotations[constant.BackupSourceTargetAnnotationKey] != "" && v.Annotations[constant.RestoreDoneAnnotationKey] != "true" {`
`83`	`88`	`restoreDoneForShardComponents = false`
`84`	`89`	`}`
`85`	`90`	`if targetName, ok := v.Annotations[constant.BackupSourceTargetAnnotationKey]; ok {`
`86`		`- compName := v.Labels[constant.KBAppComponentLabelKey]`
`87`	`91`	`allocateTargetMap[targetName] = compName`
`88`		`- c.initClusterAnnotations(compName)`
`89`		`- c.annotations[compName][constant.BackupSourceTargetAnnotationKey] = targetName`
	`92`	`+ compAnnotations[constant.BackupSourceTargetAnnotationKey] = targetName`
`90`	`93`	`}`
`91`	`94`	`}`
`92`		`- if len(allocateTargetMap) == len(backup.Status.Targets) {`
	`95`	`+ if len(allocateTargetMap) == len(targets) {`
`93`	`96`	`// check if the restore is completed when all source target have allocated.`
`94`	`97`	`if err = c.cleanupRestoreAnnotationForSharding(dag, spec.Name, restoreDoneForShardComponents); err != nil {`
`95`	`98`	`return err`
`96`	`99`	`}`
`97`		`- continue`
`98`	`100`	`}`
`99`		`- for _, target := range backup.Status.Targets {`
	`101`	`+ for _, target := range targets {`
`100`	`102`	`if _, ok = allocateTargetMap[target.Name]; ok {`
`101`	`103`	`continue`
`102`	`104`	`}`
`103`	`105`	`for _, compSpec := range c.shardingComps[spec.Name] {`
`104`		`- if _, ok = c.annotations[compSpec.Name][constant.BackupSourceTargetAnnotationKey]; ok {`
	`106`	`+ compAnnotations := c.initClusterAnnotations(compSpec.Name)`
	`107`	`+ if _, ok = compAnnotations[constant.BackupSourceTargetAnnotationKey]; ok {`
`105`	`108`	`continue`
`106`	`109`	`}`
`107`		`- c.initClusterAnnotations(compSpec.Name)`
`108`		`- c.annotations[compSpec.Name][constant.BackupSourceTargetAnnotationKey] = target.Name`
	`110`	`+ compAnnotations[constant.BackupSourceTargetAnnotationKey] = target.Name`
`109`	`111`	`break`
`110`	`112`	`}`
`111`	`113`	`}`
	`114`	`+ for _, compSpec := range c.shardingComps[spec.Name] {`
	`115`	`+ compAnnotations := c.initClusterAnnotations(compSpec.Name)`
	`116`	`+ if compAnnotations[constant.BackupSourceTargetAnnotationKey] == "" {`
	`117`	`+ compAnnotations[constant.SkipRestoreAnnotationKey] = "true"`
	`118`	`+ }`
	`119`	`+ }`
`112`	`120`	`}`
`113`	`121`	`// if component needs to do post ready restore after cluster is running, annotate component`
`114`	`122`	`if c.Cluster.Status.Phase == appsv1.RunningClusterPhase {`
`@@ -132,13 +140,14 @@ func (c clusterRestoreTransformer) Transform(ctx graph.TransformContext, dag g`
`132`	`140`	`return nil`
`133`	`141`	`}`
`134`	`142`
`135`		`-func (c *clusterRestoreTransformer) initClusterAnnotations(compName string) {`
	`143`	`+func (c *clusterRestoreTransformer) initClusterAnnotations(compName string) map[string]string {`
`136`	`144`	`if c.annotations == nil {`
`137`		`- c.annotations = map[string]map[string]string{}`
	`145`	`+ c.annotations = make(map[string]map[string]string)`
`138`	`146`	`}`
`139`	`147`	`if c.annotations[compName] == nil {`
`140`		`- c.annotations[compName] = map[string]string{}`
	`148`	`+ c.annotations[compName] = make(map[string]string)`
`141`	`149`	`}`
	`150`	`+ return c.annotations[compName]`
`142`	`151`	`}`
`143`	`152`
`144`	`153`	`func (c clusterRestoreTransformer) cleanupRestoreAnnotationForSharding(dag graph.DAG,`
Original file line number	Diff line number	Diff line change
`@@ -1071,7 +1071,7 @@ func prepare4Incremental(request dpbackup.Request) (dpbackup.Request, error) {`
`1071`	`1071`	`return nil, fmt.Errorf("backupRepo for incremental backup can't be empty")`
`1072`	`1072`	`}`
`1073`	`1073`	`// get and validate parent backup`
`1074`		`- parentBackup, err := GetParentBackup(request.Ctx, request.Client, request.Backup, request.BackupMethod, request.BackupRepo.Name)`
	`1074`	`+ parentBackup, err := GetParentBackup(request.Ctx, request.Client, request.Backup, request.BackupMethod, request.BackupPolicy, request.BackupRepo.Name)`
`1075`	`1075`	`if err != nil {`
`1076`	`1076`	`return nil, err`
`1077`	`1077`	`}`
Original file line number	Diff line number	Diff line change
`@@ -86,6 +86,9 @@ func NewRestoreManager(ctx context.Context,`
`86`	`86`	`}`
`87`	`87`
`88`	`88`	`func (r RestoreManager) DoRestore(comp component.SynthesizedComponent, compObj *appsv1.Component, postProvisionDone bool) error {`
	`89`	`+ if compObj.Annotations[constant.SkipRestoreAnnotationKey] == "true" {`
	`90`	`+ return nil`
	`91`	`+ }`
`89`	`92`	`backupObj, err := r.initFromAnnotation(comp, compObj)`
`90`	`93`	`if err != nil {`
`91`	`94`	`return err`