fix: update run-migration action

actions.yaml

@@ -1,5 +1,12 @@
 run-migration:
-  description: Run a command to create SQL schemas and apply migration plans.
+  description:
+    Run a migration, this is needed after upgrades. This is a non-reversible operation.
+    Run this after backing up the database.
+  params:
+    timeout:
+      description: Timeout after which the migration will be canceled
+      type: number
+      default: 120
 create-oauth-client:
   description: Register an oauth client
   params:

src/charm.py

@@ -350,9 +350,11 @@ def _set_version(self) -> None:
         self.unit.set_workload_version(version)
 
     def _get_database_relation_info(self) -> dict:
+        if not self.database.relations:
+            return None
+
         relation_id = self.database.relations[0].id
         relation_data = self.database.fetch_relation_data()[relation_id]
-
         return {
             "username": relation_data.get("username"),
             "password": relation_data.get("password"),
@@ -597,11 +599,25 @@ def _on_database_changed(self, event: DatabaseEndpointsChangedEvent) -> None:
 
     def _on_run_migration(self, event: ActionEvent) -> None:
         """Runs the migration as an action response."""
-        logger.info("Executing database migration initiated by user")
-        if not self._run_sql_migration():
-            event.fail("Execution failed, please inspect the logs")
+        if not self._container.can_connect():
+            event.fail("Service is not ready. Please re-run the action when the charm is active")
             return
-        event.log("Successfully ran migration")
+
+        timeout = float(event.params.get("timeout", 120))
+        event.log("Migrating database.")
+        try:
+            self._hydra_cli.run_migration(timeout=timeout, dsn=self._dsn)
+        except Error as e:
+            err_msg = e.stderr if isinstance(e, ExecError) else e
+            event.fail(f"Database migration action failed: {err_msg}")
+            return
+        event.log("Successfully migrated the database.")
+
+        if not self._peers:
+            event.fail("Peer relation not ready. Failed to store migration version")
+            return
+        self._set_peer_data(self._migration_peer_data_key, self._hydra_cli.get_version())
+        event.log("Updated migration version in peer data.")
 
     def _on_database_relation_departed(self, event: RelationDepartedEvent) -> None:
         """Event Handler for database relation departed event."""

tests/unit/test_charm.py

@@ -9,7 +9,7 @@
 import pytest
 import yaml
 from ops.model import ActiveStatus, BlockedStatus, WaitingStatus
-from ops.pebble import ExecError
+from ops.pebble import ExecError, TimeoutError
 from ops.testing import Harness
 from test_oauth_requirer import CLIENT_CONFIG  # type: ignore
 
@@ -1141,3 +1141,41 @@ def test_verify_pebble_layer_tempo_k8s(harness: Harness) -> None:
     }
 
     assert harness.charm._hydra_layer.to_dict() == expected_layer
+
+
+def test_run_migration_action(harness: Harness, mocked_run_migration: MagicMock) -> None:
+    harness.set_can_connect(CONTAINER_NAME, True)
+    setup_peer_relation(harness)
+    setup_postgres_relation(harness)
+    event = MagicMock()
+
+    harness.charm._on_run_migration(event)
+
+    mocked_run_migration.assert_called_once()
+    event.fail.assert_not_called()
+
+
+def test_error_on_run_migration_action(harness: Harness, mocked_run_migration: MagicMock) -> None:
+    harness.set_can_connect(CONTAINER_NAME, True)
+    mocked_run_migration.side_effect = ExecError(
+        command=[], exit_code=1, stdout="", stderr="Error"
+    )
+    event = MagicMock()
+
+    harness.charm._on_run_migration(event)
+
+    mocked_run_migration.assert_called_once()
+    event.fail.assert_called()
+
+
+def test_timeout_on_run_migration_action(
+    harness: Harness, mocked_run_migration: MagicMock
+) -> None:
+    harness.set_can_connect(CONTAINER_NAME, True)
+    mocked_run_migration.side_effect = TimeoutError
+    event = MagicMock()
+
+    harness.charm._on_run_migration(event)
+
+    mocked_run_migration.assert_called_once()
+    event.fail.assert_called()