From f7749c1a246e53007e2455e819209a2bd82f3d59 Mon Sep 17 00:00:00 2001
From: Quan Hoang <hdquan@protonmail.com>
Date: Mon, 19 Apr 2021 22:11:35 +0700
Subject: [PATCH] Support affinity in connection pooler deployments

---
 pkg/cluster/connection_pooler.go | 14 +++++++++++---
 1 file changed, 11 insertions(+), 3 deletions(-)

diff --git a/pkg/cluster/connection_pooler.go b/pkg/cluster/connection_pooler.go
index 4e8af610d..4103b8366 100644
--- a/pkg/cluster/connection_pooler.go
+++ b/pkg/cluster/connection_pooler.go
@@ -285,6 +285,8 @@ func (c *Cluster) generateConnectionPoolerPodTemplate(role PostgresRole) (
 		},
 	}
 
+	tolerationsSpec := tolerations(&spec.Tolerations, c.OpConfig.PodToleration)
+
 	podTemplate := &v1.PodTemplateSpec{
 		ObjectMeta: metav1.ObjectMeta{
 			Labels:      c.connectionPoolerLabels(role, true).MatchLabels,
@@ -294,12 +296,18 @@ func (c *Cluster) generateConnectionPoolerPodTemplate(role PostgresRole) (
 		Spec: v1.PodSpec{
 			TerminationGracePeriodSeconds: &gracePeriod,
 			Containers:                    []v1.Container{poolerContainer},
-			// TODO: add tolerations to scheduler pooler on the same node
-			// as database
-			//Tolerations:                   *tolerationsSpec,
+			Tolerations:                   tolerationsSpec,
 		},
 	}
 
+	nodeAffinity := nodeAffinity(c.OpConfig.NodeReadinessLabel, spec.NodeAffinity)
+	if c.OpConfig.EnablePodAntiAffinity {
+		labelsSet := labels.Set(c.connectionPoolerLabels(role, false).MatchLabels)
+		podTemplate.Spec.Affinity = generatePodAffinity(labelsSet, c.OpConfig.PodAntiAffinityTopologyKey, nodeAffinity)
+	} else if nodeAffinity != nil {
+		podTemplate.Spec.Affinity = nodeAffinity
+	}
+
 	return podTemplate, nil
 }