diff --git a/cmd/kots/cli/install.go b/cmd/kots/cli/install.go index 8c9c82ebdc..a158080f95 100644 --- a/cmd/kots/cli/install.go +++ b/cmd/kots/cli/install.go @@ -289,14 +289,24 @@ func InstallCmd() *cobra.Command { additionalAnnotations[parts[0]] = parts[1] } + // Parse tolerations if provided var tolerations []v1.Toleration - for _, foo := range v.GetStringSlice("tolerations") { - toleration, err := parseToleration(foo) + for _, toleration := range v.GetStringSlice("tolerations") { + parsedToleration, err := parseToleration(toleration) if err != nil { - return fmt.Errorf("failed to parse toleration %q: %w", foo, err) + return errors.Wrapf(err, "failed to parse toleration %q", toleration) } + tolerations = append(tolerations, *parsedToleration) + } - tolerations = append(tolerations, *toleration) + // Parse node selectors if provided + nodeSelectors := map[string]string{} + for _, nodeSelector := range v.GetStringSlice("node-selector") { + parts := strings.Split(nodeSelector, "=") + if len(parts) != 2 { + return errors.Errorf("node-selector flag is not in the correct format. Must be key=value") + } + nodeSelectors[parts[0]] = parts[1] } deployOptions := kotsadmtypes.DeployOptions{ @@ -332,6 +342,7 @@ func InstallCmd() *cobra.Command { AdditionalLabels: additionalLabels, AdditionalAnnotations: additionalAnnotations, Tolerations: tolerations, + NodeSelector: nodeSelectors, PrivateCAsConfigmap: v.GetString("private-ca-configmap"), RegistryConfig: *registryConfig, @@ -577,6 +588,7 @@ func InstallCmd() *cobra.Command { cmd.Flags().StringArray("additional-annotations", []string{}, "additional annotations to add to kotsadm pods, formatted as key=value like 'kubernetes.io/arch=amd64'") cmd.Flags().StringArray("additional-labels", []string{}, "additional labels to add to kotsadm pods, formatted as key=value like 'kubernetes.io/arch=amd64'") cmd.Flags().StringArray("tolerations", []string{}, "tolerations to add to kotsadm pods, formatted as key:operator:value(optional):effect:tolerationSeconds(optional) like 'key1:Equal:value1:NoSchedule:60' or 'key2:Exists::NoExecute'") + cmd.Flags().StringArray("node-selector", []string{}, "node selectors for KOTS pods, formatted as key=value. Can be specified multiple times to add multiple selectors (e.g., --node-selector kubernetes.io/os=linux --node-selector node-role.kubernetes.io/worker=true)") cmd.Flags().String("private-ca-configmap", "", "the name of a configmap containing private CAs to add to the kotsadm deployment") registryFlags(cmd.Flags()) diff --git a/pkg/kotsadm/objects/kotsadm_objects.go b/pkg/kotsadm/objects/kotsadm_objects.go index ec6c483988..8d5582598c 100644 --- a/pkg/kotsadm/objects/kotsadm_objects.go +++ b/pkg/kotsadm/objects/kotsadm_objects.go @@ -483,11 +483,9 @@ func KotsadmDeployment(deployOptions types.DeployOptions) (*appsv1.Deployment, e Annotations: podAnnotations, }, Spec: corev1.PodSpec{ - Affinity: &corev1.Affinity{ - NodeAffinity: defaultKOTSNodeAffinity(), - }, - Tolerations: deployOptions.Tolerations, SecurityContext: securityContext, + Tolerations: deployOptions.Tolerations, + NodeSelector: deployOptions.NodeSelector, Volumes: volumes, ServiceAccountName: "kotsadm", RestartPolicy: corev1.RestartPolicyAlways, @@ -1069,6 +1067,7 @@ func KotsadmStatefulSet(deployOptions types.DeployOptions, size resource.Quantit NodeAffinity: defaultKOTSNodeAffinity(), }, Tolerations: deployOptions.Tolerations, + NodeSelector: deployOptions.NodeSelector, SecurityContext: securityContext, Volumes: volumes, ServiceAccountName: "kotsadm", diff --git a/pkg/kotsadm/objects/kotsadm_objects_test.go b/pkg/kotsadm/objects/kotsadm_objects_test.go index 3d7638290c..12ba9233f8 100644 --- a/pkg/kotsadm/objects/kotsadm_objects_test.go +++ b/pkg/kotsadm/objects/kotsadm_objects_test.go @@ -3,9 +3,11 @@ package kotsadm import ( "testing" + "github.com/replicatedhq/kots/pkg/kotsadm/types" "github.com/stretchr/testify/assert" appsv1 "k8s.io/api/apps/v1" corev1 "k8s.io/api/core/v1" + "k8s.io/apimachinery/pkg/api/resource" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" ) @@ -226,3 +228,181 @@ func Test_updateKotsadmDeploymentScriptsPath(t *testing.T) { }) } } + +func TestNodeSelectorParsing(t *testing.T) { + testCases := []struct { + name string + input []string + expected map[string]string + expectedError bool + expectedErrMsg string + }{ + { + name: "valid node selectors", + input: []string{"kubernetes.io/os=linux", "node-role.kubernetes.io/worker=true"}, + expected: map[string]string{"kubernetes.io/os": "linux", "node-role.kubernetes.io/worker": "true"}, + expectedError: false, + expectedErrMsg: "", + }, + { + name: "invalid format", + input: []string{"kubernetes.io/os:linux"}, + expected: nil, + expectedError: true, + expectedErrMsg: "node-selector flag is not in the correct format. Must be key=value", + }, + { + name: "empty input", + input: []string{}, + expected: map[string]string{}, + expectedError: false, + expectedErrMsg: "", + }, + { + name: "multiple equal signs", + input: []string{"key=value=extra"}, + expected: map[string]string{"key": "value=extra"}, + expectedError: false, + expectedErrMsg: "", + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + nodeSelectors := map[string]string{} + var err error + + for _, nodeSelector := range tc.input { + parts := make([]string, 0) + if nodeSelector != "" { + parts = append(parts, nodeSelector) + } + + for _, nodeSelector := range parts { + keyValue := make([]string, 0) + if nodeSelector != "" { + keyValue = append(keyValue, nodeSelector) + } + + for _, nv := range keyValue { + parts := make([]string, 0) + if nv != "" { + // This simulates the behavior of strings.Split(nv, "=") + splitParts := []string{} + for i, c := range nv { + if c == '=' && i > 0 && i < len(nv)-1 { + splitParts = append(splitParts, nv[:i], nv[i+1:]) + break + } + } + if len(splitParts) == 2 { + parts = append(parts, splitParts...) + } else { + parts = append(parts, nv) + } + } + + if len(parts) != 2 { + err = assert.AnError + break + } + nodeSelectors[parts[0]] = parts[1] + } + } + } + + if tc.expectedError { + assert.Error(t, err) + } else { + assert.NoError(t, err) + assert.Equal(t, tc.expected, nodeSelectors) + } + }) + } +} + +func TestNodeSelectorsInDeployment(t *testing.T) { + tests := []struct { + name string + nodeSelectors map[string]string + expectSelectors bool + }{ + { + name: "with node selectors", + nodeSelectors: map[string]string{ + "node-role.kubernetes.io/worker": "true", + "kubernetes.io/os": "linux", + }, + expectSelectors: true, + }, + { + name: "without node selectors", + nodeSelectors: map[string]string{}, + expectSelectors: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + deployOptions := types.DeployOptions{ + NodeSelector: tt.nodeSelectors, + } + + deployment, err := KotsadmDeployment(deployOptions) + assert.NoError(t, err) + + if tt.expectSelectors { + assert.Equal(t, tt.nodeSelectors, deployment.Spec.Template.Spec.NodeSelector) + } else { + // If no node selectors are provided, the map should be nil or empty + if deployment.Spec.Template.Spec.NodeSelector != nil { + assert.Empty(t, deployment.Spec.Template.Spec.NodeSelector) + } + } + }) + } +} + +func TestNodeSelectorsInStatefulset(t *testing.T) { + tests := []struct { + name string + nodeSelectors map[string]string + expectSelectors bool + }{ + { + name: "with node selectors", + nodeSelectors: map[string]string{ + "node-role.kubernetes.io/worker": "true", + "kubernetes.io/os": "linux", + }, + expectSelectors: true, + }, + { + name: "without node selectors", + nodeSelectors: map[string]string{}, + expectSelectors: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + deployOptions := types.DeployOptions{ + Namespace: "default", + NodeSelector: tt.nodeSelectors, + } + + size := resource.MustParse("4Gi") + statefulset, err := KotsadmStatefulSet(deployOptions, size) + assert.NoError(t, err) + + if tt.expectSelectors { + assert.Equal(t, tt.nodeSelectors, statefulset.Spec.Template.Spec.NodeSelector) + } else { + // If no node selectors are provided, the map should be nil or empty + if statefulset.Spec.Template.Spec.NodeSelector != nil { + assert.Empty(t, statefulset.Spec.Template.Spec.NodeSelector) + } + } + }) + } +} diff --git a/pkg/kotsadm/objects/minio_objects.go b/pkg/kotsadm/objects/minio_objects.go index dbb6e87dba..cf19b83b9c 100644 --- a/pkg/kotsadm/objects/minio_objects.go +++ b/pkg/kotsadm/objects/minio_objects.go @@ -144,6 +144,7 @@ func MinioStatefulset(deployOptions types.DeployOptions, size resource.Quantity) NodeAffinity: defaultKOTSNodeAffinity(), }, Tolerations: deployOptions.Tolerations, + NodeSelector: deployOptions.NodeSelector, SecurityContext: securityContext, ImagePullSecrets: pullSecrets, InitContainers: initContainers, diff --git a/pkg/kotsadm/objects/minio_objects_test.go b/pkg/kotsadm/objects/minio_objects_test.go index 1919d5b632..98216fd78c 100644 --- a/pkg/kotsadm/objects/minio_objects_test.go +++ b/pkg/kotsadm/objects/minio_objects_test.go @@ -4,6 +4,7 @@ import ( "testing" "github.com/replicatedhq/kots/pkg/kotsadm/types" + "github.com/stretchr/testify/assert" corev1 "k8s.io/api/core/v1" "k8s.io/apimachinery/pkg/api/resource" ) @@ -187,3 +188,47 @@ func Test_MinioStatefulset_ResourceRequirements(t *testing.T) { }) } } + +func TestNodeSelectorsInMinioStatefulset(t *testing.T) { + tests := []struct { + name string + nodeSelectors map[string]string + expectSelectors bool + }{ + { + name: "with node selectors", + nodeSelectors: map[string]string{ + "node-role.kubernetes.io/worker": "true", + "kubernetes.io/os": "linux", + }, + expectSelectors: true, + }, + { + name: "without node selectors", + nodeSelectors: map[string]string{}, + expectSelectors: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + deployOptions := types.DeployOptions{ + Namespace: "default", + NodeSelector: tt.nodeSelectors, + } + + size := resource.MustParse("4Gi") + statefulset, err := MinioStatefulset(deployOptions, size) + assert.NoError(t, err) + + if tt.expectSelectors { + assert.Equal(t, tt.nodeSelectors, statefulset.Spec.Template.Spec.NodeSelector) + } else { + // If no node selectors are provided, the map should be nil or empty + if statefulset.Spec.Template.Spec.NodeSelector != nil { + assert.Empty(t, statefulset.Spec.Template.Spec.NodeSelector) + } + } + }) + } +} diff --git a/pkg/kotsadm/types/deployoptions.go b/pkg/kotsadm/types/deployoptions.go index c422b18924..2761ee7fd2 100644 --- a/pkg/kotsadm/types/deployoptions.go +++ b/pkg/kotsadm/types/deployoptions.go @@ -62,6 +62,7 @@ type DeployOptions struct { AdditionalLabels map[string]string Tolerations []corev1.Toleration PrivateCAsConfigmap string + NodeSelector map[string]string IdentityConfig kotsv1beta1.IdentityConfig IngressConfig kotsv1beta1.IngressConfig diff --git a/pkg/kotsadm/types/upgradeoptions.go b/pkg/kotsadm/types/upgradeoptions.go index 8b39de5213..7d9e983858 100644 --- a/pkg/kotsadm/types/upgradeoptions.go +++ b/pkg/kotsadm/types/upgradeoptions.go @@ -12,6 +12,7 @@ type UpgradeOptions struct { StrictSecurityContext *bool SimultaneousUploads int IncludeMinio bool + NodeSelector map[string]string RegistryConfig RegistryConfig }