[DPE-2648] POC Remove shards

lib/charms/mongodb/v0/mongos.py

@@ -4,11 +4,11 @@
 
 import logging
 from dataclasses import dataclass
-from typing import Optional, Set
+from typing import Dict, Optional, Set
 from urllib.parse import quote_plus
 
+from charms.mongodb.v0.mongodb import NotReadyError
 from pymongo import MongoClient
-from pymongo.errors import PyMongoError
 
 from config import Config
 
@@ -65,8 +65,16 @@ def uri(self):
         )
 
 
-class NotReadyError(PyMongoError):
-    """Raised when not all replica set members healthy or finished initial sync."""
+class RemovePrimaryShardError(Exception):
+    """Raised when there is an attempt to remove the primary shard."""
+
+
+class ShardNotInClusterError(Exception):
+    """Raised when shard is not present in cluster, but it is expected to be."""
+
+
+class ShardNotPlannedForRemovalError(Exception):
+    """Raised when it is expected that a shard is planned for removal, but it is not."""
 
 
 class MongosConnection:
@@ -158,6 +166,75 @@ def add_shard(self, shard_name, shard_hosts, shard_port=Config.MONGODB_PORT):
         logger.info("Adding shard %s", shard_name)
         self.client.admin.command("addShard", shard_url)
 
+    def remove_shard(self, shard_name: str) -> None:
+        """Removes shard from the cluster.
+
+        Raises:
+            ConfigurationError, OperationFailure, NotReadyError,
+            RemovePrimaryShardError
+        """
+        sc_status = self.client.admin.command("listShards")
+        # It is necessary to call removeShard multiple times on a shard to guarantee removal.
+        # Allow re-removal of shards that are currently draining.
+        if self._is_any_draining(sc_status, ignore_shard=shard_name):
+            cannot_remove_shard = (
+                f"cannot remove shard {shard_name} from cluster, another shard is draining"
+            )
+            logger.info(cannot_remove_shard)
+            raise NotReadyError(cannot_remove_shard)
+
+        # TODO Follow up PR, there is no MongoDB command to retrieve primary shard, this is
+        # possible with mongosh.
+        primary_shard = False
+        if primary_shard:
+            # TODO Future PR, support removing Primary Shard if there are no unsharded collections
+            # on it. All sharded collections should perform `MovePrimary`
+            cannot_remove_primary_shard = (
+                f"Shard {shard_name} is the primary shard, cannot remove."
+            )
+            logger.error(cannot_remove_primary_shard)
+            raise RemovePrimaryShardError(cannot_remove_primary_shard)
+
+        logger.info("Attempting to remove shard %s", shard_name)
+        self.client.admin.command("removeShard", shard_name)
+        logger.info("Shard %s, now draining", shard_name)
+
+    def _is_shard_draining(self, shard_name: str) -> bool:
+        """Reports if a given shard is currently in the draining state.
+
+        Raises:
+            ConfigurationError, OperationFailure, ShardNotInClusterError,
+            ShardNotPlannedForRemovalError
+        """
+        sc_status = self.client.admin.command("listShards")
+        for shard in sc_status["shards"]:
+            if shard["_id"] == shard_name:
+                if "draining" not in shard:
+                    raise ShardNotPlannedForRemovalError(
+                        f"Shard {shard_name} has not been marked for removal",
+                    )
+                return shard["draining"]
+
+        raise ShardNotInClusterError(
+            f"Shard {shard_name} not in cluster, could not retrieve draining status"
+        )
+
+    @staticmethod
+    def _is_any_draining(sc_status: Dict, ignore_shard: str = "") -> bool:
+        """Returns true if any shard members is draining.
+
+        Checks if any members in sharded cluster are draining data.
+
+        Args:
+            sc_status: current state of shard cluster status as reported by mongos.
+            ignore_shard: shard to ignore
+        """
+        return any(
+            # check draining status of all shards except the one to be ignored.
+            shard.get("draining", False) if shard["_id"] != ignore_shard else False
+            for shard in sc_status["shards"]
+        )
+
     @staticmethod
     def _hostname_from_hostport(hostname: str) -> str:
         """Return hostname part from MongoDB returned.