Merge pull request #203 from PasinduDissMrYum/main

Add pod override fields affinity and tolerations
deviceinsight · Aug 14, 2024 · 2fb456e · 2fb456e
2 parents 73a517f + 263498e
commit 2fb456e
Show file tree

Hide file tree

Showing 7 changed files with 57 additions and 7 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -5,6 +5,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## [Unreleased]
+- [#203](https://github.com/deviceinsight/kafkactl/pull/203) Add pod override fields affinity and tolerations
 
 ## 5.2.0 - 2024-08-08
 

diff --git a/README.adoc b/README.adoc
@@ -144,6 +144,25 @@ contexts:
       # optional: nodeSelector to add to the pod
       nodeSelector:
         key: value
+
+      # optional: affinity to add to the pod
+      affinity:
+        # note: other types of affinity also supported
+        nodeAffinity:
+          requiredDuringSchedulingIgnoredDuringExecution:
+            nodeSelectorTerms:
+              - matchExpressions:
+                  - key: "<key>"
+                    operator: "<operator>"
+                    values: [ "<value>" ]
+
+      # optional: tolerations to add to the pod
+      tolerations:
+        - key: "<key>"
+          operator: "<operator>"
+          value: "<value>"
+          effect: "<effect>"
+
     # optional: clientID config (defaults to kafkactl-{username})
     clientID: my-client-id
 
@@ -624,7 +643,7 @@ Producing protobuf message converted from JSON:
 kafkactl produce my-topic --key='{"keyField":123}' --key-proto-type MyKeyMessage --value='{"valueField":"value"}' --value-proto-type MyValueMessage --proto-file kafkamsg.proto
 ----
 
-A more complex protobuf message converted from a multi-line JSON string can be produced using a file input with custom separators. 
+A more complex protobuf message converted from a multi-line JSON string can be produced using a file input with custom separators.
 
 For example, if you have the following protobuf definition (`complex.proto`):
 

diff --git a/internal/common-operation.go b/internal/common-operation.go
@@ -61,6 +61,13 @@ type TLSConfig struct {
 	Insecure bool
 }
 
+type K8sToleration struct {
+	Key      string `json:"key" yaml:"key"`
+	Operator string `json:"operator" yaml:"operator"`
+	Value    string `json:"value" yaml:"value"`
+	Effect   string `json:"effect" yaml:"effect"`
+}
+
 type K8sConfig struct {
 	Enabled         bool
 	Binary          string
@@ -74,6 +81,8 @@ type K8sConfig struct {
 	Labels          map[string]string
 	Annotations     map[string]string
 	NodeSelector    map[string]string
+	Affinity        map[string]any
+	Tolerations     []K8sToleration
 }
 
 type ConsumerConfig struct {
@@ -174,6 +183,11 @@ func CreateClientContext() (ClientContext, error) {
 	context.Kubernetes.Labels = viper.GetStringMapString("contexts." + context.Name + ".kubernetes.labels")
 	context.Kubernetes.Annotations = viper.GetStringMapString("contexts." + context.Name + ".kubernetes.annotations")
 	context.Kubernetes.NodeSelector = viper.GetStringMapString("contexts." + context.Name + ".kubernetes.nodeSelector")
+	context.Kubernetes.Affinity = viper.GetStringMap("contexts." + context.Name + ".kubernetes.affinity")
+
+	if err := viper.UnmarshalKey("contexts."+context.Name+".kubernetes.tolerations", &context.Kubernetes.Tolerations); err != nil {
+		return context, err
+	}
 
 	return context, nil
 }

diff --git a/internal/consumergroupoffsets/consumer-group-offset-operation.go b/internal/consumergroupoffsets/consumer-group-offset-operation.go
@@ -29,7 +29,7 @@ type ConsumerGroupOffsetOperation struct {
 
 func (operation *ConsumerGroupOffsetOperation) ResetConsumerGroupOffset(flags ResetConsumerGroupOffsetFlags, groupName string) error {
 
-	if (flags.Topic == nil || len(flags.Topic) == 0) && (!flags.AllTopics) {
+	if (len(flags.Topic) == 0) && (!flags.AllTopics) {
 		return errors.New("no topic specified")
 	}
 

diff --git a/internal/k8s/executor.go b/internal/k8s/executor.go
@@ -37,6 +37,8 @@ type executor struct {
 	labels          map[string]string
 	annotations     map[string]string
 	nodeSelector    map[string]string
+	affinity        map[string]any
+	tolerations     []internal.K8sToleration
 }
 
 const letterBytes = "abcdefghijklmnpqrstuvwxyz123456789"
@@ -111,6 +113,8 @@ func newExecutor(context internal.ClientContext, runner Runner) *executor {
 		labels:          context.Kubernetes.Labels,
 		annotations:     context.Kubernetes.Annotations,
 		nodeSelector:    context.Kubernetes.NodeSelector,
+		affinity:        context.Kubernetes.Affinity,
+		tolerations:     context.Kubernetes.Tolerations,
 		runner:          runner,
 	}
 }

diff --git a/internal/k8s/pod_overrides.go b/internal/k8s/pod_overrides.go
@@ -1,5 +1,7 @@
 package k8s
 
+import "github.com/deviceinsight/kafkactl/v5/internal"
+
 type imagePullSecretType struct {
 	Name string `json:"name"`
 }
@@ -10,9 +12,11 @@ type metadataType struct {
 }
 
 type specType struct {
-	ImagePullSecrets   []imagePullSecretType `json:"imagePullSecrets,omitempty"`
-	ServiceAccountName *string               `json:"serviceAccountName,omitempty"`
-	NodeSelector       *map[string]string    `json:"nodeSelector,omitempty"`
+	ImagePullSecrets   []imagePullSecretType     `json:"imagePullSecrets,omitempty"`
+	ServiceAccountName *string                   `json:"serviceAccountName,omitempty"`
+	NodeSelector       *map[string]string        `json:"nodeSelector,omitempty"`
+	Affinity           *map[string]any           `json:"affinity,omitempty"`
+	Tolerations        *[]internal.K8sToleration `json:"tolerations,omitempty"`
 }
 
 type PodOverrideType struct {
@@ -29,7 +33,7 @@ func (kubectl *executor) createPodOverride() PodOverrideType {
 	var override PodOverrideType
 	override.APIVersion = "v1"
 
-	if kubectl.serviceAccount != "" || kubectl.imagePullSecret != "" || len(kubectl.nodeSelector) > 0 {
+	if kubectl.serviceAccount != "" || kubectl.imagePullSecret != "" || len(kubectl.nodeSelector) > 0 || len(kubectl.affinity) > 0 || len(kubectl.tolerations) > 0 {
 		override.Spec = &specType{}
 
 		if kubectl.serviceAccount != "" {
@@ -44,6 +48,14 @@ func (kubectl *executor) createPodOverride() PodOverrideType {
 		if len(kubectl.nodeSelector) > 0 {
 			override.Spec.NodeSelector = &kubectl.nodeSelector
 		}
+
+		if len(kubectl.affinity) > 0 {
+			override.Spec.Affinity = &kubectl.affinity
+		}
+
+		if len(kubectl.tolerations) > 0 {
+			override.Spec.Tolerations = &kubectl.tolerations
+		}
 	}
 
 	if len(kubectl.labels) > 0 || len(kubectl.annotations) > 0 {

diff --git a/internal/producer/producer-operation.go b/internal/producer/producer-operation.go
@@ -213,7 +213,7 @@ func (operation *Operation) Produce(topic string, flags Flags) error {
 			}
 
 			if inputMessage, err = inputParser.ParseLine(line); err != nil {
-				return failWithMessageCount(messageCount, err.Error())
+				return failWithMessageCount(messageCount, err.Error()) //nolint:govet
 			}
 
 			messageCount++