[DPE-2569] Extensions not enabled when database is created (#252)

* Extensions not enabled when database is created

* Bump agent version

* Bump lib

Author: dragomirp

.github/workflows/ci.yaml

@@ -86,15 +86,15 @@ jobs:
           - upgrade-from-stable-integration
         agent-versions:
           - "2.9.45"  # renovate: latest juju 2
-          - "3.1.5"  # renovate: latest juju 3
+          - "3.1.6"  # renovate: latest juju 3
         free-disk-space:
           - false
         include:
           - tox-environments: database-relation-integration
             agent-versions: "2.9.45"  # renovate: latest juju 2
             free-disk-space: true
           - tox-environments: database-relation-integration
-            agent-versions: "3.1.5"  # renovate: latest juju 3
+            agent-versions: "3.1.6"  # renovate: latest juju 3
             free-disk-space: true
           - tox-environments: db-admin-relation-integration
             agent-versions: "2.9.45"  # renovate: latest juju 2

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 = 16
+LIBPATCH = 17
 
 INVALID_EXTRA_USER_ROLE_BLOCKING_MESSAGE = "invalid role(s) for extra user roles"
 
@@ -117,12 +117,13 @@ def _connect_to_database(
         connection.autocommit = True
         return connection
 
-    def create_database(self, database: str, user: str) -> None:
+    def create_database(self, database: str, user: str, plugins: List[str] = []) -> None:
         """Creates a new database and grant privileges to a user on it.
 
         Args:
             database: database to be created.
             user: user that will have access to the database.
+            plugins: extensions to enable in the new database.
         """
         try:
             connection = self._connect_to_database()
@@ -170,6 +171,10 @@ def create_database(self, database: str, user: str) -> None:
             logger.error(f"Failed to create database: {e}")
             raise PostgreSQLCreateDatabaseError()
 
+        # Enable preset extensions
+        for plugin in plugins:
+            self.enable_disable_extension(plugin, True, database)
+
     def create_user(
         self, user: str, password: str = None, admin: bool = False, extra_user_roles: str = None
     ) -> None:

src/charm.py

@@ -930,17 +930,22 @@ def enable_disable_extensions(self, database: str = None) -> None:
         Args:
             database: optional database where to enable/disable the extension.
         """
+        original_status = self.unit.status
         for plugin in self.config.plugin_keys():
             enable = self.config[plugin]
 
             # Enable or disable the plugin/extension.
             extension = "_".join(plugin.split("_")[1:-1])
+            self.unit.status = WaitingStatus(
+                f"{'Enabling' if enable else 'Disabling'} {extension}"
+            )
             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)
                 )
+            self.unit.status = original_status
 
     def _get_ips_to_remove(self) -> Set[str]:
         """List the IPs that were part of the cluster but departed."""

src/relations/db.py

@@ -168,7 +168,13 @@ def set_up_relation(self, relation: Relation) -> bool:
             self.charm.set_secret(APP_SCOPE, f"{user}-database", database)
 
             self.charm.postgresql.create_user(user, password, self.admin)
-            self.charm.postgresql.create_database(database, user)
+            plugins = [
+                "_".join(plugin.split("_")[1:-1])
+                for plugin in self.charm.config.plugin_keys()
+                if self.charm.config[plugin]
+            ]
+
+            self.charm.postgresql.create_database(database, user, plugins=plugins)
 
             # Enable/disable extensions in the new database.
             self.charm.enable_disable_extensions(database)

src/relations/postgresql_provider.py

@@ -84,7 +84,13 @@ def _on_database_requested(self, event: DatabaseRequestedEvent) -> None:
             user = f"relation-{event.relation.id}"
             password = new_password()
             self.charm.postgresql.create_user(user, password, extra_user_roles=extra_user_roles)
-            self.charm.postgresql.create_database(database, user)
+            plugins = [
+                "_".join(plugin.split("_")[1:-1])
+                for plugin in self.charm.config.plugin_keys()
+                if self.charm.config[plugin]
+            ]
+
+            self.charm.postgresql.create_database(database, user, plugins=plugins)
 
             # Share the credentials with the application.
             self.database_provides.set_credentials(event.relation.id, user, password)

tests/unit/test_db.py

@@ -236,7 +236,7 @@ def test_set_up_relation(
             self.assertTrue(self.harness.charm.legacy_db_relation.set_up_relation(relation))
             user = f"relation-{self.rel_id}"
             postgresql_mock.create_user.assert_called_once_with(user, "test-password", False)
-            postgresql_mock.create_database.assert_called_once_with(DATABASE, user)
+            postgresql_mock.create_database.assert_called_once_with(DATABASE, user, plugins=[])
             _enable_disable_extensions.assert_called_once()
             _update_endpoints.assert_called_once()
             _update_unit_status.assert_called_once()
@@ -263,7 +263,7 @@ def test_set_up_relation(
                 )
             self.assertTrue(self.harness.charm.legacy_db_relation.set_up_relation(relation))
             postgresql_mock.create_user.assert_called_once_with(user, "test-password", False)
-            postgresql_mock.create_database.assert_called_once_with(DATABASE, user)
+            postgresql_mock.create_database.assert_called_once_with(DATABASE, user, plugins=[])
             _enable_disable_extensions.assert_called_once()
             _update_endpoints.assert_called_once()
             _update_unit_status.assert_called_once()
@@ -284,7 +284,9 @@ def test_set_up_relation(
                 )
             self.assertTrue(self.harness.charm.legacy_db_relation.set_up_relation(relation))
             postgresql_mock.create_user.assert_called_once_with(user, "test-password", False)
-            postgresql_mock.create_database.assert_called_once_with("application", user)
+            postgresql_mock.create_database.assert_called_once_with(
+                "application", user, plugins=[]
+            )
             _enable_disable_extensions.assert_called_once()
             _update_endpoints.assert_called_once()
             _update_unit_status.assert_called_once()

tests/unit/test_postgresql_provider.py

@@ -123,7 +123,7 @@ def test_on_database_requested(
             postgresql_mock.create_user.assert_called_once_with(
                 user, "test-password", extra_user_roles=EXTRA_USER_ROLES
             )
-            postgresql_mock.create_database.assert_called_once_with(DATABASE, user)
+            postgresql_mock.create_database.assert_called_once_with(DATABASE, user, plugins=[])
             postgresql_mock.get_postgresql_version.assert_called_once()
             _update_endpoints.assert_called_once()