Toggle plugins in one go (canonical#322)

Author: dragomirp

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:

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.

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(

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)