rewrite event handlers to only run data migration on database created event

src/charm.py

@@ -8,7 +8,10 @@
 
 import logging
 
-from charms.data_platform_libs.v0.database_requires import DatabaseCreatedEvent, DatabaseRequires
+from charms.data_platform_libs.v0.database_requires import (
+    DatabaseEndpointsChangedEvent,
+    DatabaseRequires,
+)
 from charms.observability_libs.v0.kubernetes_service_patch import KubernetesServicePatch
 from charms.traefik_k8s.v1.ingress import (
     IngressPerAppReadyEvent,
@@ -18,8 +21,8 @@
 from jinja2 import Template
 from ops.charm import CharmBase
 from ops.main import main
-from ops.model import ActiveStatus, BlockedStatus, MaintenanceStatus, WaitingStatus
-from ops.pebble import ChangeError, ExecError, Layer
+from ops.model import ActiveStatus, BlockedStatus, MaintenanceStatus, ModelError, WaitingStatus
+from ops.pebble import ExecError, Layer
 
 logger = logging.getLogger(__name__)
 KRATOS_ADMIN_PORT = 4434
@@ -62,7 +65,7 @@ def __init__(self, *args):
         )
 
         self.framework.observe(self.on.kratos_pebble_ready, self._on_pebble_ready)
-        self.framework.observe(self.database.on.database_created, self._on_database_changed)
+        self.framework.observe(self.database.on.database_created, self._on_database_created)
         self.framework.observe(self.database.on.endpoints_changed, self._on_database_changed)
         self.framework.observe(self.admin_ingress.on.ready, self._on_admin_ingress_ready)
         self.framework.observe(self.admin_ingress.on.revoked, self._on_ingress_revoked)
@@ -78,7 +81,7 @@ def _pebble_layer(self) -> Layer:
                 self._container_name: {
                     "override": "replace",
                     "summary": "Kratos Operator layer",
-                    "startup": "enabled",
+                    "startup": "disabled",
                     "command": f"kratos serve all --config {self._config_file_path}",
                 }
             },
@@ -111,94 +114,116 @@ def _render_conf_file(self) -> None:
         )
         return rendered
 
-    def _update_layer(self) -> None:
-        """Updates the Pebble configuration layer and kratos config if changed."""
-        config = self._render_conf_file()
-        if not self._container.get_plan().to_dict():
-            self.unit.status = MaintenanceStatus("Applying new pebble layer")
-            self._container.push(self._config_file_path, config, make_dirs=True)
-            with open("src/identity.default.schema.json", encoding="utf-8") as schema_file:
-                schema = schema_file.read()
-                self._container.push(self._identity_schema_file_path, schema, make_dirs=True)
-            self._container.add_layer(self._container_name, self._pebble_layer, combine=True)
-            logger.info("Pebble plan updated with new configuration, replanning")
-            self._container.replan()
-        else:
-            # Compare changes in kratos config
-            current_config = self._container.pull(self._config_file_path).read()
-            if current_config != config:
-                self.unit.status = MaintenanceStatus("Updating Kratos Config")
-                self._container.push(self._config_file_path, config, make_dirs=True)
-                logger.info("Updated kratos config")
-                self._container.restart(self._container_name)
-
     def _get_database_relation_info(self) -> dict:
         """Get database info from relation data bag."""
         relation_id = self.database.relations[0].id
         relation_data = self.database.fetch_relation_data()[relation_id]
 
         return {
-            "username": relation_data["username"],
-            "password": relation_data["password"],
-            "endpoints": relation_data["endpoints"],
+            "username": relation_data.get("username"),
+            "password": relation_data.get("password"),
+            "endpoints": relation_data.get("endpoints"),
             "database_name": self._db_name,
         }
 
