diff --git a/lib/charms/postgresql_k8s/v0/postgresql.py b/lib/charms/postgresql_k8s/v0/postgresql.py index 099047b01f..4eb23f3315 100644 --- a/lib/charms/postgresql_k8s/v0/postgresql.py +++ b/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 = 19 +LIBPATCH = 20 INVALID_EXTRA_USER_ROLE_BLOCKING_MESSAGE = "invalid role(s) for extra user roles" @@ -172,8 +172,7 @@ def create_database(self, database: str, user: str, plugins: List[str] = []) -> raise PostgreSQLCreateDatabaseError() # Enable preset extensions - for plugin in plugins: - self.enable_disable_extension(plugin, True, database) + self.enable_disable_extensions({plugin: True for plugin in plugins}, database) def create_user( self, user: str, password: str = None, admin: bool = False, extra_user_roles: str = None @@ -270,22 +269,16 @@ 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: + def enable_disable_extensions(self, extensions: Dict[str, 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. + extensions: the name of the extensions. 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: @@ -301,7 +294,12 @@ def enable_disable_extension(self, extension: str, enable: bool, database: str = with self._connect_to_database( database=database ) as connection, connection.cursor() as cursor: - cursor.execute(statement) + for extension, enable in extensions.items(): + cursor.execute( + f"CREATE EXTENSION IF NOT EXISTS {extension};" + if enable + else f"DROP EXTENSION IF EXISTS {extension};" + ) except psycopg2.errors.UniqueViolation: pass except psycopg2.Error: diff --git a/src/charm.py b/src/charm.py index 7c57423b24..aaec92aab9 100755 --- a/src/charm.py +++ b/src/charm.py @@ -552,24 +552,22 @@ def enable_disable_extensions(self, database: str = None) -> None: database: optional database where to enable/disable the extension. """ original_status = self.unit.status + extensions = {} + # collect extensions plugins_exception = {"uuid_ossp": '"uuid-ossp"'} for plugin in self.config.plugin_keys(): enable = self.config[plugin] # Enable or disable the plugin/extension. extension = "_".join(plugin.split("_")[1:-1]) - if extension in plugins_exception: - extension = plugins_exception[extension] - 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 + extension = plugins_exception.get(extension, extension) + extensions[extension] = enable + self.unit.status = WaitingStatus("Updating extensions") + try: + self.postgresql.enable_disable_extensions(extensions, database) + except PostgreSQLEnableDisableExtensionError as e: + logger.exception("failed to change plugins: %s", str(e)) + self.unit.status = original_status def _add_members(self, event) -> None: """Add new cluster members. diff --git a/src/relations/db.py b/src/relations/db.py index 3261351392..48f073c460 100644 --- a/src/relations/db.py +++ b/src/relations/db.py @@ -152,9 +152,6 @@ def set_up_relation(self, relation: Relation) -> bool: self.charm.postgresql.create_database(database, user, plugins=plugins) - # Enable/disable extensions in the new database. - self.charm.enable_disable_extensions(database) - # Build the primary's connection string. primary = str( ConnectionString( diff --git a/tests/unit/test_db.py b/tests/unit/test_db.py index dafa52d226..e80f6b0df9 100644 --- a/tests/unit/test_db.py +++ b/tests/unit/test_db.py @@ -183,14 +183,12 @@ def test_get_extensions(self): ) @patch("relations.db.DbProvides._update_unit_status") - @patch("charm.PostgresqlOperatorCharm.enable_disable_extensions") @patch("relations.db.new_password", return_value="test-password") @patch("relations.db.DbProvides._get_extensions") def test_set_up_relation( self, _get_extensions, _new_password, - _enable_disable_extensions, _update_unit_status, ): with patch.object(PostgresqlOperatorCharm, "postgresql", Mock()) as postgresql_mock: @@ -228,7 +226,6 @@ def test_set_up_relation( self.assertFalse(self.harness.charm.legacy_db_relation.set_up_relation(relation)) postgresql_mock.create_user.assert_not_called() postgresql_mock.create_database.assert_not_called() - _enable_disable_extensions.assert_not_called() postgresql_mock.get_postgresql_version.assert_not_called() _update_unit_status.assert_not_called() @@ -244,7 +241,6 @@ def test_set_up_relation( user = f"relation_id_{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, plugins=[]) - _enable_disable_extensions.assert_called_once() self.assertEqual(postgresql_mock.get_postgresql_version.call_count, 2) _update_unit_status.assert_called_once() expected_data = { @@ -271,7 +267,6 @@ def test_set_up_relation( # provided only in the unit databag. postgresql_mock.create_user.reset_mock() postgresql_mock.create_database.reset_mock() - _enable_disable_extensions.reset_mock() postgresql_mock.get_postgresql_version.reset_mock() _update_unit_status.reset_mock() with self.harness.hooks_disabled(): @@ -289,7 +284,6 @@ 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, plugins=[]) - _enable_disable_extensions.assert_called_once() self.assertEqual(postgresql_mock.get_postgresql_version.call_count, 2) _update_unit_status.assert_called_once() self.assertEqual(self.harness.get_relation_data(self.rel_id, self.app), expected_data) @@ -299,7 +293,6 @@ def test_set_up_relation( # Assert that the correct calls were made when the database name is not provided. postgresql_mock.create_user.reset_mock() postgresql_mock.create_database.reset_mock() - _enable_disable_extensions.reset_mock() postgresql_mock.get_postgresql_version.reset_mock() _update_unit_status.reset_mock() with self.harness.hooks_disabled(): @@ -312,7 +305,6 @@ def test_set_up_relation( self.assertFalse(self.harness.charm.legacy_db_relation.set_up_relation(relation)) postgresql_mock.create_user.assert_not_called() postgresql_mock.create_database.assert_not_called() - _enable_disable_extensions.assert_not_called() postgresql_mock.get_postgresql_version.assert_not_called() _update_unit_status.assert_not_called() # No data is set in the databags by the database. @@ -329,7 +321,6 @@ def test_set_up_relation( ) self.assertFalse(self.harness.charm.legacy_db_relation.set_up_relation(relation)) postgresql_mock.create_database.assert_not_called() - _enable_disable_extensions.assert_not_called() postgresql_mock.get_postgresql_version.assert_not_called() _update_unit_status.assert_not_called() self.assertIsInstance(self.harness.model.unit.status, BlockedStatus) @@ -340,7 +331,6 @@ def test_set_up_relation( # BlockedStatus due to a PostgreSQLCreateDatabaseError. self.harness.charm.unit.status = ActiveStatus() self.assertFalse(self.harness.charm.legacy_db_relation.set_up_relation(relation)) - _enable_disable_extensions.assert_not_called() postgresql_mock.get_postgresql_version.assert_not_called() _update_unit_status.assert_not_called() self.assertIsInstance(self.harness.model.unit.status, BlockedStatus)