From 571e609b27b57482d2720aa8509bb8c2fdd5ef20 Mon Sep 17 00:00:00 2001
From: Joseph Sawaya <jsawaya@redhat.com>
Date: Tue, 31 Aug 2021 12:14:06 -0400
Subject: [PATCH] mgr/rook: host add/rm label in rook orchestrator

This commit adds the functionality for adding/removing labels
to hosts using the rook orchestrator and orch host label add/rm.
The labels are added to the kubernetes node object as kubernetes
labels prefixed by "ceph-label/".

This commit also changes ceph orch host ls to display only the ceph
labels.

Signed-off-by: Joseph Sawaya <jsawaya@redhat.com>
---
 src/pybind/mgr/rook/module.py       |  7 +++++-
 src/pybind/mgr/rook/rook_cluster.py | 35 +++++++++++++++++++++++++++++
 2 files changed, 41 insertions(+), 1 deletion(-)

diff --git a/src/pybind/mgr/rook/module.py b/src/pybind/mgr/rook/module.py
index bf6f673b302d9..6303355877c7f 100644
--- a/src/pybind/mgr/rook/module.py
+++ b/src/pybind/mgr/rook/module.py
@@ -222,7 +222,7 @@ class RookOrchestrator(MgrModule, orchestrator.Orchestrator):
     @handle_orch_error
     def get_hosts(self):
         # type: () -> List[orchestrator.HostSpec]
-        return [orchestrator.HostSpec(n) for n in self.rook_cluster.get_node_names()]
+        return self.rook_cluster.get_hosts()
 
     @handle_orch_error
     def describe_service(self,
@@ -484,6 +484,11 @@ class RookOrchestrator(MgrModule, orchestrator.Orchestrator):
         res = self._rook_cluster.remove_osds(osd_ids, replace, force, self.mon_command)
         return OrchResult(res)
 
+    def add_host_label(self, host: str, label: str) -> OrchResult[str]:
+        return self.rook_cluster.add_host_label(host, label)
+    
+    def remove_host_label(self, host: str, label: str) -> OrchResult[str]:
+        return self.rook_cluster.remove_host_label(host, label)
     """
     @handle_orch_error
     def create_osds(self, drive_group):
diff --git a/src/pybind/mgr/rook/rook_cluster.py b/src/pybind/mgr/rook/rook_cluster.py
index 5a13d39294877..c0386c849d77a 100644
--- a/src/pybind/mgr/rook/rook_cluster.py
+++ b/src/pybind/mgr/rook/rook_cluster.py
@@ -12,6 +12,7 @@ import logging
 from contextlib import contextmanager
 from time import sleep
 import re
+from orchestrator import OrchResult
 
 import jsonpatch
 from urllib.parse import urljoin
@@ -870,6 +871,29 @@ class RookCluster(object):
             cfs.CephFilesystem, 'cephfilesystems', spec.service_id,
             _update_fs, _create_fs)
 
+    def get_matching_node(self, host: str) -> Any:
+        matching_node = None
+        for node in self.nodes.items:
+            if node.metadata.labels['kubernetes.io/hostname'] == host:
+                matching_node = node
+        return matching_node
+
+    def add_host_label(self, host: str, label: str) -> OrchResult[str]:
+        matching_node = self.get_matching_node(host)
+        if matching_node == None:
+            return OrchResult(None, RuntimeError(f"Cannot add {label} label to {host}: host not found in cluster")) 
+        matching_node.metadata.labels['ceph-label/'+ label] = ""
+        self.coreV1_api.patch_node(host, matching_node)
+        return OrchResult(f'Added {label} label to {host}')
+
+    def remove_host_label(self, host: str, label: str) -> OrchResult[str]:
+        matching_node = self.get_matching_node(host)
+        if matching_node == None:
+            return OrchResult(None, RuntimeError(f"Cannot remove {label} label from {host}: host not found in cluster")) 
+        matching_node.metadata.labels.pop('ceph-label/' + label, None)
+        self.coreV1_api.patch_node(host, matching_node)
+        return OrchResult(f'Removed {label} label from {host}')
+
     def apply_objectstore(self, spec: RGWSpec) -> str:
         assert spec.service_id is not None
 
@@ -1040,6 +1064,17 @@ class RookCluster(object):
         )
         return self.remover.remove()
 
+    def get_hosts(self) -> List[orchestrator.HostSpec]:
+        ret = []
+        for node in self.nodes.items:
+            spec = orchestrator.HostSpec(
+                node.metadata.name, 
+                addr='/'.join([addr.address for addr in node.status.addresses]), 
+                labels=[label.split('/')[1] for label in node.metadata.labels if label.startswith('ceph-label')],
+            )
+            ret.append(spec)
+        return ret
+
     def _patch(self, crd: Type, crd_name: str, cr_name: str, func: Callable[[CrdClassT, CrdClassT], CrdClassT]) -> str:
         current_json = self.rook_api_get(
             "{}/{}".format(crd_name, cr_name)
-- 
2.39.5