[DPE-1372] Add extensions/plugins support (#170)

* Add plugins support

* Improve db relation test

* Implement more robust tests with Indico

* Fix list of requested extensions

* Remove unstable mark

* Use custom application name in test

* Update based on PR feedback

* Improve code for CI

* Add helper to wait for relation removal

* Fix change of the order of the requested extensions

Author: marceloneppel

.github/workflows/ci.yaml

@@ -71,6 +71,7 @@ jobs:
       matrix:
         tox-environments:
           - password-rotation-integration
+          - plugins-integration
           - tls-integration
         agents-version: ["2.9.42"]  # renovate: latest
         include:

config.yaml

@@ -0,0 +1,28 @@
+# Copyright 2023 Canonical Ltd.
+# See LICENSE file for licensing details.
+
+options:
+  plugin_citext_enable:
+    default: false
+    type: boolean
+    description: Enable citext extension
+  plugin_debversion_enable:
+    default: false
+    type: boolean
+    description: Enable debversion extension
+  plugin_hstore_enable:
+    default: false
+    type: boolean
+    description: Enable hstore extension
+  plugin_pg_trgm_enable:
+    default: false
+    type: boolean
+    description: Enable pg_trgm extension
+  plugin_plpython3u_enable:
+    default: false
+    type: boolean
+    description: Enable PL/Python extension
+  plugin_unaccent_enable:
+    default: false
+    type: boolean
+    description: Enable unaccent extension

lib/charms/postgresql_k8s/v0/postgresql.py

@@ -32,7 +32,7 @@
 
 # Increment this PATCH version before using `charmcraft publish-lib` or reset
 # to 0 if you are raising the major API version
-LIBPATCH = 8
+LIBPATCH = 9
 
 
 logger = logging.getLogger(__name__)
@@ -50,6 +50,10 @@ class PostgreSQLDeleteUserError(Exception):
     """Exception raised when deleting a user fails."""
 
 
+class PostgreSQLEnableDisableExtensionError(Exception):
+    """Exception raised when enabling/disabling an extension fails."""
+
+
 class PostgreSQLGetPostgreSQLVersionError(Exception):
     """Exception raised when retrieving PostgreSQL version fails."""
 
@@ -237,6 +241,46 @@ def delete_user(self, user: str) -> None:
             logger.error(f"Failed to delete user: {e}")
             raise PostgreSQLDeleteUserError()
 
+    def enable_disable_extension(self, extension: str, enable: bool, database: str = None) -> None:
+        """Enables or disables a PostgreSQL extension.
+
+        Args:
+            extension: the name of the extensions.
+            enable: whether the extension should be enabled or disabled.
+            database: optional database where to enable/disable the extension.
+
+        Raises:
+            PostgreSQLEnableDisableExtensionError if the operation fails.
+        """
+        statement = (
+            f"CREATE EXTENSION IF NOT EXISTS {extension};"
+            if enable
+            else f"DROP EXTENSION IF EXISTS {extension};"
+        )
+        connection = None
+        try:
+            if database is not None:
+                databases = [database]
+            else:
+                # Retrieve all the databases.
+                with self._connect_to_database() as connection, connection.cursor() as cursor:
+                    cursor.execute("SELECT datname FROM pg_database WHERE NOT datistemplate;")
+                    databases = {database[0] for database in cursor.fetchall()}
+
+            # Enable/disabled the extension in each database.
+            for database in databases:
+                with self._connect_to_database(
+                    database=database
+                ) as connection, connection.cursor() as cursor:
+                    cursor.execute(statement)
+        except psycopg2.errors.UniqueViolation:
+            pass
+        except psycopg2.Error:
+            raise PostgreSQLEnableDisableExtensionError()
+        finally:
+            if connection is not None:
+                connection.close()
+
     def get_postgresql_version(self) -> str:
         """Returns the PostgreSQL version.
 

src/charm.py

@@ -12,6 +12,7 @@
 from charms.observability_libs.v1.kubernetes_service_patch import KubernetesServicePatch
 from charms.postgresql_k8s.v0.postgresql import (
     PostgreSQL,
+    PostgreSQLEnableDisableExtensionError,
     PostgreSQLUpdateUserPasswordError,
 )
 from charms.postgresql_k8s.v0.postgresql_tls import PostgreSQLTLS
@@ -63,7 +64,7 @@
     WORKLOAD_OS_USER,
 )
 from patroni import NotReadyError, Patroni
-from relations.db import DbProvides
+from relations.db import EXTENSIONS_BLOCKING_MESSAGE, DbProvides
 from relations.postgresql_provider import PostgreSQLProvider
 from utils import new_password
 
@@ -165,6 +166,11 @@ def set_secret(self, scope: str, key: str, value: Optional[str]) -> None:
         else:
             raise RuntimeError("Unknown secret scope.")
 
+    @property
+    def is_cluster_initialised(self) -> bool:
+        """Returns whether the cluster is already initialised."""
+        return "cluster_initialised" in self.app_peer_data
+
     @property
     def postgresql(self) -> PostgreSQL:
         """Returns an instance of the object used to interact with the database."""
@@ -286,9 +292,51 @@ def _on_peer_relation_changed(self, event: RelationChangedEvent) -> None:
             self.unit.status = ActiveStatus()
 
     def _on_config_changed(self, _) -> None:
-        """Handle the config-changed event."""
-        # TODO: placeholder method to implement logic specific to configuration change.
-        pass
+        """Handle configuration changes, like enabling plugins."""
+        if not self.is_cluster_initialised:
+            logger.debug("Early exit on_config_changed: cluster not initialised yet")
+            return
+
+        if not self.unit.is_leader():
+            return
+
+        # Enable and/or disable the extensions.
+        self.enable_disable_extensions()
+
+        # Unblock the charm after extensions are enabled (only if it's blocked due to application
+        # charms requesting extensions).
+        if self.unit.status.message != EXTENSIONS_BLOCKING_MESSAGE:
+            return
+
+        for relation in [
+            *self.model.relations.get("db", []),
+            *self.model.relations.get("db-admin", []),
+        ]:
+            if not self.legacy_db_relation.set_up_relation(relation):
+                logger.debug(
+                    "Early exit on_config_changed: legacy relation requested extensions that are still disabled"
+                )
+                return
+
+    def enable_disable_extensions(self, database: str = None) -> None:
+        """Enable/disable PostgreSQL extensions set through config options.
+
+        Args:
+            database: optional database where to enable/disable the extension.
+        """
+        for config, enable in self.model.config.items():
+            # Filter config option not related to plugins.
+            if not config.startswith("plugin_"):
+                continue
+
+            # Enable or disable the plugin/extension.
+            extension = "_".join(config.split("_")[1:-1])
+            try:
+                self.postgresql.enable_disable_extension(extension, enable, database)
+            except PostgreSQLEnableDisableExtensionError as e:
+                logger.exception(
+                    f"failed to {'enable' if enable else 'disable'} {extension} plugin: %s", str(e)
+                )
 
     def _add_members(self, event) -> None:
         """Add new cluster members.
@@ -463,6 +511,10 @@ def _on_postgresql_pebble_ready(self, event: WorkloadEvent) -> None:
         # Update the archive command and replication configurations.
         self.update_config()
 
+        # Enable/disable PostgreSQL extensions if they were set before the cluster
+        # was fully initialised.
+        self.enable_disable_extensions()
+
         # All is well, set an ActiveStatus.
         self.unit.status = ActiveStatus()
 

src/relations/db.py

@@ -5,7 +5,7 @@
 
 
 import logging
-from typing import Iterable
+from typing import Iterable, List, Set, Tuple
 
 from charms.postgresql_k8s.v0.postgresql import (
     PostgreSQLCreateDatabaseError,
@@ -89,34 +89,59 @@ def _on_relation_changed(self, event: RelationChangedEvent) -> None:
 
         logger.warning(f"DEPRECATION WARNING - `{self.relation_name}` is a legacy interface")
 
-        unit_relation_databag = event.relation.data[self.charm.unit]
-        application_relation_databag = event.relation.data[self.charm.app]
+        self.set_up_relation(event.relation)
 
-        # Do not allow apps requesting extensions to be installed.
-        if "extensions" in event.relation.data.get(
-            event.app, {}
-        ) or "extensions" in event.relation.data.get(event.unit, {}):
+    def _get_extensions(self, relation: Relation) -> Tuple[List, Set]:
+        """Returns the list of required and disabled extensions."""
+        requested_extensions = relation.data.get(relation.app, {}).get("extensions", "").split(",")
+        for unit in relation.units:
+            requested_extensions.extend(
+                relation.data.get(unit, {}).get("extensions", "").split(",")
+            )
+        required_extensions = []
+        for extension in requested_extensions:
+            if extension != "" and extension not in required_extensions:
+                required_extensions.append(extension)
+        disabled_extensions = set()
+        if required_extensions:
+            for extension in required_extensions:
+                extension_name = extension.split(":")[0]
+                if not self.charm.model.config.get(f"plugin_{extension_name}_enable"):
+                    disabled_extensions.add(extension_name)
+        return required_extensions, disabled_extensions
+
+    def set_up_relation(self, relation: Relation) -> bool:
+        """Set up the relation to be used by the application charm."""
+        # Do not allow apps requesting extensions to be installed
+        # (let them now about config options).
+        required_extensions, disabled_extensions = self._get_extensions(relation)
+        if disabled_extensions:
             logger.error(
-                "ERROR - `extensions` cannot be requested through relations"
-                " - they should be installed through a database charm config in the future"
+                f"ERROR - `extensions` ({', '.join(disabled_extensions)}) cannot be requested through relations"
+                " - Please enable extensions through `juju config` and add the relation again."
             )
             self.charm.unit.status = BlockedStatus(EXTENSIONS_BLOCKING_MESSAGE)
-            return
+            return False
 
-        # Sometimes a relation changed event is triggered,
-        # and it doesn't have a database name in it.
-        database = event.relation.data.get(event.app, {}).get(
-            "database", event.relation.data.get(event.unit, {}).get("database")
-        )
+        database = relation.data.get(relation.app, {}).get("database")
         if not database:
-            logger.warning("Deferring on_relation_changed: No database name provided")
-            event.defer()
-            return
+            for unit in relation.units:
+                unit_database = relation.data.get(unit, {}).get("database")
+                if unit_database:
+                    database = unit_database
+                    break
+
+        if not database:
+            logger.warning("Early exit on_relation_changed: No database name provided")
+            return False
 
         try:
+            unit_relation_databag = relation.data[self.charm.unit]
+            application_relation_databag = relation.data[self.charm.app]
+
             # Creates the user and the database for this specific relation if it was not already
             # created in a previous relation changed event.
-            user = f"relation_id_{event.relation.id}"
+            user = f"relation_id_{relation.id}"
             password = unit_relation_databag.get("password", new_password())
             self.charm.postgresql.create_user(user, password, self.admin)
             self.charm.postgresql.create_database(database, user)
@@ -129,7 +154,7 @@ def _on_relation_changed(self, event: RelationChangedEvent) -> None:
                     port=DATABASE_PORT,
                     user=user,
                     password=password,
-                    fallback_application_name=event.app.name,
+                    fallback_application_name=relation.app.name,
                 )
             )
 
