fix: use node selector param (#296)

* fix: use node selector param Signed-off-by: chenk <hen.keinan@gmail.com> * fix: use node selector param Signed-off-by: chenk <hen.keinan@gmail.com> --------- Signed-off-by: chenk <hen.keinan@gmail.com>
aquasecurity · Jan 29, 2024 · af2de58 · af2de58
1 parent 9ff67c5
commit af2de58
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 11 deletions.
diff --git a/pkg/jobs/builder.go b/pkg/jobs/builder.go
@@ -17,9 +17,9 @@ func WithTemplate(template string) JobOption {
 	}
 }
 
-func WithNodeSelector(nodeSelector string) JobOption {
+func WithNodeName(nodeName string) JobOption {
 	return func(j *JobBuilder) {
-		j.nodeSelector = nodeSelector
+		j.nodeName = nodeName
 	}
 }
 func WithJobName(name string) JobOption {
@@ -115,6 +115,12 @@ func WithNodeConfiguration(nodeConfig bool) JobOption {
 	}
 }
 
+func WithUseNodeSelectorParam(useNodeSelector bool) JobOption {
+	return func(j *JobBuilder) {
+		j.useNodeSelector = useNodeSelector
+	}
+}
+
 func GetJob(opts ...JobOption) (*batchv1.Job, error) {
 	jb := &JobBuilder{}
 	for _, opt := range opts {
@@ -125,7 +131,7 @@ func GetJob(opts ...JobOption) (*batchv1.Job, error) {
 
 type JobBuilder struct {
 	template             string
-	nodeSelector         string
+	nodeName             string
 	namespace            string
 	imageRef             string
 	serviceAccount       string
@@ -142,6 +148,7 @@ type JobBuilder struct {
 	resourceRequirements *corev1.ResourceRequirements
 	timeout              time.Duration
 	nodeConfig           bool
+	useNodeSelector      bool
 }
 
 func (b *JobBuilder) build() (*batchv1.Job, error) {
@@ -159,12 +166,12 @@ func (b *JobBuilder) build() (*batchv1.Job, error) {
 	if len(b.imageRef) > 0 {
 		job.Spec.Template.Spec.Containers[0].Image = b.imageRef
 	}
-	if len(b.nodeSelector) > 0 && b.nodeConfig {
-		job.Spec.Template.Spec.Containers[0].Args = append(job.Spec.Template.Spec.Containers[0].Args, "--node", b.nodeSelector)
+	if b.nodeConfig {
+		job.Spec.Template.Spec.Containers[0].Args = append(job.Spec.Template.Spec.Containers[0].Args, "--node", b.nodeName)
 	}
-	if b.nodeSelector != "" {
+	if b.useNodeSelector {
 		job.Spec.Template.Spec.NodeSelector = map[string]string{
-			corev1.LabelHostname: b.nodeSelector,
+			corev1.LabelHostname: b.nodeName,
 		}
 	}
 	// append lables

diff --git a/pkg/jobs/collector.go b/pkg/jobs/collector.go
@@ -239,7 +239,7 @@ func (jb *jobCollector) ApplyAndCollect(ctx context.Context, nodeName string) (s
 	JobOptions := []JobOption{
 		WithTemplate(jb.templateName),
 		WithNamespace(jb.namespace),
-		WithNodeSelector(nodeName),
+		WithNodeName(nodeName),
 		WithAnnotation(jb.annotation),
 		WithLabels(jb.labels),
 		WithJobTimeout(jb.collectorTimeout),
@@ -253,6 +253,7 @@ func (jb *jobCollector) ApplyAndCollect(ctx context.Context, nodeName string) (s
 		WithNodeConfiguration(jb.nodeConfig),
 		WithPriorityClassName(jb.priorityClassName),
 		WithResourceRequirements(jb.resourceRequirements),
+		WithUseNodeSelectorParam(true),
 		WithJobName(fmt.Sprintf("%s-%s", jb.templateName, ComputeHash(
 			ObjectRef{
 				Kind:      "Node-Info",
@@ -322,11 +323,11 @@ func (jb *jobCollector) Apply(ctx context.Context, nodeName string) (*batchv1.Jo
 		WithImagePullSecrets(jb.imagePullSecrets),
 		WithContainerVolumeMounts(jb.volumeMounts),
 		WithPriorityClassName(jb.priorityClassName),
+		WithNodeName(nodeName),
 		WithJobName(jb.name),
+		WithUseNodeSelectorParam(jb.useNodeSelector),
 		WithResourceRequirements(jb.resourceRequirements)}
-	if jb.useNodeSelector {
-		jobOptions = append(jobOptions, WithNodeSelector(nodeName))
-	}
+
 	job, err := GetJob(jobOptions...)
 	if err != nil {
 		return nil, fmt.Errorf("running node-collector job: %w", err)