[DPE-5372] Add safeguard hooks for upgrades (#327)

Author: Gu1nness

src/charm.py

@@ -43,6 +43,7 @@
     CharmBase,
     ConfigChangedEvent,
     RelationDepartedEvent,
+    RelationEvent,
     StartEvent,
     UpdateStatusEvent,
 )
@@ -93,6 +94,7 @@ def __init__(self, *args):
         super().__init__(*args)
 
         self.framework.observe(self.on.mongod_pebble_ready, self._on_mongod_pebble_ready)
+        self.framework.observe(self.on.config_changed, self._on_config_changed)
         self.framework.observe(self.on.start, self._on_start)
         self.framework.observe(self.on.update_status, self._on_update_status)
         self.framework.observe(
@@ -659,25 +661,40 @@ def _on_start(self, event) -> None:
             event.defer()
             return
 
-    def _relation_changes_handler(self, event) -> None:
+    def _relation_changes_handler(self, event: RelationEvent) -> None:
         """Handles different relation events and updates MongoDB replica set."""
         self._connect_mongodb_exporter()
         self._connect_pbm_agent()
 
-        if type(event) is RelationDepartedEvent:
+        if isinstance(event, RelationDepartedEvent):
             if event.departing_unit.name == self.unit.name:
                 self.unit_peer_data.setdefault("unit_departed", "True")
 
         if not self.unit.is_leader():
             return
 
+        if self.upgrade_in_progress:
+            logger.warning(
+                "Adding replicas during an upgrade is not supported. The charm may be in a broken, unrecoverable state"
+            )
+            event.defer()
+            return
+
         # Admin password and keyFile should be created before running MongoDB.
         # This code runs on leader_elected event before mongod_pebble_ready
         self._generate_secrets()
 
         if not self.db_initialised:
             return
 
+        self._reconcile_mongo_hosts_and_users(event)
+
+    def _reconcile_mongo_hosts_and_users(self, event: RelationEvent) -> None:
+        """Auxiliary function to reconcile mongo data for relation events.
+
+        Args:
+            event: The relation event
+        """
         with MongoDBConnection(self.mongodb_config) as mongo:
             try:
                 replset_members = mongo.get_replset_members()

tests/integration/upgrades/__init__.py

@@ -0,0 +1,2 @@
+# Copyright 2024 Canonical Ltd.
+# See LICENSE file for licensing details.

tests/integration/upgrades/test_upgrades.py

@@ -0,0 +1,58 @@
+#!/usr/bin/env python3
+# Copyright 2024 Canonical Ltd.
+# See LICENSE file for licensing details.
+
+import pytest
+from pytest_operator.plugin import OpsTest
+
+from ..ha_tests.helpers import find_unit
+from ..helpers import (
+    APP_NAME,
+    check_or_scale_app,
+    get_app_name,
+    get_password,
+    set_password,
+)
+
+
+@pytest.mark.skip("Missing upgrade code for now")
+@pytest.mark.group(1)
+@pytest.mark.abort_on_fail
+async def test_build_and_deploy(ops_test: OpsTest):
+    app_name = await get_app_name(ops_test)
+
+    if app_name:
+        await check_or_scale_app(ops_test, app_name, required_units=3)
+        return
+
+    app_name = APP_NAME
+
+    await ops_test.model.deploy(
+        app_name,
+        application_name=app_name,
+        num_units=3,
+        series="jammy",
+        channel="6/edge",
+    )
+    await ops_test.model.wait_for_idle(
+        apps=[app_name], status="active", timeout=1000, idle_period=120
+    )
+
+
+@pytest.mark.skip("Missing upgrade code for now")
+@pytest.mark.group(1)
+@pytest.mark.abort_on_fail
+async def test_upgrade_password_change_fail(ops_test: OpsTest):
+    app_name = await get_app_name(ops_test)
+    leader_id = await find_unit(ops_test, leader=True, app_name=app_name)
+
+    current_password = await get_password(ops_test, leader_id, app_name=app_name)
+    new_charm = await ops_test.build_charm(".")
+    await ops_test.model.applications[app_name].refresh(path=new_charm)
+    results = await set_password(ops_test, leader_id, password="0xdeadbeef", app_name=app_name)
+
+    assert results == "Cannot set passwords while an upgrade is in progress."
+
+    after_action_password = await get_password(ops_test, leader_id, app_name=app_name)
+
+    assert current_password == after_action_password

tests/unit/test_upgrade.py

@@ -0,0 +1,69 @@
+# Copyright 2024 Canonical Ltd.
+# See LICENSE file for licensing details.
+import unittest
+from unittest.mock import Mock, PropertyMock, patch
+
+from ops.model import ActiveStatus, Relation
+from ops.testing import ActionFailed, Harness
+from parameterized import parameterized
+
+from charm import MongoDBCharm
+from config import Config
+
+from .helpers import patch_network_get
+
+
+class TestUpgrades(unittest.TestCase):
+    @patch_network_get(private_address="1.1.1.1")
+    def setUp(self, *unused):
+        self.harness = Harness(MongoDBCharm)
+        self.addCleanup(self.harness.cleanup)
+        mongo_resource = {
+            "registrypath": "mongo:4.4",
+        }
+        self.harness.add_oci_resource("mongodb-image", mongo_resource)
+        self.harness.begin()
+        self.harness.set_leader(True)
+        self.peer_rel_id = self.harness.add_relation("database-peers", "mongodb-peers")
+
+    @patch("ops.framework.EventBase.defer")
+    @patch("charm.MongoDBCharm.upgrade_in_progress", new_callable=PropertyMock)
+    def test_on_config_changed_during_upgrade_fails(self, mock_upgrade, defer):
+        def is_role_changed_mock(*args):
+            return True
+
+        self.harness.charm.is_role_changed = is_role_changed_mock
+
+        mock_upgrade.return_value = True
+        self.harness.charm.on.config_changed.emit()
+
+        defer.assert_called()
+
+    @parameterized.expand([("relation_joined"), ("relation_changed")])
+    @patch("charm.MongoDBCharm._connect_pbm_agent")
+    @patch("charm.MongoDBCharm._connect_mongodb_exporter")
+    @patch("ops.framework.EventBase.defer")
+    @patch("charm.MongoDBCharm.upgrade_in_progress", new_callable=PropertyMock)
+    def test_on_relation_handler(self, handler, mock_upgrade, defer, *unused):
+        relation: Relation = self.harness.charm.model.get_relation("database-peers")
+        mock_upgrade.return_value = True
+        getattr(self.harness.charm.on[Config.Relations.PEERS], handler).emit(relation)
+        defer.assert_called()
+
+    @patch("charm.MongoDBCharm.upgrade_in_progress", new_callable=PropertyMock)
+    def test_pass_pre_set_password_check_fails(self, mock_upgrade):
+        def mock_shard_role(*args):
+            return args != ("shard",)
+
+        mock_pbm_status = Mock(return_value=ActiveStatus())
+        self.harness.charm.is_role = mock_shard_role
+        mock_upgrade.return_value = True
+        self.harness.charm.backups.get_pbm_status = mock_pbm_status
+
+        with self.assertRaises(ActionFailed) as action_failed:
+            self.harness.run_action("set-password")
+
+        assert (
+            action_failed.exception.message
+            == "Cannot set passwords while an upgrade is in progress."
+        )