feat: inject NodeAffinity into DataBackup status after successful backup

pkg/controllers/v1alpha1/databackup/status_handler.go

@@ -21,6 +21,7 @@ import (
 
 	"github.com/fluid-cloudnative/fluid/api/v1alpha1"
 	"github.com/fluid-cloudnative/fluid/pkg/common"
+	"github.com/fluid-cloudnative/fluid/pkg/dataflow"
 	"github.com/fluid-cloudnative/fluid/pkg/dataoperation"
 	"github.com/fluid-cloudnative/fluid/pkg/runtime"
 	"github.com/fluid-cloudnative/fluid/pkg/utils"
@@ -52,7 +53,13 @@ func (o *OnceHandler) GetOperationStatus(ctx runtime.ReconcileRequestContext, op
 		return
 	}
 
-	// TODO: inject nodeaffinity like other data operations when using job instead of pod
+	if kubeclient.IsSucceededPod(backupPod) && result.NodeAffinity == nil {
+		result.NodeAffinity, err = dataflow.GenerateNodeAffinityFromPod(backupPod)
+		if err != nil {
+			ctx.Log.V(1).Info("NodeAffinity not injected", "reason", err.Error())
+			err = nil
+		}
+	}
 
 	var finishTime time.Time
 	if len(backupPod.Status.Conditions) != 0 {

pkg/controllers/v1alpha1/databackup/status_handler_test.go

@@ -184,5 +184,76 @@ var _ = Describe("OnceHandler", func() {
 			Expect(opStatus.Conditions).To(HaveLen(1))
 			Expect(opStatus.Conditions[0].LastTransitionTime.Time).To(BeTemporally("~", conditionTime.Time, time.Second))
 		})
+
+		It("should inject NodeAffinity when backup pod succeeds with dataflow annotations", func() {
+			pod := &corev1.Pod{
+				ObjectMeta: v1.ObjectMeta{
+					Name:      mockDataBackup.GetName() + "-pod",
+					Namespace: mockDataBackup.GetNamespace(),
+					Annotations: map[string]string{
+						common.AnnotationDataFlowAffinityInject:                          "true",
+						common.AnnotationDataFlowCustomizedAffinityPrefix + "node-label": "node-value",
+					},
+				},
+				Status: corev1.PodStatus{
+					Phase: corev1.PodSucceeded,
+					Conditions: []corev1.PodCondition{
+						{
+							Type:               corev1.PodReady,
+							Status:             corev1.ConditionFalse,
+							LastTransitionTime: v1.Now(),
+						},
+					},
+				},
+			}
+			c := fake.NewFakeClientWithScheme(testScheme, pod, mockDataBackup)
+			handler := &OnceHandler{dataBackup: mockDataBackup}
+			ctx := cruntime.ReconcileRequestContext{
+				NamespacedName: types.NamespacedName{Namespace: "default", Name: "test"},
+				Log:            fake.NullLogger(),
+				Client:         c,
+			}
+			result, err := handler.GetOperationStatus(ctx, &mockDataBackup.Status)
+			Expect(err).NotTo(HaveOccurred())
+			Expect(result.Phase).To(Equal(common.PhaseComplete))
+			Expect(result.NodeAffinity).NotTo(BeNil())
+			Expect(result.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution).NotTo(BeNil())
+			terms := result.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution.NodeSelectorTerms
+			Expect(terms).To(HaveLen(1))
+			Expect(terms[0].MatchExpressions).To(HaveLen(1))
+			Expect(terms[0].MatchExpressions[0].Key).To(Equal("node-label"))
+			Expect(terms[0].MatchExpressions[0].Values).To(ContainElement("node-value"))
+		})
+
+		It("should not inject NodeAffinity when backup pod succeeds without dataflow annotations", func() {
+			pod := &corev1.Pod{
+				ObjectMeta: v1.ObjectMeta{
+					Name:      mockDataBackup.GetName() + "-pod",
+					Namespace: mockDataBackup.GetNamespace(),
+				},
+				Status: corev1.PodStatus{
+					Phase: corev1.PodSucceeded,
+					Conditions: []corev1.PodCondition{
+						{
+							Type:               corev1.PodReady,
+							Status:             corev1.ConditionFalse,
+							LastTransitionTime: v1.Now(),
+						},
+					},
+				},
+			}
+			c := fake.NewFakeClientWithScheme(testScheme, pod, mockDataBackup)
+			handler := &OnceHandler{dataBackup: mockDataBackup}
+			ctx := cruntime.ReconcileRequestContext{
+				NamespacedName: types.NamespacedName{Namespace: "default", Name: "test"},
+				Log:            fake.NullLogger(),
+				Client:         c,
+			}
+			result, err := handler.GetOperationStatus(ctx, &mockDataBackup.Status)
+			Expect(err).NotTo(HaveOccurred())
+			Expect(result.Phase).To(Equal(common.PhaseComplete))
+			Expect(result.NodeAffinity).To(BeNil())
+		})
+
 	})
 })

