Add database migration script support (#50)

.github/workflows/integration_test.yaml

@@ -8,4 +8,4 @@ jobs:
     uses: canonical/operator-workflows/.github/workflows/integration_test.yaml@main
     secrets: inherit
     with:
-      modules: '["test_charm", "test_proxy", "test_cos", "test_database", "test_default", "test_wsgi_path"]'
+      modules: '["test_charm", "test_proxy", "test_cos", "test_database", "test_db_migration", "test_default", "test_wsgi_path"]'

config.yaml

@@ -1,6 +1,18 @@
 # Copyright 2023 Canonical Ltd.
 # See LICENSE file for licensing details.
 options:
+  database_migration_script:
+    type: string
+    description: >-
+      Specifies the relative path from /srv/flask/app that points to a shell script
+      executing database migrations for the Flask application. This script is
+      designed to run once for each Flask container unit. However, users must ensure:
+      1. The script can be executed multiple times without issues;
+      2. Concurrent migrations from different units are safe.
+      In case of migration failure, the charm will re-attempt during the 
+      update-status event. Successful database migration in a container ensures that
+      any configuration updates won't trigger another migration unless 
+      the Flask container is upgraded or restarted.
   flask_application_root:
     type: string
     description: >-

requirements.txt

@@ -1,4 +1,4 @@
 cosl
 jsonschema >=4.19,<4.20
-ops >= 2.2
+ops >= 2.6
 pydantic >= 1.10,<2

src/charm.py

@@ -9,12 +9,12 @@
 
 import ops
 from charms.traefik_k8s.v1.ingress import IngressPerAppRequirer
-from ops.main import main
 
 from charm_state import CharmState
 from constants import FLASK_CONTAINER_NAME
+from database_migration import DatabaseMigration, DatabaseMigrationStatus
 from databases import Databases, get_uris, make_database_requirers
-from exceptions import CharmConfigInvalidError, PebbleNotReadyError
+from exceptions import CharmConfigInvalidError
 from flask_app import FlaskApp
 from observability import Observability
 from secret_storage import SecretStorage
@@ -35,6 +35,7 @@ def __init__(self, *args: typing.Any) -> None:
         super().__init__(*args)
         self._secret_storage = SecretStorage(charm=self)
         database_requirers = make_database_requirers(self)
+
         try:
             self._charm_state = CharmState.from_charm(
                 charm=self,
@@ -44,12 +45,20 @@ def __init__(self, *args: typing.Any) -> None:
         except CharmConfigInvalidError as exc:
             self._update_app_and_unit_status(ops.BlockedStatus(exc.msg))
             return
-        self._webserver = GunicornWebserver(
+
+        self._database_migration = DatabaseMigration(
+            flask_container=self.unit.get_container(FLASK_CONTAINER_NAME),
+            charm_state=self._charm_state,
+        )
+        webserver = GunicornWebserver(
             charm_state=self._charm_state,
             flask_container=self.unit.get_container(FLASK_CONTAINER_NAME),
         )
         self._flask_app = FlaskApp(
-            charm=self, charm_state=self._charm_state, webserver=self._webserver
+            charm=self,
+            charm_state=self._charm_state,
+            webserver=webserver,
+            database_migration=self._database_migration,
         )
         self._databases = Databases(
             charm=self,
@@ -77,22 +86,6 @@ def __init__(self, *args: typing.Any) -> None:
             self.on.secret_storage_relation_changed, self._on_secret_storage_relation_changed
         )
 
-    def container(self) -> ops.Container:
-        """Get the flask application workload container controller.
-
-        Return:
-            The controller of the flask application workload container.
-
-        Raises:
-            PebbleNotReadyError: if the pebble service inside the container is not ready while the
-                ``require_connected`` is set to True.
-        """
-        if not self.unit.get_container(FLASK_CONTAINER_NAME).can_connect():
-            raise PebbleNotReadyError("pebble inside flask-app container is not ready")
-
-        container = self.unit.get_container(FLASK_CONTAINER_NAME)
-        return container
-
     def _on_config_changed(self, _event: ops.EventBase) -> None:
         """Configure the flask pebble service layer.
 
@@ -169,6 +162,11 @@ def _restart_flask(self) -> None:
         except CharmConfigInvalidError as exc:
             self._update_app_and_unit_status(ops.BlockedStatus(exc.msg))
 
+    def _on_update_status(self, _: ops.HookEvent) -> None:
+        """Handle the update-status event."""
+        if self._database_migration.get_status() == DatabaseMigrationStatus.FAILED:
+            self._restart_flask()
+
 
 if __name__ == "__main__":  # pragma: nocover
-    main(FlaskCharm)
+    ops.main.main(FlaskCharm)

src/charm_state.py

@@ -21,13 +21,15 @@
 )
 
 from charm_types import WebserverConfig
+from constants import FLASK_APP_DIR
 from exceptions import CharmConfigInvalidError
 from secret_storage import SecretStorage
 
 if typing.TYPE_CHECKING:
     from charm import FlaskCharm
 
 KNOWN_CHARM_CONFIG = (
+    "database_migration_script",
     "flask_application_root",
     "flask_debug",
     "flask_env",
@@ -105,6 +107,7 @@ class CharmState:  # pylint: disable=too-many-instance-attributes
         flask_config: the value of the flask_config charm configuration.
         app_config: user-defined configurations for the Flask application.
         base_dir: the base directory of the Flask application.
+        database_migration_script: The database migration script path.
         database_uris: a mapping of available database environment variable to database uris.
         flask_dir: the path to the Flask directory.
         flask_wsgi_app_path: the path to the Flask directory.
@@ -121,6 +124,7 @@ def __init__(
         self,
         *,
         app_config: dict[str, int | str | bool] | None = None,
+        database_migration_script: str | None = None,
         database_uris: dict[str, str] | None = None,
         flask_config: dict[str, int | str] | None = None,
         flask_secret_key: str | None = None,
@@ -137,6 +141,7 @@ def __init__(
             app_config: User-defined configuration values for the Flask configuration.
             flask_config: The value of the flask_config charm configuration.
             flask_secret_key: The secret storage manager associated with the charm.
+            database_migration_script: The database migration script path
             database_uris: The database uri environment variables.
             is_secret_storage_ready: whether the secret storage system is ready.
             webserver_workers: The number of workers to use for the web server,
@@ -161,6 +166,7 @@ def __init__(
         self._is_secret_storage_ready = is_secret_storage_ready
         self._flask_secret_key = flask_secret_key
         self.database_uris = database_uris if database_uris is not None else {}
+        self.database_migration_script = database_migration_script
 
     @property
     def proxy(self) -> "ProxyConfig":
@@ -213,9 +219,17 @@ def from_charm(
             )
             error_field_str = " ".join(f"flask_{f}" for f in error_fields)
             raise CharmConfigInvalidError(f"invalid configuration: {error_field_str}") from exc
+        database_migration_script = charm.config.get("database_migration_script")
+        if database_migration_script:
+            database_migration_script = os.path.normpath(FLASK_APP_DIR / database_migration_script)
+            if not database_migration_script.startswith(str(FLASK_APP_DIR)):
+                raise CharmConfigInvalidError(
+                    f"database_migration_script is not inside {FLASK_APP_DIR}"
+                )
         return cls(
             flask_config=valid_flask_config.dict(exclude_unset=True, exclude_none=True),
             app_config=typing.cast(dict[str, str | int | bool], app_config),
+            database_migration_script=database_migration_script,
             database_uris=database_uris,
             webserver_workers=int(workers) if workers is not None else None,
             webserver_threads=int(threads) if threads is not None else None,

src/constants.py

@@ -3,8 +3,13 @@
 
 """This module defines constants used throughout the Flask application."""
 
+import pathlib
+
 FLASK_CONTAINER_NAME = "flask-app"
 FLASK_SERVICE_NAME = "flask-app"
 FLASK_ENV_CONFIG_PREFIX = "FLASK_"
 FLASK_DATABASE_NAME = "flask-app"
 FLASK_SUPPORTED_DB_INTERFACES = {"mysql_client": "mysql", "postgresql_client": "postgresql"}
+FLASK_BASE_DIR = pathlib.Path("/srv/flask")
+FLASK_APP_DIR = FLASK_BASE_DIR / "app"
+FLASK_STATE_DIR = FLASK_BASE_DIR / "state"

src/database_migration.py

@@ -0,0 +1,132 @@
+# Copyright 2023 Canonical Ltd.
+# See LICENSE file for licensing details.
+
+"""Provide the DatabaseMigration class to manage database migrations."""
+import enum
+import logging
+from typing import cast
+
+import ops
+from ops.pebble import ExecError
+
+from charm_state import CharmState
+from constants import FLASK_APP_DIR, FLASK_STATE_DIR
+from exceptions import CharmConfigInvalidError
+
+logger = logging.getLogger(__name__)
+
+
+class DatabaseMigrationStatus(str, enum.Enum):
+    """Database migration status.
+
+    Attrs:
+        COMPLETED: A status denoting a successful database migration.
+        FAILED: A status denoting an unsuccessful database migration.
+        PENDING: A status denoting a pending database migration.
+    """
+
+    COMPLETED = "COMPLETED"
+    FAILED = "FAILED"
+    PENDING = "PENDING"
+
+
+class DatabaseMigration:
+    """The DatabaseMigration class that manages database migrations.
+
+    Attrs:
+        script: the database migration script.
+    """
+
+    _STATUS_FILE = FLASK_STATE_DIR / "database-migration-status"
+    _COMPLETED_SCRIPT_FILE = FLASK_STATE_DIR / "completed-database-migration"
+
+    def __init__(self, flask_container: ops.Container, charm_state: CharmState):
+        """Initialize the DatabaseMigration instance.
+
+        Args:
+            flask_container: The flask container object.
+            charm_state: The charm state.
+        """
+        self._container = flask_container
+        self._charm_state = charm_state
+
+    @property
+    def script(self) -> str | None:
+        """Get the database migration script."""
+        return self._charm_state.database_migration_script
+
+    def get_status(self) -> DatabaseMigrationStatus:
+        """Get the database migration run status.
+
+        Returns:
+            One of "PENDING", "COMPLETED", or "FAILED".
+        """
+        return (
+            DatabaseMigrationStatus.PENDING
+            if not self._container.exists(self._STATUS_FILE)
+            else DatabaseMigrationStatus(cast(str, self._container.pull(self._STATUS_FILE).read()))
+        )
+
+    def _set_status(self, status: DatabaseMigrationStatus) -> None:
+        """Set the database migration run status.
+
+        Args:
+            status: One of "PENDING", "COMPLETED", or "FAILED".
+        """
+        self._container.push(self._STATUS_FILE, source=status, make_dirs=True)
+
+    def get_completed_script(self) -> str | None:
+        """Get the database migration script that has completed in the current container.
+
+        Returns:
+            The completed database migration script in the current container.
+        """
+        if self._container.exists(self._COMPLETED_SCRIPT_FILE):
+            return cast(str, self._container.pull(self._COMPLETED_SCRIPT_FILE).read())
+        return None
+
+    def _set_completed_script(self, script_path: str) -> None:
+        """Set the database migration script that has completed in the current container.
+
+        Args:
+            script_path: The completed database migration script in the current container.
+        """
+        self._container.push(self._COMPLETED_SCRIPT_FILE, script_path, make_dirs=True)
+
+    def run(self, environment: dict[str, str]) -> None:
+        """Run the database migration script if database migration is still pending.
+
+        Args:
+            environment: Environment variables that's required for the run.
+
+        Raises:
+            CharmConfigInvalidError: if the database migration run failed.
+        """
+        if self.get_status() not in (
+            DatabaseMigrationStatus.PENDING,
+            DatabaseMigrationStatus.FAILED,
+        ):
+            return
+        if not self.script:
+            return
+        logger.info("execute database migration script: %s", repr(self.script))
+        try:
+            self._container.exec(
+                ["/bin/bash", "-xeo", "pipefail", self.script],
+                environment=environment,
+                working_dir=str(FLASK_APP_DIR),
+            ).wait()
+            self._set_status(DatabaseMigrationStatus.COMPLETED)
+            self._set_completed_script(self.script)
+        except ExecError as exc:
+            self._set_status(DatabaseMigrationStatus.FAILED)
+            logger.error(
+                "database migration script %s failed, stdout: %s, stderr: %s",
+                repr(self.script),
+                exc.stdout,
+                exc.stderr,
+            )
+            raise CharmConfigInvalidError(
+                f"database migration script {self.script!r} failed, "
+                "will retry in next update-status"
+            ) from exc

src/flask_app.py

@@ -10,6 +10,8 @@
 
 from charm_state import KNOWN_CHARM_CONFIG, CharmState
 from constants import FLASK_ENV_CONFIG_PREFIX, FLASK_SERVICE_NAME
+from database_migration import DatabaseMigration
+from exceptions import CharmConfigInvalidError
 from webserver import GunicornWebserver
 
 logger = logging.getLogger(__name__)
@@ -19,18 +21,24 @@ class FlaskApp:  # pylint: disable=too-few-public-methods
     """Flask application manager."""
 
     def __init__(
-        self, charm: ops.CharmBase, charm_state: CharmState, webserver: GunicornWebserver
+        self,
+        charm: ops.CharmBase,
+        charm_state: CharmState,
+        webserver: GunicornWebserver,
+        database_migration: DatabaseMigration,
     ):
         """Construct the FlaskApp instance.
 
         Args:
             charm: The main charm object.
             charm_state: The state of the charm.
             webserver: The webserver manager object.
+            database_migration: The database migration manager object.
         """
         self._charm = charm
         self._charm_state = charm_state
         self._webserver = webserver
+        self._database_migration = database_migration
 
     def _flask_environment(self) -> dict[str, str]:
         """Generate a Flask environment dictionary from the charm Flask configurations.
@@ -90,7 +98,8 @@ def _flask_layer(self) -> ops.pebble.LayerDict:
     def restart_flask(self) -> None:
         """Restart or start the flask service if not started with the latest configuration.
 
-        Raise CharmConfigInvalidError if the configuration is not valid.
+        Raises:
+             CharmConfigInvalidError: if the configuration is not valid.
         """
         container = self._charm.unit.get_container("flask-app")
         if not container.can_connect():
@@ -105,4 +114,15 @@ def restart_flask(self) -> None:
             flask_environment=self._flask_environment(),
             is_webserver_running=is_webserver_running,
         )
+        self._database_migration.run(self._flask_environment())
         container.replan()
+        if (
+            self._database_migration.get_completed_script() is not None
+            and self._database_migration.script is not None
+            and self._database_migration.script != self._database_migration.get_completed_script()
+        ):
+            raise CharmConfigInvalidError(
+                f"database migration script {self._database_migration.get_completed_script()!r} "
+                f"has been executed successfully in the current flask container,"
+                f"updating database-migration-script in config has no effect"
+            )

src/webserver.py

@@ -26,7 +26,11 @@ class GunicornWebserver:
 
     """
 
-    def __init__(self, charm_state: CharmState, flask_container: ops.Container):
+    def __init__(
+        self,
+        charm_state: CharmState,
+        flask_container: ops.Container,
+    ):
         """Initialize a new instance of the GunicornWebserver class.
 
         Args: