Legacy interface fails to relate if event attributes are None

src/relations/db.py

@@ -20,14 +20,16 @@
     RelationDepartedEvent,
 )
 from ops.framework import Object
-from ops.model import BlockedStatus, Relation, Unit
+from ops.model import ActiveStatus, BlockedStatus, Relation, Unit
 from pgconnstr import ConnectionString
 
 from constants import DATABASE_PORT
 from utils import new_password
 
 logger = logging.getLogger(__name__)
 
+EXTENSIONS_BLOCKING_MESSAGE = "extensions requested through relation"
+
 
 class DbProvides(Object):
     """Defines functionality for the 'provides' side of the 'db' relation.
@@ -66,6 +68,21 @@ def __init__(self, charm: CharmBase, admin: bool = False):
         self.admin = admin
         self.charm = charm
 
+    def _check_for_blocking_relations(self, relation_id: int) -> bool:
+        """Checks if there are relations with extensions.
+
+        Args:
+            relation_id: current relation to be skipped
+        """
+        for relname in ["db", "db-admin"]:
+            for relation in self.charm.model.relations.get(relname, []):
+                if relation.id == relation_id:
+                    continue
+                for data in relation.data.values():
+                    if "extensions" in data:
+                        return True
+        return False
+
     def _on_relation_changed(self, event: RelationChangedEvent) -> None:
         """Handle the legacy db/db-admin relation changed event.
 
@@ -92,20 +109,20 @@ def _on_relation_changed(self, event: RelationChangedEvent) -> None:
         application_relation_databag = event.relation.data[self.charm.app]
 
         # Do not allow apps requesting extensions to be installed.
-        if "extensions" in event.relation.data[
-            event.app
-        ] or "extensions" in event.relation.data.get(event.unit, {}):
+        if "extensions" in event.relation.data.get(
+            event.app, {}
+        ) or "extensions" in event.relation.data.get(event.unit, {}):
             logger.error(
                 "ERROR - `extensions` cannot be requested through relations"
                 " - they should be installed through a database charm config in the future"
             )
-            self.charm.unit.status = BlockedStatus("extensions requested through relation")
+            self.charm.unit.status = BlockedStatus(EXTENSIONS_BLOCKING_MESSAGE)
             return
 
         # Sometimes a relation changed event is triggered,
         # and it doesn't have a database name in it.
-        database = event.relation.data[event.app].get(
-            "database", event.relation.data[event.unit].get("database")
+        database = event.relation.data.get(event.app, {}).get(
+            "database", event.relation.data.get(event.unit, {}).get("database")
         )
         if not database:
             logger.warning("No database name provided")
@@ -228,6 +245,17 @@ def _on_relation_broken(self, event: RelationBrokenEvent) -> None:
                 f"Failed to delete user during {self.relation_name} relation broken event"
             )
 
+        # Clean up Blocked status if caused by the departed relation
+        if (
+            self.charm._has_blocked_status
+            and self.charm.unit.status.message == EXTENSIONS_BLOCKING_MESSAGE
+        ):
+            if "extensions" in event.relation.data.get(
+                event.app, {}
+            ) or "extensions" in event.relation.data.get(event.unit, {}):
+                if not self._check_for_blocking_relations(event.relation.id):
+                    self.charm.unit.status = ActiveStatus()
+
     def update_endpoints(self, event: RelationChangedEvent = None) -> None:
         """Set the read/write and read-only endpoints."""
         if not self.charm.unit.is_leader():

tests/integration/test_db.py

@@ -193,4 +193,31 @@ async def test_nextcloud_db_blocked(ops_test: OpsTest, charm: str) -> None:
         except JujuUnitError:
             pass
 
