fix: handle database relation departed

Closes #137

src/charm.py

@@ -372,6 +372,12 @@ def _run_sql_migration(self, timeout: float = 60) -> bool:
     def _oauth_relation_peer_data_key(self, relation_id: int) -> str:
         return f"oauth_{relation_id}"
 
+    @property
+    def _migration_peer_data_key(self) -> Optional[str]:
+        if not self.database.relations:
+            return None
+        return f"{DB_MIGRATION_VERSION_KEY}_{self.database.relations[0].id}"
+
     @property
     def _peers(self) -> Optional[Relation]:
         """Fetch the peer relation."""
@@ -474,13 +480,15 @@ def _handle_status_update_config(self, event: HookEvent) -> None:
             event.defer()
             return
 
+        self._cleanup_peer_data()
         self._container.push(self._hydra_config_path, self._render_conf_file(), make_dirs=True)
         try:
             self._container.restart(self._container_name)
         except ChangeError as err:
             logger.error(str(err))
             self.unit.status = BlockedStatus("Failed to restart, please consult the logs")
             return
+
         self.unit.status = ActiveStatus()
 
     def _on_leader_elected(self, event: LeaderElectedEvent) -> None:
@@ -527,16 +535,12 @@ def _migration_is_needed(self):
         if not self._peers:
             return
 
-        return self._get_peer_data(DB_MIGRATION_VERSION_KEY) != self._hydra_cli.get_version()
+        return self._get_peer_data(self._migration_peer_data_key) != self._hydra_cli.get_version()
 
     def _on_database_created(self, event: DatabaseCreatedEvent) -> None:
         """Event Handler for database created event."""
         logger.info("Retrieved database details")
 
-        if not self.model.relations[self._db_relation_name]:
-            self.unit.status = BlockedStatus("Missing required relation with postgresql")
-            return
-
         if not self._container.can_connect():
             event.defer()
             logger.info("Cannot connect to Hydra container. Deferring the event.")
@@ -595,11 +599,21 @@ def _on_run_migration(self, event: ActionEvent) -> None:
 
     def _on_database_relation_departed(self, event: RelationDepartedEvent) -> None:
         """Event Handler for database relation departed event."""
-        logger.error("Missing required relation with postgresql")
-        self.model.unit.status = BlockedStatus("Missing required relation with postgresql")
-        self._pop_peer_data(DB_MIGRATION_VERSION_KEY)
-        if self._container.can_connect():
-            self._container.stop(self._container_name)
+        self.unit.status = BlockedStatus("Missing required relation with postgresql")
+
+    def _cleanup_peer_data(self) -> None:
+        if not self._peers:
+            return
+        # We need to remove the migration key from peer data. We can't do that in relation
+        # departed as we can't tell if the event was triggered from a unit dying of the
+        # relation being actually departed.
+        extra_keys = [
+            k
+            for k in self._peers.data[self.app].keys()
+            if k.startswith(DB_MIGRATION_VERSION_KEY) and k != self._migration_peer_data_key
+        ]
+        for k in extra_keys:
+            self._pop_peer_data(k)
 
     def _on_admin_ingress_ready(self, event: IngressPerAppReadyEvent) -> None:
         if self.unit.is_leader():

tests/integration/test_charm.py

@@ -330,3 +330,16 @@ async def test_hydra_scale_up(ops_test: OpsTest) -> None:
         timeout=1000,
         wait_for_exact_units=3,
     )
+
+
+async def test_hydra_scale_down(ops_test: OpsTest) -> None:
+    """Check that hydra works after it is scaled down."""
+    app = ops_test.model.applications[HYDRA_APP]
+
+    await app.scale(1)
+
+    await ops_test.model.wait_for_idle(
+        apps=[HYDRA_APP],
+        status="active",
+        timeout=1000,
+    )

tests/unit/test_charm.py

@@ -199,7 +199,7 @@ def test_relation_data(harness: Harness, mocked_run_migration: MagicMock) -> Non
     assert relation_data["endpoints"] == "postgresql-k8s-primary.namespace.svc.cluster.local:5432"
 
 
-def test_relation_departed(harness: Harness, mocked_run_migration: MagicMock) -> None:
+def test_relation_departed(harness: Harness) -> None:
     db_relation_id = setup_postgres_relation(harness)
 
     harness.remove_relation_unit(db_relation_id, "postgresql-k8s/0")
@@ -250,6 +250,7 @@ def test_postgres_created_when_migration_has_run(
     harness.set_leader(False)
     harness.set_can_connect(CONTAINER_NAME, True)
     harness.charm.on.hydra_pebble_ready.emit(CONTAINER_NAME)
+    setup_ingress_relation(harness, "public")
     setup_peer_relation(harness)
 
     setup_postgres_relation(harness)