From 25e140a17127e6cb7558b4b934944a6157a9de2d Mon Sep 17 00:00:00 2001
From: David Boreham <david@bozemanpass.com>
Date: Sun, 11 Aug 2024 20:52:17 -0600
Subject: [PATCH] Support for pod/node affinity and tolerations

---
 stack_orchestrator/constants.py               |  2 +
 stack_orchestrator/deploy/k8s/cluster_info.py | 44 ++++++++++++++++++-
 stack_orchestrator/deploy/spec.py             |  6 +++
 3 files changed, 51 insertions(+), 1 deletion(-)

diff --git a/stack_orchestrator/constants.py b/stack_orchestrator/constants.py
index aee36ad8..07fc68f4 100644
--- a/stack_orchestrator/constants.py
+++ b/stack_orchestrator/constants.py
@@ -35,5 +35,7 @@ security_key = "security"
 annotations_key = "annotations"
 labels_key = "labels"
 replicas_key = "replicas"
+node_affinities_key = "node-affinities"
+node_tolerations_key = "node-tolerations"
 kind_config_filename = "kind-config.yml"
 kube_config_filename = "kubeconfig.yml"
diff --git a/stack_orchestrator/deploy/k8s/cluster_info.py b/stack_orchestrator/deploy/k8s/cluster_info.py
index 05443f72..7342816e 100644
--- a/stack_orchestrator/deploy/k8s/cluster_info.py
+++ b/stack_orchestrator/deploy/k8s/cluster_info.py
@@ -365,6 +365,8 @@ class ClusterInfo:
 
         annotations = None
         labels = {"app": self.app_name}
+        affinity = None
+        tolerations = None
 
         if self.spec.get_annotations():
             annotations = {}
@@ -377,12 +379,52 @@ class ClusterInfo:
                 for service_name in services:
                     labels[key.replace("{name}", service_name)] = value
 
+        if self.spec.get_node_affinities():
+            affinities = []
+            for rule in self.spec.get_node_affinities():
+                # TODO add some input validation here
+                label_name = rule['label']
+                label_value = rule['value']
+                affinities.append(client.V1NodeSelectorTerm(
+                            match_expressions=[client.V1NodeSelectorRequirement(
+                                key=label_name,
+                                operator="In",
+                                values=[label_value]
+                            )]
+                        )
+                    )
+            affinity = client.V1Affinity(
+                node_affinity=client.V1NodeAffinity(
+                    required_during_scheduling_ignored_during_execution=client.V1NodeSelector(
+                        node_selector_terms=affinities
+                    ))
+                )
+
+        if self.spec.get_node_tolerations():
+            tolerations = []
+            for toleration in self.spec.get_node_tolerations():
+                # TODO add some input validation here
+                toleration_key = toleration['key']
+                toleration_value = toleration['value']
+                tolerations.append(client.V1Toleration(
+                    effect="NoSchedule",
+                    key=toleration_key,
+                    operator="Equal",
+                    value=toleration_value
+                ))
+
         template = client.V1PodTemplateSpec(
             metadata=client.V1ObjectMeta(
                 annotations=annotations,
                 labels=labels
             ),
-            spec=client.V1PodSpec(containers=containers, image_pull_secrets=image_pull_secrets, volumes=volumes),
+            spec=client.V1PodSpec(
+                containers=containers,
+                image_pull_secrets=image_pull_secrets,
+                volumes=volumes,
+                affinity=affinity,
+                tolerations=tolerations
+                ),
         )
         spec = client.V1DeploymentSpec(
             replicas=self.spec.get_replicas(),
diff --git a/stack_orchestrator/deploy/spec.py b/stack_orchestrator/deploy/spec.py
index e8d293e3..99dfb16f 100644
--- a/stack_orchestrator/deploy/spec.py
+++ b/stack_orchestrator/deploy/spec.py
@@ -120,6 +120,12 @@ class Spec:
     def get_replicas(self):
         return self.obj.get(constants.replicas_key, 1)
 
+    def get_node_affinities(self):
+        return self.obj.get(constants.node_affinities_key, [])
+
+    def get_node_tolerations(self):
+        return self.obj.get(constants.node_tolerations_key, [])
+
     def get_labels(self):
         return self.obj.get(constants.labels_key, {})