-    def _update_container(self, event) -> None:
-        """Update configs, pebble layer and run database migration."""
+    def _run_sql_migration(self) -> bool:
+        """Runs database migration.
+
+        Returns True if migration was run successfully, else returns false.
+        """
+        try:
+            process = self._container.exec(
+                ["kratos", "migrate", "sql", "-e", "--config", self._config_file_path, "--yes"],
+            )
+            stdout, _ = process.wait_output()
+            logger.info(f"Successfully executed automigration: {stdout}")
+        except ExecError as err:
+            logger.error(f"Exited with code {err.exit_code}. Stderr: {err.stderr}")
+            self.unit.status = BlockedStatus("Database migration job failed")
+            return False
+
+        return True
+
+    def _on_pebble_ready(self, event) -> None:
+        """Event Handler for pebble ready event."""
         if not self._container.can_connect():
             event.defer()
             logger.info("Cannot connect to Kratos container. Deferring event.")
             self.unit.status = WaitingStatus("Waiting to connect to Kratos container")
             return
 
-        if not self.model.relations["pg-database"]:
-            logger.error("Missing required relation with postgresql")
-            self.model.unit.status = BlockedStatus("Missing required relation with postgresql")
+        self.unit.status = MaintenanceStatus("Configuring/deploying resources")
+
+        with open("src/identity.default.schema.json", encoding="utf-8") as schema_file:
+            schema = schema_file.read()
+            self._container.push(self._identity_schema_file_path, schema, make_dirs=True)
+
+        self._container.add_layer(self._container_name, self._pebble_layer, combine=True)
+        logger.info("Pebble plan updated with new configuration, replanning")
+        self._container.replan()
+
+        # in case container was terminated unexpectedly
+        if self.database.is_database_created():
+            self._container.push(self._config_file_path, self._render_conf_file(), make_dirs=True)
+            self._container.start(self._container_name)
+            self.unit.status = ActiveStatus()
             return
 
-        if not self.database.is_database_created():
-            event.defer()
-            logger.info("Missing database details. Deferring event.")
+        if self.model.relations["pg-database"]:
             self.unit.status = WaitingStatus("Waiting for database creation")
+        else:
+            self.unit.status = BlockedStatus("Missing postgres database relation")
+
+    def _on_database_created(self, event) -> None:
+        """Event Handler for database created event."""
+        if not self._container.can_connect():
+            event.defer()
+            logger.info("Cannot connect to Kratos container. Deferring event.")
+            self.unit.status = WaitingStatus("Waiting to connect to Kratos container")
             return
 
+        self.unit.status = MaintenanceStatus(
+            "Configuring container and resources for database connection"
+        )
+
         try:
-            self._update_layer()
-        except ChangeError as err:
-            logger.error(str(err))
-            self.unit.status = BlockedStatus("Failed to replan")
+            self._container.get_service(self._container_name)
+        except (ModelError, RuntimeError):
+            event.defer()
+            self.unit.status = WaitingStatus("Waiting for Kratos service")
+            logger.info("Kratos service is absent. Deferring database created event.")
             return
 
-        if not self.unit.is_leader():
-            return
+        logger.info("Updating Kratos config and restarting service")
+        self._container.push(self._config_file_path, self._render_conf_file(), make_dirs=True)
 
-        self._run_sql_migration()
+        if self.unit.is_leader() and not self._run_sql_migration():
+            self.unit.status = BlockedStatus("Database migration failed.")
+        else:
+            self._container.start(self._container_name)
+            self.unit.status = ActiveStatus()
 
-        self.unit.status = ActiveStatus()
+    def _on_database_changed(self, event: DatabaseEndpointsChangedEvent) -> None:
+        """Event Handler for database changed event."""
+        if not self._container.can_connect():
+            event.defer()
+            logger.info("Cannot connect to Kratos container. Deferring event.")
+            self.unit.status = WaitingStatus("Waiting to connect to Kratos container")
+            return
 
-    def _run_sql_migration(self) -> None:
-        """Runs database migration."""
-        process = self._container.exec(
-            ["kratos", "migrate", "sql", "-e", "--config", self._config_file_path, "--yes"],
-            timeout=20.0,
-        )
-        try:
-            stdout, _ = process.wait_output()
-            logger.info(f"Successfully executed automigration: {stdout}")
-        except ExecError as err:
-            logger.error(f"Exited with code {err.exit_code}. Stderr: {err.stderr}")
-            self.unit.status = BlockedStatus("Database migration job failed")
+        self.unit.status = MaintenanceStatus("Updating database details")
 
