Repository: thelastpickle/cassandra-medusa

Issue:
Modify k8s restores to use an env variable to store the restore mapping
Currently in `restore.py`, we read from a file the restore mapping using a fairly brittle process.
We need to evolve this into reading the mapping from an env variable, using a mapping generated by k8ssandra-operator instead of medusa itself.
We'll rely on the `POD_NAME` env variable to identify the current node and find the source node for the restore, which will be used to set the `POD_IP` env variable and enforce the fqdn for the restore operation.

Comments:

Patch:
diff --git a/medusa/service/grpc/restore.py b/medusa/service/grpc/restore.py
index 0d5ccf6..ea6e79e 100644
--- a/medusa/service/grpc/restore.py
+++ b/medusa/service/grpc/restore.py
@@ -23,7 +23,6 @@ from pathlib import Path
 import medusa.config
 import medusa.restore_node
 import medusa.listing
-from medusa.service.grpc.server import RESTORE_MAPPING_LOCATION
 
 
 def create_config(config_file_path):
@@ -49,43 +48,34 @@ def configure_console_logging(config):
             logging.getLogger(logger_name).setLevel(logging.WARN)
 
 
-if __name__ == '__main__':
-    if len(sys.argv) > 3:
-        config_file_path = sys.argv[2]
-        restore_key = sys.argv[3]
-    else:
-        logging.error("Usage: {} <config_file_path> <restore_key>".format(sys.argv[0]))
-        sys.exit(1)
-
+def apply_mapping_env():
+    # By default we consider that we're restoring in place.
     in_place = True
-    if os.path.exists(f"{RESTORE_MAPPING_LOCATION}/{restore_key}"):
-        logging.info(f"Reading mapping file {RESTORE_MAPPING_LOCATION}/{restore_key}")
-        with open(f"{RESTORE_MAPPING_LOCATION}/{restore_key}", 'r') as f:
-            mapping = json.load(f)
-            # Mapping json structure will look like:
-            # {'in_place': true,
-            #  'host_map':
-            #       {'172.24.0.3': {'source': ['172.24.0.3'], 'seed': False},
-            #        '127.0.0.1': {'source': ['172.24.0.4'], 'seed': False},
-            #        '172.24.0.6': {'source': ['172.24.0.6'], 'seed': False}}}
-            # As each mapping is specific to a Cassandra node, we're looking for the node that maps to 127.0.0.1,
-            # which will be different for each pod.
-            # If hostname resolving is turned on, we're looking for the localhost key instead.
+    if "RESTORE_MAPPING" in os.environ.keys():
+        logging.info("Reading restore mapping from environment variable")
+        mapping = json.loads(os.environ["RESTORE_MAPPING"])
+        # Mapping json structure will look like:
+        # {'in_place': true,
+        #  'host_map':
+        #       {'test-dc1-sts-0': {'source': ['172.24.0.3'], 'seed': False},
+        #        'test-dc1-sts-1': {'source': ['172.24.0.4'], 'seed': False},
+        #        'test-dc1-sts-2': {'source': ['172.24.0.6'], 'seed': False}}}
+        # As each mapping is specific to a Cassandra node, we're looking for
+        # the node that maps to the value of the POD_NAME var.
+        in_place = mapping["in_place"]
+        if not in_place:
             print(f"Mapping: {mapping}")
-            if "localhost" in mapping["host_map"].keys():
-                os.environ["POD_IP"] = mapping["host_map"]["localhost"]["source"][0]
-            elif "127.0.0.1" in mapping["host_map"].keys():
-                os.environ["POD_IP"] = mapping["host_map"]["127.0.0.1"]["source"][0]
-            elif "::1" in mapping["host_map"].keys():
-                os.environ["POD_IP"] = mapping["host_map"]["::1"]["source"][0]
-            in_place = mapping["in_place"]
-            if not in_place and "POD_IP" not in os.environ.keys():
-                print("Could not find target node mapping for this pod while performing remote restore. Exiting.")
-                sys.exit(1)
+            # While POD_IP isn't a great name, it's the env variable that is used to enforce the fqdn of the node.
+            # This allows us to specify which node we're restoring from.
+            if os.environ["POD_NAME"] in mapping["host_map"].keys():
+                os.environ["POD_IP"] = mapping["host_map"][os.environ["POD_NAME"]]["source"][0]
+                print(f"Restoring from {os.environ['POD_IP']}")
+            else:
+                return False, f"POD_NAME {os.environ['POD_NAME']} not found in mapping"
+    return in_place, None
 