pkg/dataflow/helper.go

@@ -65,3 +65,47 @@
 
 	return nodeAffinity, nil
 }
+
+// GenerateNodeAffinityFromPod generates a NodeAffinity from a Pod's annotations,
+// using the same annotation-based logic as GenerateNodeAffinity for Jobs.
+// This is used by DataBackup which runs as a Pod rather than a Job.
+func GenerateNodeAffinityFromPod(pod *corev1.Pod) (*corev1.NodeAffinity, error) {
+	if pod == nil {
+		return nil, nil
+	}
+	// not inject, i.e. feature gate not enabled
+	if v := pod.Annotations[common.AnnotationDataFlowAffinityInject]; v != "true" {
+		return nil, nil
+	}
+
+	annotations := pod.Annotations
+
+	nodeAffinity := &corev1.NodeAffinity{
+		RequiredDuringSchedulingIgnoredDuringExecution: &corev1.NodeSelector{
+			NodeSelectorTerms: []corev1.NodeSelectorTerm{
+				{
+					MatchExpressions: nil,
+				},
+			},
+		},
+	}
+
+	hasInjectedLabels := false
+	for key, value := range annotations {
+		if strings.HasPrefix(key, common.AnnotationDataFlowCustomizedAffinityPrefix) {
+			nodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution.NodeSelectorTerms[0].MatchExpressions =
+				append(nodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution.NodeSelectorTerms[0].MatchExpressions,
+					corev1.NodeSelectorRequirement{
+						Key:      strings.TrimPrefix(key, common.AnnotationDataFlowCustomizedAffinityPrefix),
+						Operator: corev1.NodeSelectorOpIn,
+						Values:   []string{value},
+					})
+			hasInjectedLabels = true
+		}
+	}
+	if !hasInjectedLabels {
+		return nil, errors.New("the affinity label is not set, wait for next reconcile")
+	}
+
+	return nodeAffinity, nil
+}

pkg/dataflow/helper_test.go

@@ -128,3 +128,86 @@
 		})
 	}
 }
+
+func TestGenerateNodeAffinityFromPod(t *testing.T) {
+	tests := []struct {
+		name    string
+		pod     *v1.Pod
+		wantNil bool
+		wantErr bool
+		wantLen int
+	}{
+		{
+			name:    "nil pod returns nil",
+			pod:     nil,
+			wantNil: true,
+			wantErr: false,
+		},
+		{
+			name: "pod without inject annotation returns nil",
+			pod: &v1.Pod{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:        "backup-pod",
+					Annotations: map[string]string{},
+				},
+			},
+			wantNil: true,
+			wantErr: false,
+		},
+		{
+			name: "pod with inject annotation but no affinity labels returns error",
+			pod: &v1.Pod{
+				ObjectMeta: metav1.ObjectMeta{
+					Name: "backup-pod",
+					Annotations: map[string]string{
+						common.AnnotationDataFlowAffinityInject: "true",
+					},
+				},
+			},
+			wantNil: true,
+			wantErr: true,
+		},
+		{
+			name: "pod with inject annotation and affinity labels returns NodeAffinity",
+			pod: &v1.Pod{
+				ObjectMeta: metav1.ObjectMeta{
+					Name: "backup-pod",
+					Annotations: map[string]string{
+						common.AnnotationDataFlowAffinityInject:                                        "true",
+						common.AnnotationDataFlowCustomizedAffinityPrefix + common.K8sNodeNameLabelKey: "node01",
+						common.AnnotationDataFlowCustomizedAffinityPrefix + common.K8sZoneLabelKey:     "zone01",
+					},
+				},
+			},
+			wantNil: false,
+			wantErr: false,
+			wantLen: 2,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := GenerateNodeAffinityFromPod(tt.pod)
+			if (err != nil) != tt.wantErr {
+				t.Errorf("GenerateNodeAffinityFromPod() error = %v, wantErr %v", err, tt.wantErr)
+				return
+			}
+			if tt.wantNil {
+				if got != nil {
+					t.Errorf("GenerateNodeAffinityFromPod() expected nil, got %v", got)
+				}
+				return
+			}
+			if got == nil {
+				t.Fatal("GenerateNodeAffinityFromPod() expected non-nil NodeAffinity")
+			}
+			terms := got.RequiredDuringSchedulingIgnoredDuringExecution.NodeSelectorTerms
+			if len(terms) != 1 {
+				t.Fatalf("expected 1 NodeSelectorTerm, got %d", len(terms))
+			}
+			if len(terms[0].MatchExpressions) != tt.wantLen {
+				t.Errorf("expected %d MatchExpressions, got %d", tt.wantLen, len(terms[0].MatchExpressions))
+			}
+		})
+	}
+}