[DPE-4801] Block charm if plugin disable fails due to dependent objec…

…ts (#567)

* block charm if plugin disable fails

* nits

* address issues

* refactor functions

* add tests

lib/charms/postgresql_k8s/v0/postgresql.py

@@ -37,7 +37,7 @@
 
 # Increment this PATCH version before using `charmcraft publish-lib` or reset
 # to 0 if you are raising the major API version
-LIBPATCH = 31
+LIBPATCH = 32
 
 INVALID_EXTRA_USER_ROLE_BLOCKING_MESSAGE = "invalid role(s) for extra user roles"
 
@@ -330,6 +330,8 @@ def enable_disable_extensions(self, extensions: Dict[str, bool], database: str =
                         )
         except psycopg2.errors.UniqueViolation:
             pass
+        except psycopg2.errors.DependentObjectsStillExist:
+            raise
         except psycopg2.Error:
             raise PostgreSQLEnableDisableExtensionError()
         finally:

src/charm.py

@@ -99,6 +99,7 @@
 logger = logging.getLogger(__name__)
 
 EXTENSIONS_DEPENDENCY_MESSAGE = "Unsatisfied plugin dependencies. Please check the logs"
+EXTENSION_OBJECT_MESSAGE = "Cannot disable plugins: Existing objects depend on it. See logs"
 
 # http{x,core} clutter the logs with debug messages
 logging.getLogger("httpcore").setLevel(logging.ERROR)
@@ -633,12 +634,27 @@ def enable_disable_extensions(self, database: str = None) -> None:
         if self.is_blocked and self.unit.status.message == EXTENSIONS_DEPENDENCY_MESSAGE:
             self._set_active_status()
             original_status = self.unit.status
+
+        self._handle_enable_disable_extensions(original_status, extensions, database)
+
+    def _handle_enable_disable_extensions(self, original_status, extensions, database) -> None:
+        """Try enablind/disabling Postgresql extensions and handle exceptions appropriately."""
         if not isinstance(original_status, UnknownStatus):
             self.unit.status = WaitingStatus("Updating extensions")
         try:
             self.postgresql.enable_disable_extensions(extensions, database)
+        except psycopg2.errors.DependentObjectsStillExist as e:
+            logger.error(
+                "Failed to disable plugin: %s\nWas the plugin enabled manually? If so, update charm config with `juju config postgresql-k8s plugin_<plugin_name>_enable=True`",
+                str(e),
+            )
+            self.unit.status = BlockedStatus(EXTENSION_OBJECT_MESSAGE)
+            return
         except PostgreSQLEnableDisableExtensionError as e:
             logger.exception("failed to change plugins: %s", str(e))
+        if original_status.message == EXTENSION_OBJECT_MESSAGE:
+            self._set_active_status()
+            return
         if not isinstance(original_status, UnknownStatus):
             self.unit.status = original_status
 
@@ -1239,6 +1255,9 @@ def _on_update_status(self, _) -> None:
             return
 
         if self._has_blocked_status or self._has_waiting_status:
+            # If charm was failing to disable plugin, try again (user may have removed the objects)
+            if self.unit.status.message == EXTENSION_OBJECT_MESSAGE:
+                self.enable_disable_extensions()
             logger.debug("on_update_status early exit: Unit is in Blocked/Waiting status")
             return
 
@@ -1248,25 +1267,10 @@ def _on_update_status(self, _) -> None:
             logger.debug("on_update_status early exit: Service has not been added nor started yet")
             return
 
-        if "restoring-backup" in self.app_peer_data:
-            if services[0].current != ServiceStatus.ACTIVE:
-                logger.error("Restore failed: database service failed to start")
-                self.unit.status = BlockedStatus("Failed to restore backup")
-                return
-
-            if not self._patroni.member_started:
-                logger.debug("on_update_status early exit: Patroni has not started yet")
-                return
-
-            # Remove the restoring backup flag and the restore stanza name.
-            self.app_peer_data.update({"restoring-backup": "", "restore-stanza": ""})
-            self.update_config()
-            logger.info("Restore succeeded")
-
-            can_use_s3_repository, validation_message = self.backup.can_use_s3_repository()
-            if not can_use_s3_repository:
-                self.unit.status = BlockedStatus(validation_message)
-                return
+        if "restoring-backup" in self.app_peer_data and not self._was_restore_successful(
+            services[0]
+        ):
+            return
 
         if self._handle_processes_failures():
             return
@@ -1276,6 +1280,29 @@ def _on_update_status(self, _) -> None:
 
         self._set_active_status()
 
+    def _was_restore_successful(self, service) -> bool:
+        """Checks if restore operation succeeded and S3 is properly configured."""
+        if service.current != ServiceStatus.ACTIVE:
+            logger.error("Restore failed: database service failed to start")
+            self.unit.status = BlockedStatus("Failed to restore backup")
+            return False
+
+        if not self._patroni.member_started:
+            logger.debug("on_update_status early exit: Patroni has not started yet")
+            return False
+
+        # Remove the restoring backup flag and the restore stanza name.
+        self.app_peer_data.update({"restoring-backup": "", "restore-stanza": ""})
+        self.update_config()
+        logger.info("Restore succeeded")
+
+        can_use_s3_repository, validation_message = self.backup.can_use_s3_repository()
+        if not can_use_s3_repository:
+            self.unit.status = BlockedStatus(validation_message)
+            return False
+
+        return True
+
     def _handle_processes_failures(self) -> bool:
         """Handle Patroni and PostgreSQL OS processes failures.
 

tests/integration/test_plugins.py

@@ -201,3 +201,55 @@ def enable_disable_config(enabled: False):
             else:
                 connection.cursor().execute(query)
     connection.close()
+
+
+@pytest.mark.group(1)
+@pytest.mark.abort_on_fail
+async def test_plugin_objects(ops_test: OpsTest) -> None:
+    """Checks if charm gets blocked when trying to disable a plugin in use."""
+    primary = await get_primary(ops_test)
+    password = await get_password(ops_test)
+    address = await get_unit_address(ops_test, primary)
+
+    logger.info("Creating an index object which depends on the pg_trgm config")
+    with db_connect(host=address, password=password) as connection:
+        connection.autocommit = True
+        connection.cursor().execute("CREATE TABLE example (value VARCHAR(10));")
+        connection.cursor().execute(
+            "CREATE INDEX example_idx ON example USING gin(value gin_trgm_ops);"
+        )
+    connection.close()
+
+    logger.info("Disabling the plugin on charm config, waiting for blocked status")
+    await ops_test.model.applications[DATABASE_APP_NAME].set_config({
+        "plugin_pg_trgm_enable": "False"
+    })
+    await ops_test.model.block_until(
+        lambda: ops_test.model.units[primary].workload_status == "blocked",
+        timeout=100,
+    )
+
+    logger.info("Enabling the plugin back on charm config, status should resolve")
+    await ops_test.model.applications[DATABASE_APP_NAME].set_config({
+        "plugin_pg_trgm_enable": "True"
+    })
+    await ops_test.model.wait_for_idle(apps=[DATABASE_APP_NAME], status="active")
+
+    logger.info("Re-disabling plugin, waiting for blocked status to come back")
+    await ops_test.model.applications[DATABASE_APP_NAME].set_config({
+        "plugin_pg_trgm_enable": "False"
+    })
+    await ops_test.model.block_until(
+        lambda: ops_test.model.units[primary].workload_status == "blocked",
+        timeout=100,
+    )
+
+    logger.info("Delete the object which depends on the plugin")
+    with db_connect(host=address, password=password) as connection:
+        connection.autocommit = True
+        connection.cursor().execute("DROP INDEX example_idx;")
+    connection.close()
+
+    logger.info("Waiting for status to resolve again")
+    async with ops_test.fast_forward(fast_interval="60s"):
+        await ops_test.model.wait_for_idle(apps=[DATABASE_APP_NAME], status="active")

tests/unit/test_charm.py

@@ -598,6 +598,14 @@ def test_enable_disable_extensions(harness):
             None,
         )
 
+        # Block if extension-dependent object error is raised
+        _enable_disable_extensions.side_effect = [psycopg2.errors.DependentObjectsStillExist, None]
+        harness.charm.enable_disable_extensions()
+        assert isinstance(harness.charm.unit.status, BlockedStatus)
+        # Should resolve afterwards
+        harness.charm.enable_disable_extensions()
+        assert isinstance(harness.charm.unit.status, ActiveStatus)
+
 
 def test_on_peer_relation_departed(harness):
     with (