-    config = create_config(config_file_path)
-    configure_console_logging(config.logging)
 
+def restore_backup(in_place, config):
     backup_name = os.environ["BACKUP_NAME"]
     tmp_dir = Path("/tmp") if "MEDUSA_TMP_DIR" not in os.environ else Path(os.environ["MEDUSA_TMP_DIR"])
     print(f"Downloading backup {backup_name} to {tmp_dir}")
@@ -98,17 +88,33 @@ if __name__ == '__main__':
 
     cluster_backups = list(medusa.listing.get_backups(config, True))
     logging.info(f"Found {len(cluster_backups)} backups in the cluster")
-    backup_found = False
     # Checking if the backup exists for the node we're restoring.
     # Skipping restore if it doesn't exist.
     for cluster_backup in cluster_backups:
         if cluster_backup.name == backup_name:
-            backup_found = True
             logging.info("Starting restore of backup {}".format(backup_name))
             medusa.restore_node.restore_node(config, tmp_dir, backup_name, in_place, keep_auth,
                                              seeds, verify, keyspaces, tables, use_sstableloader)
-            logging.info("Finished restore of backup {}".format(backup_name))
-            break
+            return f"Finished restore of backup {backup_name}"
+
+    return f"Skipped restore of missing backup {backup_name}"
+
+
+if __name__ == '__main__':
+    if len(sys.argv) > 3:
+        config_file_path = sys.argv[2]
+        restore_key = sys.argv[3]
+    else:
+        logging.error("Usage: {} <config_file_path> <restore_key>".format(sys.argv[0]))
+        sys.exit(1)
+
+    (in_place, error_message) = apply_mapping_env()
+    if error_message:
+        print(error_message)
+        sys.exit(1)
+
+    config = create_config(config_file_path)
+    configure_console_logging(config.logging)
 
-    if not backup_found:
-        logging.info("Skipped restore of missing backup {}".format(backup_name))
+    output_message = restore_backup(in_place, config)
+    logging.info(output_message)
diff --git a/medusa/service/grpc/server.py b/medusa/service/grpc/server.py
index 119a31e..9a456a9 100644
--- a/medusa/service/grpc/server.py
+++ b/medusa/service/grpc/server.py
@@ -43,6 +43,7 @@ TIMESTAMP_FORMAT = '%Y-%m-%d %H:%M:%S'
 BACKUP_MODE_DIFFERENTIAL = "differential"
 BACKUP_MODE_FULL = "full"
 RESTORE_MAPPING_LOCATION = "/var/lib/cassandra/.restore_mapping"
+RESTORE_MAPPING_ENV = "RESTORE_MAPPING"
 
 
 class Server:


