Toggle plugins in one go (#279)

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

@@ -948,21 +948,20 @@ 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
         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
+            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 _get_ips_to_remove(self) -> Set[str]:
         """List the IPs that were part of the cluster but departed."""

src/relations/db.py

@@ -176,8 +176,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)
         except (PostgreSQLCreateDatabaseError, PostgreSQLCreateUserError) as e:
             logger.exception(e)
             self.charm.unit.status = BlockedStatus(

tests/integration/test_plugins.py

@@ -38,7 +38,7 @@ async def test_plugins(ops_test: OpsTest) -> None:
             series=CHARM_SERIES,
             config={"profile": "testing"},
         )
-        await ops_test.model.wait_for_idle(apps=[DATABASE_APP_NAME], status="active")
+        await ops_test.model.wait_for_idle(apps=[DATABASE_APP_NAME], status="active", timeout=1000)
 
     # Check that the available plugins are disabled.
     primary = await get_primary(ops_test, f"{DATABASE_APP_NAME}/0")

tests/unit/test_charm.py

@@ -288,22 +288,22 @@ def test_enable_disable_extensions(self, _):
             postgresql_mock.enable_disable_extension.side_effect = None
             with self.assertNoLogs("charm", "ERROR"):
                 self.charm.enable_disable_extensions()
-                self.assertEqual(postgresql_mock.enable_disable_extension.call_count, 6)
+                self.assertEqual(postgresql_mock.enable_disable_extensions.call_count, 1)
 
             # Test when one extension install/uninstall fails.
             postgresql_mock.reset_mock()
-            postgresql_mock.enable_disable_extension.side_effect = (
+            postgresql_mock.enable_disable_extensions.side_effect = (
                 PostgreSQLEnableDisableExtensionError
             )
             with self.assertLogs("charm", "ERROR") as logs:
                 self.charm.enable_disable_extensions()
-                self.assertEqual(postgresql_mock.enable_disable_extension.call_count, 6)
-                self.assertIn("failed to disable citext plugin", "".join(logs.output))
+                self.assertEqual(postgresql_mock.enable_disable_extensions.call_count, 1)
+                self.assertIn("failed to change plugins", "".join(logs.output))
 
             # Test when one config option should be skipped (because it's not related
             # to a plugin/extension).
             postgresql_mock.reset_mock()
-            postgresql_mock.enable_disable_extension.side_effect = None
+            postgresql_mock.enable_disable_extensions.side_effect = None
             with self.assertNoLogs("charm", "ERROR"):
                 config = """options:
   plugin_citext_enable:
@@ -331,7 +331,7 @@ def test_enable_disable_extensions(self, _):
                 self.addCleanup(harness.cleanup)
                 harness.begin()
                 harness.charm.enable_disable_extensions()
-                self.assertEqual(postgresql_mock.enable_disable_extension.call_count, 6)
+                self.assertEqual(postgresql_mock.enable_disable_extensions.call_count, 1)
 
     @patch("charm.PostgresqlOperatorCharm.enable_disable_extensions")
     @patch("charm.snap.SnapCache")

tests/unit/test_db.py

@@ -185,14 +185,12 @@ def test_get_extensions(self):
     @patch("subprocess.check_output", return_value=b"C")
     @patch("relations.db.DbProvides._update_unit_status")
     @patch("relations.db.DbProvides.update_endpoints")
-    @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_endpoints,
         _update_unit_status,
         _,
@@ -222,7 +220,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_endpoints.assert_not_called()
             _update_unit_status.assert_not_called()
@@ -239,7 +236,6 @@ def test_set_up_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, plugins=[])
-            _enable_disable_extensions.assert_called_once()
             _update_endpoints.assert_called_once()
             _update_unit_status.assert_called_once()
             self.assertNotIsInstance(self.harness.model.unit.status, BlockedStatus)
@@ -248,7 +244,6 @@ def test_set_up_relation(
             # provided in both application and unit databags.
             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_endpoints.reset_mock()
             _update_unit_status.reset_mock()
@@ -266,15 +261,13 @@ 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()
             _update_endpoints.assert_called_once()
             _update_unit_status.assert_called_once()
             self.assertNotIsInstance(self.harness.model.unit.status, BlockedStatus)
 
             # 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_endpoints.reset_mock()
             _update_unit_status.reset_mock()
@@ -289,28 +282,24 @@ def test_set_up_relation(
             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()
             self.assertNotIsInstance(self.harness.model.unit.status, BlockedStatus)
 
             # BlockedStatus due to a PostgreSQLCreateUserError.
             postgresql_mock.create_database.reset_mock()
-            _enable_disable_extensions.reset_mock()
             postgresql_mock.get_postgresql_version.reset_mock()
             _update_endpoints.reset_mock()
             _update_unit_status.reset_mock()
             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()
             _update_endpoints.assert_not_called()
             _update_unit_status.assert_not_called()
             self.assertIsInstance(self.harness.model.unit.status, BlockedStatus)
 
             # 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()
             _update_endpoints.assert_not_called()
             _update_unit_status.assert_not_called()
             self.assertIsInstance(self.harness.model.unit.status, BlockedStatus)