-    def _on_pebble_ready(self, event) -> None:
-        """Event Handler for pebble ready event."""
-        self.unit.status = MaintenanceStatus("Configuring/deploying resources")
-        self._update_container(event)
+        try:
+            self._container.get_service(self._container_name)
+        except (ModelError, RuntimeError):
+            event.defer()
+            self.unit.status = WaitingStatus("Waiting for Kratos service")
+            logger.info("Kratos service is absent. Deferring database created event.")
+            return
 
-    def _on_database_changed(self, event: DatabaseCreatedEvent) -> None:
-        """Event Handler for database created event."""
-        self.unit.status = MaintenanceStatus("Retrieving database details")
-        self._update_container(event)
+        self._container.push(self._config_file_path, self._render_conf_file(), make_dirs=True)
+        self._container.restart(self._container_name)
+        self.unit.status = ActiveStatus()
 
     def _on_admin_ingress_ready(self, event: IngressPerAppReadyEvent) -> None:
         if self.unit.is_leader():

templates/kratos.yaml.j2

@@ -3,15 +3,16 @@ log:
 identity:
     default_schema_id: default
     schemas:
-    - id: default
-      url: "file://{{ identity_schema_file_path }}"
+        - id: default
+          url: 'file://{{ identity_schema_file_path }}'
 selfservice:
-    default_browser_return_url: {{ default_browser_return_url | d("http://127.0.0.1:9999/")}}
+    default_browser_return_url:
+        {{ default_browser_return_url | d("http://127.0.0.1:9999/") }}
     flows:
         registration:
             enabled: True
             ui_url: {{ registration_ui_url }}
-dsn: postgres://{{ db_info['username'] }}:{{ db_info['password'] }}@{{ db_info['endpoints'] }}/{{ db_info['database_name'] }}
+dsn: postgres://{{ db_info.get('username') }}:{{ db_info.get('password') }}@{{ db_info.get('endpoints') }}/{{ db_info.get('database_name') }}
 courier:
     smtp:
         connection_uri: {{ smtp_connection_uri }}

tests/unit/conftest.py

@@ -2,16 +2,19 @@
 # See LICENSE file for licensing details.
 
 import pytest
+from ops.pebble import ExecError
 from ops.testing import Harness
 
 from charm import KratosCharm
 
 
 @pytest.fixture()
-def harness() -> None:
+def harness(mocked_kubernetes_service_patcher) -> None:
     harness = Harness(KratosCharm)
     harness.set_model_name("kratos-model")
+    harness.set_can_connect("kratos", True)
     harness.set_leader(True)
+    harness.begin()
     return harness
 
 
@@ -30,12 +33,22 @@ def mocked_fqdn(mocker):
 
 
 @pytest.fixture()
-def mocked_sql_migration(mocker):
-    mocked_sql_migration = mocker.patch("charm.KratosCharm._run_sql_migration")
-    yield mocked_sql_migration
+def mocked_pebble_exec(mocker):
+    mocked_pebble_exec = mocker.patch("ops.model.Container.exec")
+    yield mocked_pebble_exec
 
 
 @pytest.fixture()
-def mocked_update_container(mocker):
-    mocked_update_container = mocker.patch("charm.KratosCharm._update_container")
-    yield mocked_update_container
+def mocked_pebble_exec_success(mocker, mocked_pebble_exec):
+    mocked_process = mocker.patch("ops.pebble.ExecProcess")
+    mocked_process.wait_output.return_value = ("Success", None)
+    mocked_pebble_exec.return_value = mocked_process
+    yield mocked_pebble_exec
+
+
+@pytest.fixture()
+def mocked_pebble_exec_failed(mocked_pebble_exec):
+    mocked_pebble_exec.side_effect = ExecError(
+        exit_code=400, stderr="Failed to execute", stdout="Failed", command="test command"
+    )
+    yield