Test patch:
diff --git a/tests/service/grpc/restore_test.py b/tests/service/grpc/restore_test.py
new file mode 100644
index 0000000..831dcea
--- /dev/null
+++ b/tests/service/grpc/restore_test.py
@@ -0,0 +1,129 @@
+# -*- coding: utf-8 -*-
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import unittest
+import os
+from unittest.mock import MagicMock, patch
+from pathlib import PosixPath
+
+from medusa.service.grpc.restore import apply_mapping_env, restore_backup
+
+
+class ServiceRestoreTest(unittest.TestCase):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def setUp(self):
+        os.environ.pop('POD_IP', None)
+        os.environ.pop('POD_NAME', None)
+        os.environ.pop('RESTORE_MAPPING', None)
+
+    def test_restore_inplace(self):
+        os.environ['POD_NAME'] = 'test-dc1-sts-0'
+        os.environ['RESTORE_MAPPING'] = '{"in_place": true, "host_map": {' \
+            + '"test-dc1-sts-0": {"source": ["test-dc1-sts-0"], "seed": false},' \
+            + '"test-dc1-sts-1": {"source": ["test-dc1-sts-1"], "seed": false},' \
+            + '"test-dc1-sts-2": {"source": "prod-dc1-sts-2", "seed": false}}}'
+        (in_place, error_message) = apply_mapping_env()
+
+        assert in_place is True
+        assert error_message is None
+        assert "POD_IP" not in os.environ.keys()
+
+    def test_restore_remote(self):
+        os.environ.update({'POD_NAME': 'test-dc1-sts-0'})
+        os.environ['RESTORE_MAPPING'] = '{"in_place": false, "host_map": {' \
+            + '"test-dc1-sts-0": {"source": ["prod-dc1-sts-3"], "seed": false},' \
+            + '"test-dc1-sts-1": {"source": ["prod-dc1-sts-1"], "seed": false},' \
+            + '"test-dc1-sts-2": {"source": "prod-dc1-sts-2", "seed": false}}}'
+        (in_place, error_message) = apply_mapping_env()
+
+        assert in_place is False
+        assert error_message is None
+        assert "POD_IP" in os.environ.keys()
+        assert os.environ['POD_IP'] == 'prod-dc1-sts-3'
+
+    def test_restore_no_match(self):
+        os.environ['POD_NAME'] = 'test-dc1-sts-0'
+        os.environ['RESTORE_MAPPING'] = '{"in_place": false, "host_map": {' \
+            + '"test-dc1-sts-3": {"source": ["prod-dc1-sts-3"], "seed": false},' \
+            + '"test-dc1-sts-1": {"source": ["prod-dc1-sts-1"], "seed": false},' \
+            + '"test-dc1-sts-2": {"source": "prod-dc1-sts-2", "seed": false}}}'
+        (in_place, error_message) = apply_mapping_env()
+
+        assert in_place is False
+        assert error_message is not None
+        assert "POD_IP" not in os.environ.keys()
+
+    def test_success_restore_backup(self):
+        # Define test inputs
+        in_place = True
+        config = {'some': 'config'}
+
+        # Define expected output
+        expected_output = 'Finished restore of backup test_backup'
+
+        # Set up mock environment variables
+        os.environ["BACKUP_NAME"] = "test_backup"
+        os.environ["MEDUSA_TMP_DIR"] = "/tmp"
+
+        # Set up mock for medusa.listing.get_backups()
+        with patch('medusa.listing.get_backups') as mock_get_backups:
+            mock_cluster_backup = MagicMock()
+            mock_cluster_backup.name = "test_backup"
+            mock_get_backups.return_value = [mock_cluster_backup]
+
+            # Set up mock for medusa.restore_node.restore_node()
+            with patch('medusa.restore_node.restore_node') as mock_restore_node:
+                mock_restore_node.return_value = None
+
+                # Call the function
+                result = restore_backup(in_place, config)
+
+                # Assertions
+                assert result == expected_output
+                mock_get_backups.assert_called_once_with(config, True)
+                mock_restore_node.assert_called_once_with(config, PosixPath('/tmp'),
+                                                          'test_backup', True, False, None, False, {}, {}, False)
+
+    def test_fail_restore_backup(self):
+        # Define test inputs
+        in_place = True
+        config = {'some': 'config'}
+
+        # Define expected output
+        expected_output = 'Skipped restore of missing backup test_backup'
+
+        # Set up mock environment variables
+        os.environ["BACKUP_NAME"] = "test_backup"
+        os.environ["MEDUSA_TMP_DIR"] = "/tmp"
+
+        # Set up mock for medusa.listing.get_backups()
+        with patch('medusa.listing.get_backups') as mock_get_backups:
+            mock_cluster_backup = MagicMock()
+            mock_cluster_backup.name = "test_backup10"
+            mock_get_backups.return_value = [mock_cluster_backup]
+
+            # Set up mock for medusa.restore_node.restore_node()
+            with patch('medusa.restore_node.restore_node') as mock_restore_node:
+                mock_restore_node.return_value = None
+
+                result = restore_backup(in_place, config)
+
+                assert result == expected_output
+                mock_get_backups.assert_called_once_with(config, True)
+                mock_restore_node.assert_not_called()
+
+
+if __name__ == '__main__':
+    unittest.main()