Add More Deploy Options

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())

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",

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)
+				}
+			}
+		})
+	}
+}

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,

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)
+				}
+			}
+		})
+	}
+}

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

pkg/kotsadm/types/upgradeoptions.go

@@ -12,6 +12,7 @@ type UpgradeOptions struct {
 	StrictSecurityContext *bool
 	SimultaneousUploads   int
 	IncludeMinio          bool
+	NodeSelector          map[string]string
 
 	RegistryConfig RegistryConfig
 }