Fix Spec.Drain.PodSelector

Use helpers instead of string ops to handle pod selector generation Signed-off-by: Brad Davidson <[email protected]>
rancher · May 4, 2023 · f1b085e · f1b085e
1 parent b95d7cd
commit f1b085e
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 5 deletions.
diff --git a/e2e/suite/job_generate_test.go b/e2e/suite/job_generate_test.go
@@ -114,7 +114,8 @@ var _ = Describe("Job Generation", func() {
 			Expect(jobs[0].Status.Succeeded).To(BeNumerically("==", 1))
 			Expect(jobs[0].Status.Failed).To(BeNumerically("==", 0))
 			Expect(jobs[0].Spec.Template.Spec.InitContainers).To(HaveLen(1))
-			Expect(jobs[0].Spec.Template.Spec.InitContainers[0].Args).To(ContainSubstring("csi-attacher"))
+			Expect(jobs[0].Spec.Template.Spec.InitContainers[0].Args).To(ContainSubstring("!upgrade.cattle.io/controller"))
+			Expect(jobs[0].Spec.Template.Spec.InitContainers[0].Args).To(ContainSubstring("app notin (csi-attacher,csi-provisioner)"))
 		})
 	})
 })
diff --git a/pkg/upgrade/job/job.go b/pkg/upgrade/job/job.go
@@ -17,6 +17,7 @@ import (
 	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/labels"
+	"k8s.io/apimachinery/pkg/selection"
 )
 
 const (
@@ -240,12 +241,22 @@ func New(plan *upgradeapiv1.Plan, node *corev1.Node, controllerName string) *bat
 	// then we cordon/drain
 	cordon, drain := plan.Spec.Cordon, plan.Spec.Drain
 	if drain != nil {
-		podSelector := `!` + upgradeapi.LabelController
+		controllerRequirement, _ := labels.NewRequirement(upgradeapi.LabelController, selection.DoesNotExist, nil)
+		podSelector := labels.NewSelector().Add(*controllerRequirement)
+
 		if drain.PodSelector != nil {
-			podSelector = podSelector + `,` + drain.PodSelector.String()
+			if selector, err := metav1.LabelSelectorAsSelector(drain.PodSelector); err != nil {
+				logrus.Warnf("failed to convert Spec.Drain.PodSelector to selector: %v", err)
+			} else {
+				if requirements, ok := selector.Requirements(); !ok {
+					logrus.Warnf("Spec.Drain.PodSelector requirements are not selectable")
+				} else {
+					podSelector = podSelector.Add(requirements...)
+				}
+			}
 		}
 
-		args := []string{"drain", node.Name, "--pod-selector", podSelector}
+		args := []string{"drain", node.Name, "--pod-selector", podSelector.String()}
 		if drain.IgnoreDaemonSets == nil || *plan.Spec.Drain.IgnoreDaemonSets {
 			args = append(args, "--ignore-daemonsets")
 		}

diff --git a/pkg/upgrade/plan/plan.go b/pkg/upgrade/plan/plan.go
@@ -31,7 +31,8 @@ const (
 )
 
 var (
-	ErrDrainDeleteConflict = fmt.Errorf("spec.drain cannot specify both deleteEmptydirData and deleteLocalData")
+	ErrDrainDeleteConflict           = fmt.Errorf("spec.drain cannot specify both deleteEmptydirData and deleteLocalData")
+	ErrDrainPodSelectorNotSelectable = fmt.Errorf("spec.drain.podSelector is not selectable")
 
 	PollingInterval = func(defaultValue time.Duration) time.Duration {
 		if str, ok := os.LookupEnv("SYSTEM_UPGRADE_PLAN_POLLING_INTERVAL"); ok {
@@ -237,6 +238,15 @@ func Validate(plan *upgradeapiv1.Plan) error {
 		if drainSpec.DeleteEmptydirData != nil && drainSpec.DeleteLocalData != nil {
 			return ErrDrainDeleteConflict
 		}
+		if drainSpec.PodSelector != nil {
+			selector, err := metav1.LabelSelectorAsSelector(drainSpec.PodSelector)
+			if err != nil {
+				return err
+			}
+			if _, ok := selector.Requirements(); !ok {
+				return ErrDrainPodSelectorNotSelectable
+			}
+		}
 	}
 	return nil
 }