From 2374dfbd10de3cbb3244daf19ea6e0c9d1889f28 Mon Sep 17 00:00:00 2001
From: Vicente Zepeda Mas <vzepedam@redhat.com>
Date: Fri, 21 Jul 2023 15:29:51 +0200
Subject: [PATCH] Added cluster zones methods

Signed-off-by: Vicente Zepeda Mas <vzepedam@redhat.com>
---
 k8s-commons/factory.go |  5 ++++
 k8s-commons/zones.go   | 52 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 57 insertions(+)
 create mode 100644 k8s-commons/zones.go

diff --git a/k8s-commons/factory.go b/k8s-commons/factory.go
index b87ec32..1451603 100644
--- a/k8s-commons/factory.go
+++ b/k8s-commons/factory.go
@@ -10,8 +10,13 @@ import (
 
 type KClient interface {
 	GetVersion() string
+	// Pods
 	GetPod(ctx context.Context, namespace string, name string, version string) (*apiv1.Pod, error)
 	GetPods(ctx context.Context, namespace string, podParams PodParams) (*apiv1.PodList, error)
+	// Zones
+	GetZones(ctx context.Context) (map[string]int, error)
+	IsClusterMultiAZ(ctx context.Context) (bool, error)
+	GetZoneForNode(ctx context.Context, nodeName string) (string, error)
 }
 
 func NewKClient(configOverrides clientcmd.ConfigOverrides) (KClient, error) {
diff --git a/k8s-commons/zones.go b/k8s-commons/zones.go
new file mode 100644
index 0000000..257539d
--- /dev/null
+++ b/k8s-commons/zones.go
@@ -0,0 +1,52 @@
+package k8scommons
+
+import (
+	"context"
+	"fmt"
+
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+)
+
+// GetZone returns a map with all the zones an the number of
+// worker nodes on each one
+func (c *K8SClient) GetZones(ctx context.Context) (map[string]int, error) {
+	return c.getZones(ctx)
+}
+
+// IsClusterMultiAZ returns true if the cluster has more than
+// one zone configured in their worker nodes
+func (c *K8SClient) IsClusterMultiAZ(ctx context.Context) (bool, error) {
+	zones, err := c.getZones(ctx)
+	if err != nil {
+		return false, err
+	}
+	return len(zones) > 1, nil
+}
+
+func (c *K8SClient) getZones(ctx context.Context) (map[string]int, error) {
+	zones := make(map[string]int)
+
+	n, err := c.client.CoreV1().Nodes().List(context.TODO(), metav1.ListOptions{LabelSelector: "node-role.kubernetes.io/worker="})
+	if err != nil {
+		return zones, err
+	}
+
+	for _, l := range n.Items {
+		if len(l.GetLabels()["topology.kubernetes.io/zone"]) < 1 {
+			return zones, fmt.Errorf("no zone label on Node: %s", l.Name)
+		}
+		zones[l.GetLabels()["topology.kubernetes.io/zone"]]++
+	}
+	return zones, nil
+}
+
+func (c *K8SClient) GetZoneForNode(ctx context.Context, nodeName string) (string, error) {
+	n, err := c.client.CoreV1().Nodes().Get(ctx, nodeName, metav1.GetOptions{})
+	if err != nil {
+		return "", err
+	}
+	if len(n.GetLabels()["topology.kubernetes.io/zone"]) < 1 {
+		return "", fmt.Errorf("no zone label on Node: %s", nodeName)
+	}
+	return n.GetLabels()["topology.kubernetes.io/zone"], nil
+}