@@ -141,7 +166,7 @@ def _on_relation_changed(self, event: RelationChangedEvent) -> None:
                     port=DATABASE_PORT,
                     user=user,
                     password=password,
-                    fallback_application_name=event.app.name,
+                    fallback_application_name=relation.app.name,
                 )
             )
 
@@ -153,8 +178,8 @@ def _on_relation_changed(self, event: RelationChangedEvent) -> None:
             # application and this charm will not work.
             for databag in [application_relation_databag, unit_relation_databag]:
                 updates = {
-                    "allowed-subnets": self._get_allowed_subnets(event.relation),
-                    "allowed-units": self._get_allowed_units(event.relation),
+                    "allowed-subnets": self._get_allowed_subnets(relation),
+                    "allowed-units": self._get_allowed_units(relation),
                     "host": self.charm.endpoint,
                     "master": primary,
                     "port": DATABASE_PORT,
@@ -163,6 +188,7 @@ def _on_relation_changed(self, event: RelationChangedEvent) -> None:
                     "user": user,
                     "password": password,
                     "database": database,
+                    "extensions": ",".join(required_extensions),
                 }
                 databag.update(updates)
         except (
@@ -173,7 +199,9 @@ def _on_relation_changed(self, event: RelationChangedEvent) -> None:
             self.charm.unit.status = BlockedStatus(
                 f"Failed to initialize {self.relation_name} relation"
             )
-            return
+            return False
+
+        return True
 
     def _check_for_blocking_relations(self, relation_id: int) -> bool:
         """Checks if there are relations with extensions.

tests/integration/helpers.py

@@ -23,7 +23,9 @@
     retry_if_exception,
     retry_if_result,
     stop_after_attempt,
+    stop_after_delay,
     wait_exponential,
+    wait_fixed,
 )
 
 CHARM_SERIES = "jammy"
@@ -529,6 +531,15 @@ async def check_tls_patroni_api(ops_test: OpsTest, unit_name: str, enabled: bool
         return False
 
 
+def has_relation_exited(ops_test: OpsTest, endpoint_one: str, endpoint_two: str) -> bool:
+    """Returns true if the relation between endpoint_one and endpoint_two has been removed."""
+    for rel in ops_test.model.relations:
+        endpoints = [endpoint.name for endpoint in rel.endpoints]
+        if endpoint_one not in endpoints and endpoint_two not in endpoints:
+            return True
+    return False
+
+
 @retry(
     retry=retry_if_result(lambda x: not x),
     stop=stop_after_attempt(10),
@@ -640,3 +651,22 @@ async def wait_for_idle_on_blocked(
         ),
         ops_test.model.block_until(lambda: unit.workload_status_message == status_message),
     )
+
+
+def wait_for_relation_removed_between(
+    ops_test: OpsTest, endpoint_one: str, endpoint_two: str
+) -> None:
+    """Wait for relation to be removed before checking if it's waiting or idle.
+
+    Args:
+        ops_test: running OpsTest instance
+        endpoint_one: one endpoint of the relation. Doesn't matter if it's provider or requirer.
+        endpoint_two: the other endpoint of the relation.
+    """
+    try:
+        for attempt in Retrying(stop=stop_after_delay(3 * 60), wait=wait_fixed(3)):
+            with attempt:
+                if has_relation_exited(ops_test, endpoint_one, endpoint_two):
+                    break
+    except RetryError:
+        assert False, "Relation failed to exit after 3 minutes."