-        leader_unit.workload_status_message == "extensions requested through relation"
+        assert leader_unit.workload_status_message == "extensions requested through relation"
+
+        await ops_test.model.remove_application("nextcloud", block_until_done=True)
+
+
+@pytest.mark.db_relation_tests
+async def test_weebl_db(ops_test: OpsTest, charm: str) -> None:
+    async with ops_test.fast_forward():
+        await ops_test.model.deploy(
+            "weebl",
+            application_name="weebl",
+            num_units=APPLICATION_UNITS,
+        )
+        await ops_test.model.wait_for_idle(
+            apps=["weebl"],
+            status="blocked",
+            raise_on_blocked=False,
+            timeout=1000,
+        )
+
+        await ops_test.model.relate("weebl:database", f"{DATABASE_APP_NAME}:db")
+
+        await ops_test.model.wait_for_idle(
+            apps=["weebl", DATABASE_APP_NAME],
+            status="active",
+            raise_on_blocked=False,
+            timeout=1000,
+        )

tests/unit/test_db.py

@@ -211,6 +211,70 @@ def test_on_relation_broken(self, _on_relation_departed, _member_started, _prima
             self.harness.remove_relation(self.rel_id)
             self.assertTrue(isinstance(self.harness.model.unit.status, BlockedStatus))
 
+    @patch(
+        "charm.PostgresqlOperatorCharm.primary_endpoint",
+        new_callable=PropertyMock,
+    )
+    @patch("charm.PostgresqlOperatorCharm._has_blocked_status", new_callable=PropertyMock)
+    @patch("charm.Patroni.member_started", new_callable=PropertyMock)
+    @patch("charm.DbProvides._on_relation_departed")
+    def test_on_relation_broken_extensions_unblock(
+        self, _on_relation_departed, _member_started, _primary_endpoint, _has_blocked_status
+    ):
+        with patch.object(PostgresqlOperatorCharm, "postgresql", Mock()) as postgresql_mock:
+            # Set some side effects to test multiple situations.
+            _has_blocked_status.return_value = True
+            _member_started.return_value = True
+            _primary_endpoint.return_value = {"1.1.1.1"}
+            postgresql_mock.delete_user = PropertyMock(return_value=None)
+            self.harness.model.unit.status = BlockedStatus("extensions requested through relation")
+            with self.harness.hooks_disabled():
+                self.rel_id = self.harness.add_relation(RELATION_NAME, "application")
+                self.harness.update_relation_data(
+                    self.rel_id,
+                    "application",
+                    {"database": DATABASE, "extensions": ["test"]},
+                )
+
+            # Break the relation before the database is ready.
+            self.harness.remove_relation(self.rel_id)
+            self.assertTrue(isinstance(self.harness.model.unit.status, ActiveStatus))
+
+    @patch(
+        "charm.PostgresqlOperatorCharm.primary_endpoint",
+        new_callable=PropertyMock,
+    )
+    @patch("charm.PostgresqlOperatorCharm._has_blocked_status", new_callable=PropertyMock)
+    @patch("charm.Patroni.member_started", new_callable=PropertyMock)
+    @patch("charm.DbProvides._on_relation_departed")
+    def test_on_relation_broken_extensions_keep_block(
+        self, _on_relation_departed, _member_started, _primary_endpoint, _has_blocked_status
+    ):
+        with patch.object(PostgresqlOperatorCharm, "postgresql", Mock()) as postgresql_mock:
+            # Set some side effects to test multiple situations.
+            _has_blocked_status.return_value = True
+            _member_started.return_value = True
+            _primary_endpoint.return_value = {"1.1.1.1"}
+            postgresql_mock.delete_user = PropertyMock(return_value=None)
+            self.harness.model.unit.status = BlockedStatus("extensions requested through relation")
+            with self.harness.hooks_disabled():
+                self.rel_id = self.harness.add_relation(RELATION_NAME, "application")
+                self.harness.update_relation_data(
+                    self.rel_id,
+                    "application",
+                    {"database": DATABASE, "extensions": ["test"]},
+                )
+                second_rel_id = self.harness.add_relation(RELATION_NAME, "application2")
+                self.harness.update_relation_data(
+                    second_rel_id,
+                    "application2",
+                    {"database": DATABASE, "extensions": ["test"]},
+                )
+
+            # Break the relation before the database is ready.
+            self.harness.remove_relation(self.rel_id)
+            self.assertTrue(isinstance(self.harness.model.unit.status, BlockedStatus))
+
     @patch(
         "charm.DbProvides._get_state",
         side_effect="postgresql/0",