diff --git a/lib/charms/postgresql_k8s/v0/postgresql.py b/lib/charms/postgresql_k8s/v0/postgresql.py index bfda780e8c..199ae3fec5 100644 --- a/lib/charms/postgresql_k8s/v0/postgresql.py +++ b/lib/charms/postgresql_k8s/v0/postgresql.py @@ -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: diff --git a/src/relations/db.py b/src/relations/db.py index bc63529e5a..cd59d06aaa 100644 --- a/src/relations/db.py +++ b/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) diff --git a/src/relations/postgresql_provider.py b/src/relations/postgresql_provider.py index 9824061e5a..99661dcd62 100644 --- a/src/relations/postgresql_provider.py +++ b/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) diff --git a/tests/unit/test_db.py b/tests/unit/test_db.py index 7eb2f2b9c4..2d240d4ede 100644 --- a/tests/unit/test_db.py +++ b/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,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("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() diff --git a/tests/unit/test_postgresql_provider.py b/tests/unit/test_postgresql_provider.py index 5bf5321fcb..28d20093c3 100644 --- a/tests/unit/test_postgresql_provider.py +++ b/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()