From 474359e3c0efb9938d69d28604355079676a7002 Mon Sep 17 00:00:00 2001
From: Anthony Mahanna <43019056+aMahanna@users.noreply.github.com>
Date: Fri, 2 Feb 2024 15:42:21 -0500
Subject: [PATCH] new: `get/put /_admin/cluster/maintenance/` (#306)

* new: `server_maintenance_mode` and `toggle_server_maintenance_mode`

* fix lint

* increase sleep

* increase sleep (again)

* revert 8e71d82

* finding a dbserver

* Changing timeouts

---------

Co-authored-by: Alex Petenchea <alex.petenchea@gmail.com>
---
 arango/cluster.py     | 52 +++++++++++++++++++++++++++++++++++++++++++
 docs/cluster.rst      |  7 +++++-
 tests/test_cluster.py | 32 ++++++++++++++++++++++++++
 3 files changed, 90 insertions(+), 1 deletion(-)

diff --git a/arango/cluster.py b/arango/cluster.py
index 6ebdb69b..dffaeb04 100644
--- a/arango/cluster.py
+++ b/arango/cluster.py
@@ -162,6 +162,58 @@ def response_handler(resp: Response) -> Json:
 
         return self._execute(request, response_handler)
 
+    def server_maintenance_mode(self, server_id: str) -> Result[Json]:
+        """Return the maintenance status for the given server.
+
+        :param server_id: Server ID.
+        :type server_id: str
+        :return: Maintenance status for the given server.
+        :rtype: dict
+        :raise arango.exceptions.ClusterMaintenanceModeError: If retrieval fails.
+        """
+        request = Request(
+            method="get",
+            endpoint=f"/_admin/cluster/maintenance/{server_id}",
+        )
+
+        def response_handler(resp: Response) -> Json:
+            if resp.is_success:
+                result: Json = resp.body.get("result", {})
+                return result
+
+            raise ClusterMaintenanceModeError(resp, request)
+
+        return self._execute(request, response_handler)
+
+    def toggle_server_maintenance_mode(
+        self, server_id: str, mode: str, timeout: Optional[int] = None
+    ) -> Result[Json]:
+        """Enable or disable the maintenance mode for the given server.
+
+        :param server_id: Server ID.
+        :type server_id: str
+        :param mode: Maintenance mode. Allowed values are "normal" and "maintenance".
+        :type mode: str
+        :param timeout: Timeout in seconds.
+        :type timeout: Optional[int]
+        :return: Result of the operation.
+        :rtype: dict
+        :raise arango.exceptions.ClusterMaintenanceModeError: If toggle fails.
+        """
+        request = Request(
+            method="put",
+            endpoint=f"/_admin/cluster/maintenance/{server_id}",
+            data={"mode": mode, "timeout": timeout},
+        )
+
+        def response_handler(resp: Response) -> Json:
+            if resp.is_success:
+                return format_body(resp.body)
+
+            raise ClusterMaintenanceModeError(resp, request)
+
+        return self._execute(request, response_handler)
+
     def health(self) -> Result[Json]:
         """Return the cluster health.
 
diff --git a/docs/cluster.rst b/docs/cluster.rst
index fbb3bb5e..fdb45bca 100644
--- a/docs/cluster.rst
+++ b/docs/cluster.rst
@@ -86,8 +86,13 @@ Below is an example on how to manage clusters using python-arango.
     cluster.server_engine(server_id)
     cluster.server_version(server_id)
     cluster.server_statistics(server_id)
+    cluster.server_maintenance_mode(server_id)
 
-    # Toggle maintenance mode (allowed values are "on" and "off").
+    # Toggle Server maintenance mode (allowed values are "normal" and "maintenance").
+    cluster.toggle_server_maintenance_mode(server_id, 'normal')
+    cluster.toggle_server_maintenance_mode(server_id, 'maintenance', timeout=30)
+
+    # Toggle Cluster maintenance mode (allowed values are "on" and "off").
     cluster.toggle_maintenance_mode('on')
     cluster.toggle_maintenance_mode('off')
 
diff --git a/tests/test_cluster.py b/tests/test_cluster.py
index 3eda0bb0..39c2b8cf 100644
--- a/tests/test_cluster.py
+++ b/tests/test_cluster.py
@@ -1,3 +1,4 @@
+import time
 import warnings
 
 import pytest
@@ -112,6 +113,37 @@ def test_cluster_server_statistics(sys_db, bad_db, cluster):
     assert err.value.error_code in {FORBIDDEN, DATABASE_NOT_FOUND}
 
 
+def test_cluster_server_maintenance_mode(sys_db, bad_db, cluster):
+    if not cluster:
+        pytest.skip("Only tested in a cluster setup")
+
+    # Must be a DBServer
+    health = sys_db.cluster.health()
+    server_id = None
+    for server_id, info in health["Health"].items():
+        if info["Role"] == "DBServer":
+            server_id = server_id
+            break
+    if server_id is None:
+        pytest.skip("No DBServer found in cluster")
+
+    result = sys_db.cluster.server_maintenance_mode(server_id)
+    assert result == {}
+
+    with assert_raises(ClusterMaintenanceModeError) as err:
+        bad_db.cluster.server_maintenance_mode(server_id)
+    assert err.value.error_code in {FORBIDDEN, DATABASE_NOT_FOUND}
+
+    sys_db.cluster.toggle_server_maintenance_mode(server_id, "maintenance", timeout=2)
+    result = sys_db.cluster.server_maintenance_mode(server_id)
+    assert "Mode" in result
+    assert "Until" in result
+
+    time.sleep(5)
+    result = sys_db.cluster.server_maintenance_mode(server_id)
+    assert result == {}
+
+
 def test_cluster_toggle_maintenance_mode(sys_db, bad_db, cluster):
     if not cluster:
         pytest.skip("Only tested in a cluster setup")