fix(dataprotection): pass target metadata to post-ready jobs (#10838)

Author: leon-ape

controllers/dataprotection/restore_controller_test.go

@@ -64,6 +64,7 @@ var _ = Describe("Restore Controller test", func() {
 
 			// namespaced
 			testapps.ClearResourcesWithRemoveFinalizerOption(&testCtx, generics.ClusterSignature, true, inNS, ml)
+			testapps.ClearResources(&testCtx, generics.ComponentSignature, inNS, ml)
 			testapps.ClearResources(&testCtx, generics.PodSignature, inNS, ml)
 			testapps.ClearResourcesWithRemoveFinalizerOption(&testCtx, generics.BackupSignature, true, inNS)
 
@@ -487,6 +488,10 @@ var _ = Describe("Restore Controller test", func() {
 			BeforeEach(func() {
 				By("fake a new cluster")
 				_ = testdp.NewFakeCluster(&testCtx)
+				testapps.NewComponentFactory(testCtx.DefaultNamespace,
+					constant.GenerateClusterComponentName(testdp.ClusterName, testdp.ComponentName), "test-cmpd").
+					SetServiceVersion("8.0.30").
+					Create(&testCtx)
 			})
 
 			It("test post ready actions", func() {

pkg/dataprotection/restore/builder.go

@@ -310,6 +310,23 @@ func (r *restoreJobBuilder) addTargetPodAndCredentialEnv(pod *corev1.Pod,
 	return r
 }
 
+func (r *restoreJobBuilder) overridePostReadyTargetEnv(targetEnv []corev1.EnvVar) *restoreJobBuilder {
+	if len(targetEnv) == 0 {
+		return r
+	}
+	env := make([]corev1.EnvVar, 0, len(r.env)+len(targetEnv))
+	for i := range r.env {
+		if r.env[i].Name == dptypes.DPTargetClusterTopology ||
+			r.env[i].Name == dptypes.DPTargetServiceVersion {
+			continue
+		}
+		env = append(env, r.env[i])
+	}
+	env = append(env, targetEnv...)
+	r.env = env
+	return r
+}
+
 // builderRestoreJobName builds restore job name.
 func (r *restoreJobBuilder) builderRestoreJobName(jobIndex int) string {
 	jobName := fmt.Sprintf("restore-%s-%s-%s-%d", strings.ToLower(string(r.stage)), r.restore.UID[:8], r.backupSet.Backup.Name, jobIndex)

pkg/dataprotection/restore/builder_test.go

@@ -0,0 +1,55 @@
+/*
+Copyright (C) 2022-2026 ApeCloud Co., Ltd
+
+This file is part of KubeBlocks project
+
+This program is free software: you can redistribute it and/or modify
+it under the terms of the GNU Affero General Public License as published by
+the Free Software Foundation, either version 3 of the License, or
+(at your option) any later version.
+
+This program is distributed in the hope that it will be useful,
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+GNU Affero General Public License for more details.
+
+You should have received a copy of the GNU Affero General Public License
+along with this program.  If not, see <http://www.gnu.org/licenses/>.
+*/
+
+package restore
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	corev1 "k8s.io/api/core/v1"
+
+	dptypes "github.com/apecloud/kubeblocks/pkg/dataprotection/types"
+)
+
+func TestRestoreJobBuilderOverridePostReadyTargetEnv(t *testing.T) {
+	builder := &restoreJobBuilder{env: []corev1.EnvVar{
+		{Name: "KEEP", Value: "kept"},
+		{Name: dptypes.DPTargetClusterTopology, Value: "restore-value"},
+		{Name: dptypes.DPTargetClusterTopology, Value: "pod-value"},
+		{Name: dptypes.DPTargetServiceVersion, Value: "spoofed-version"},
+	}}
+	builder.overridePostReadyTargetEnv([]corev1.EnvVar{
+		{Name: dptypes.DPTargetClusterTopology, Value: "shared-nothing"},
+		{Name: dptypes.DPTargetServiceVersion, Value: "3.4.0"},
+	})
+
+	values := func(name string) []string {
+		var result []string
+		for i := range builder.env {
+			if builder.env[i].Name == name {
+				result = append(result, builder.env[i].Value)
+			}
+		}
+		return result
+	}
+	require.Equal(t, []string{"kept"}, values("KEEP"))
+	require.Equal(t, []string{"shared-nothing"}, values(dptypes.DPTargetClusterTopology))
+	require.Equal(t, []string{"3.4.0"}, values(dptypes.DPTargetServiceVersion))
+}

pkg/dataprotection/restore/manager.go

@@ -36,6 +36,7 @@ import (
 	"sigs.k8s.io/controller-runtime/pkg/client"
 	"sigs.k8s.io/controller-runtime/pkg/controller/controllerutil"
 
+	appsv1 "github.com/apecloud/kubeblocks/apis/apps/v1"
 	dpv1alpha1 "github.com/apecloud/kubeblocks/apis/dataprotection/v1alpha1"
 	"github.com/apecloud/kubeblocks/pkg/constant"
 	"github.com/apecloud/kubeblocks/pkg/controller/instanceset"
@@ -661,6 +662,50 @@ func (r *RestoreManager) isJobForRestoreAction(job *batchv1.Job) bool {
 	return restoreNamespace == "" || restoreNamespace == r.Restore.Namespace
 }
 
+func (r *RestoreManager) postReadyTargetEnv(reqCtx intctrlutil.RequestCtx, cli client.Client, pod *corev1.Pod) ([]corev1.EnvVar, error) {
+	clusterName := pod.Labels[constant.AppInstanceLabelKey]
+	componentName := pod.Labels[constant.KBAppComponentLabelKey]
+	if clusterName == "" || componentName == "" {
+		return nil, nil
+	}
+	cluster := &appsv1.Cluster{}
+	if err := cli.Get(reqCtx.Ctx, types.NamespacedName{Namespace: pod.Namespace, Name: clusterName}, cluster); err != nil {
+		return nil, err
+	}
+	component := &appsv1.Component{}
+	componentKey := types.NamespacedName{
+		Namespace: pod.Namespace,
+		Name:      constant.GenerateClusterComponentName(clusterName, componentName),
+	}
+	if err := cli.Get(reqCtx.Ctx, componentKey, component); err != nil {
+		return nil, err
+	}
+
+	serviceVersion := component.Spec.ServiceVersion
+	if templateName := pod.Labels[constant.KBAppInstanceTemplateLabelKey]; templateName != "" {
+		found := false
+		for i := range component.Spec.Instances {
+			instance := component.Spec.Instances[i]
+			if instance.Name != templateName {
+				continue
+			}
+			found = true
+			if instance.ServiceVersion != "" {
+				serviceVersion = instance.ServiceVersion
+			}
+			break
+		}
+		if !found {
+			return nil, intctrlutil.NewFatalError(fmt.Sprintf(
+				"target Pod %s/%s references unknown instance template %q", pod.Namespace, pod.Name, templateName))
+		}
+	}
+	return []corev1.EnvVar{
+		{Name: dptypes.DPTargetClusterTopology, Value: cluster.Spec.Topology},
+		{Name: dptypes.DPTargetServiceVersion, Value: serviceVersion},
+	}, nil
+}
+
 // BuildPostReadyActionJobs builds the post ready jobs.
 func (r *RestoreManager) BuildPostReadyActionJobs(reqCtx intctrlutil.RequestCtx, cli client.Client, backupSet BackupActionSet, target *dpv1alpha1.BackupStatusTarget, step int) ([]*batchv1.Job, error) {
 	readyConfig := r.Restore.Spec.ReadyConfig
@@ -705,7 +750,7 @@ func (r *RestoreManager) BuildPostReadyActionJobs(reqCtx intctrlutil.RequestCtx,
 			return nil, err
 		}
 		sort.Sort(intctrlutil.ByPodName(targetPodList.Items))
-		buildJob := func(targetPod *corev1.Pod, sourceTargetPodName string, index int) *batchv1.Job {
+		buildJob := func(targetPod *corev1.Pod, sourceTargetPodName string, index int) (*batchv1.Job, error) {
 			if boolptr.IsSetToTrue(actionSpec.Job.RunOnTargetPodNode) {
 				jobBuilder.resetSpecificVolumesAndMounts()
 				jobBuilder.setNodeNameToNodeSelector(targetPod.Spec.NodeName)
@@ -719,15 +764,20 @@ func (r *RestoreManager) BuildPostReadyActionJobs(reqCtx intctrlutil.RequestCtx,
 					}
 				}
 			}
+			targetEnv, err := r.postReadyTargetEnv(reqCtx, cli, targetPod)
+			if err != nil {
+				return nil, err
+			}
 			return jobBuilder.setImage(actionSpec.Job.Image).
 				setJobName(buildJobName(index)).
 				addCommonEnv(sourceTargetPodName).
 				attachBackupRepo().
 				setCommand(actionSpec.Job.Command).
 				setToleration(targetPod.Spec.Tolerations).
 				addTargetPodAndCredentialEnv(targetPod, readyConfig.ConnectionCredential, &target.BackupTarget).
+				overridePostReadyTargetEnv(targetEnv).
 				setServiceAccount(r.WorkerServiceAccount).
-				build()
+				build(), nil
 		}
 
 		if podSelector.Strategy == dpv1alpha1.PodSelectionStrategyAny {
@@ -747,7 +797,11 @@ func (r *RestoreManager) BuildPostReadyActionJobs(reqCtx intctrlutil.RequestCtx,
 				// no need to recover the volume when the pod selection policy is 'All' and sourceTargetPodName is not found.
 				continue
 			}
-			jobs = append(jobs, buildJob(&targetPodList.Items[i], sourceTargetPodName, i))
+			job, err := buildJob(&targetPodList.Items[i], sourceTargetPodName, i)
+			if err != nil {
+				return nil, err
+			}
+			jobs = append(jobs, job)
 		}
 		return jobs, nil
 	}

pkg/dataprotection/restore/manager_test.go

@@ -20,21 +20,27 @@ along with this program.  If not, see <http://www.gnu.org/licenses/>.
 package restore
 
 import (
+	"context"
 	"fmt"
 	"strconv"
 	"strings"
+	"testing"
 	"time"
 
 	. "github.com/onsi/ginkgo/v2"
 	. "github.com/onsi/gomega"
+	"github.com/stretchr/testify/require"
 	batchv1 "k8s.io/api/batch/v1"
 	corev1 "k8s.io/api/core/v1"
 	apierrors "k8s.io/apimachinery/pkg/api/errors"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/apimachinery/pkg/types"
 	ctrl "sigs.k8s.io/controller-runtime"
 	"sigs.k8s.io/controller-runtime/pkg/client"
+	"sigs.k8s.io/controller-runtime/pkg/client/fake"
 
+	appsv1 "github.com/apecloud/kubeblocks/apis/apps/v1"
 	dpv1alpha1 "github.com/apecloud/kubeblocks/apis/dataprotection/v1alpha1"
 	"github.com/apecloud/kubeblocks/pkg/constant"
 	intctrlutil "github.com/apecloud/kubeblocks/pkg/controllerutil"
@@ -47,6 +53,82 @@ import (
 	viper "github.com/apecloud/kubeblocks/pkg/viperx"
 )
 
+func postReadyJobEnvValues(job *batchv1.Job, name string) []string {
+	var values []string
+	for i := range job.Spec.Template.Spec.Containers[0].Env {
+		env := job.Spec.Template.Spec.Containers[0].Env[i]
+		if env.Name == name {
+			values = append(values, env.Value)
+		}
+	}
+	return values
+}
+
+func TestRestoreManagerPostReadyTargetEnv(t *testing.T) {
+	const (
+		namespace     = "default"
+		clusterName   = "target"
+		componentName = "mysql"
+	)
+	scheme := runtime.NewScheme()
+	require.NoError(t, appsv1.AddToScheme(scheme))
+	cluster := &appsv1.Cluster{
+		ObjectMeta: metav1.ObjectMeta{Namespace: namespace, Name: clusterName},
+		Spec:       appsv1.ClusterSpec{Topology: "shared-nothing"},
+	}
+	component := &appsv1.Component{
+		ObjectMeta: metav1.ObjectMeta{
+			Namespace: namespace,
+			Name:      constant.GenerateClusterComponentName(clusterName, componentName),
+		},
+		Spec: appsv1.ComponentSpec{
+			ServiceVersion: "3.3.2",
+			Instances: []appsv1.InstanceTemplate{
+				{Name: "canary", ServiceVersion: "3.4.0"},
+				{Name: "inherited"},
+			},
+		},
+	}
+	cli := fake.NewClientBuilder().WithScheme(scheme).WithObjects(cluster, component).Build()
+	manager := &RestoreManager{Restore: &dpv1alpha1.Restore{}}
+	reqCtx := intctrlutil.RequestCtx{Ctx: context.Background()}
+
+	tests := []struct {
+		name             string
+		instanceTemplate string
+		wantVersion      string
+	}{
+		{name: "component default", wantVersion: "3.3.2"},
+		{name: "instance override", instanceTemplate: "canary", wantVersion: "3.4.0"},
+		{name: "instance inherits component", instanceTemplate: "inherited", wantVersion: "3.3.2"},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			pod := &corev1.Pod{ObjectMeta: metav1.ObjectMeta{
+				Namespace: namespace,
+				Name:      "target-mysql-0",
+				Labels: map[string]string{
+					constant.AppInstanceLabelKey:           clusterName,
+					constant.KBAppComponentLabelKey:        componentName,
+					constant.KBAppInstanceTemplateLabelKey: tt.instanceTemplate,
+				},
+			}}
+
+			env, err := manager.postReadyTargetEnv(reqCtx, cli, pod)
+
+			require.NoError(t, err)
+			require.Equal(t, []corev1.EnvVar{
+				{Name: dptypes.DPTargetClusterTopology, Value: "shared-nothing"},
+				{Name: dptypes.DPTargetServiceVersion, Value: tt.wantVersion},
+			}, env)
+		})
+	}
+
+	env, err := manager.postReadyTargetEnv(reqCtx, cli, &corev1.Pod{})
+	require.NoError(t, err)
+	require.Empty(t, env, "non-KubeBlocks target Pods must keep their existing behavior")
+}
+
 var _ = Describe("RestoreManager Test", func() {
 
 	cleanEnv := func() {
@@ -58,6 +140,7 @@ var _ = Describe("RestoreManager Test", func() {
 		// namespaced
 		testapps.ClearResources(&testCtx, generics.PodSignature, inNS, ml)
 		testapps.ClearResources(&testCtx, generics.ClusterSignature, inNS, ml)
+		testapps.ClearResources(&testCtx, generics.ComponentSignature, inNS, ml)
 		testapps.ClearResourcesWithRemoveFinalizerOption(&testCtx, generics.BackupSignature, true, inNS)
 
 		// wait all backup to be deleted, otherwise the controller maybe create
@@ -538,9 +621,23 @@ var _ = Describe("RestoreManager Test", func() {
 			restoreMGR, backupSet := initResources(reqCtx, 0, false, func(f *testdp.MockRestoreFactory) {
 				f.SetConnectCredential(testdp.ClusterName).SetJobActionConfig(matchLabels).SetExecActionConfig(matchLabels)
 			})
+			restoreMGR.Restore.Spec.Env = append(restoreMGR.Restore.Spec.Env,
+				corev1.EnvVar{Name: dptypes.DPTargetClusterTopology, Value: "spoofed-topology"},
+				corev1.EnvVar{Name: dptypes.DPTargetServiceVersion, Value: "spoofed-version"})
 
 			By("create cluster to restore")
-			testdp.NewFakeCluster(&testCtx)
+			clusterInfo := testdp.NewFakeCluster(&testCtx)
+			Expect(testapps.ChangeObj(&testCtx, clusterInfo.Cluster, func(cluster *appsv1.Cluster) {
+				cluster.Spec.Topology = "shared-nothing"
+			})).Should(Succeed())
+			testapps.NewComponentFactory(testCtx.DefaultNamespace,
+				constant.GenerateClusterComponentName(testdp.ClusterName, testdp.ComponentName), "test-cmpd").
+				SetServiceVersion("3.3.2").
+				AddInstances(appsv1.InstanceTemplate{Name: "canary", ServiceVersion: "3.4.0"}).
+				Create(&testCtx)
+			Expect(testapps.ChangeObj(&testCtx, clusterInfo.TargetPod, func(pod *corev1.Pod) {
+				pod.Labels[constant.KBAppInstanceTemplateLabelKey] = "canary"
+			})).Should(Succeed())
 
 			By("test with execAction and expect for creating 2 exec job")
 			target := utils.GetBackupStatusTarget(backupSet.Backup, restoreMGR.Restore.Spec.Backup.SourceTargetName)
@@ -558,6 +655,8 @@ var _ = Describe("RestoreManager Test", func() {
 			Expect(err).ShouldNot(HaveOccurred())
 			// count of job should equal to 1
 			Expect(len(jobs)).Should(Equal(1))
+			Expect(postReadyJobEnvValues(jobs[0], dptypes.DPTargetClusterTopology)).Should(Equal([]string{"shared-nothing"}))
+			Expect(postReadyJobEnvValues(jobs[0], dptypes.DPTargetServiceVersion)).Should(Equal([]string{"3.4.0"}))
 			// test timeZone transform
 			var backupStopTimeEnv string
 			for _, v := range jobs[0].Spec.Template.Spec.Containers[0].Env {

pkg/dataprotection/types/constant.go

@@ -120,6 +120,10 @@ const (
 	DPTargetPodName = "DP_TARGET_POD_NAME"
 	// DPTargetPodRole the target pod role
 	DPTargetPodRole = "DP_TARGET_POD_ROLE"
+	// DPTargetClusterTopology is the topology of the restore target Cluster.
+	DPTargetClusterTopology = "DP_TARGET_CLUSTER_TOPOLOGY"
+	// DPTargetServiceVersion is the expected serviceVersion of the restore target instance.
+	DPTargetServiceVersion = "DP_TARGET_SERVICE_VERSION"
 	// DPBackupBasePath the base path for backup data in the storage
 	// In a backup action pod, it equals ${DP_BACKUP_ROOT_PATH}/${DP_BACKUP_NAME}/${DP_TARGET_RELATIVE_PATH}
 	DPBackupBasePath = "DP_BACKUP_BASE_PATH"