use feedback in roles too, block cancellation if commit

Author: 3nids

pum/feedback.py

@@ -16,6 +16,7 @@ class Feedback(abc.ABC):
     def __init__(self) -> None:
         """Initialize the feedback instance."""
         self._is_cancelled = False
+        self._cancellation_locked = False
         self._current_step = 0
         self._total_steps = 0
 
@@ -56,17 +57,32 @@ def is_cancelled(self) -> bool:
 
         Returns:
             True if the operation should be cancelled, False otherwise.
+            Always returns False if cancellation has been locked (after commit).
         """
+        if self._cancellation_locked:
+            return False
         return self._is_cancelled
 
     def cancel(self) -> None:
-        """Cancel the operation."""
-        self._is_cancelled = True
+        """Cancel the operation.
+
+        Note: This will have no effect if cancellation has been locked (after commit).
+        """
+        if not self._cancellation_locked:
+            self._is_cancelled = True
 
     def reset(self) -> None:
         """Reset the cancellation status."""
         self._is_cancelled = False
 
+    def lock_cancellation(self) -> None:
+        """Lock cancellation to prevent it after a commit.
+
+        Once locked, is_cancelled() will always return False and cancel() will have no effect.
+        This should be called immediately before committing database changes.
+        """
+        self._cancellation_locked = True
+
 
 class LogFeedback(Feedback):
     """Feedback implementation that logs progress messages.

pum/role_manager.py

@@ -5,6 +5,7 @@
 import logging
 
 from .sql_content import SqlContent
+from .exceptions import PumException
 
 if TYPE_CHECKING:
     from .feedback import Feedback
@@ -43,12 +44,14 @@ def grant(
         role: str,
         connection: psycopg.Connection,
         commit: bool = False,
+        feedback: Optional["Feedback"] = None,
     ) -> None:
         """Grant the permission to the specified role.
         Args:
             role: The name of the role to grant the permission to.
             connection: The database connection to execute the SQL statements.
             commit: Whether to commit the transaction. Defaults to False.
+            feedback: Optional feedback object for progress reporting.
         """
         if not isinstance(role, str):
             raise TypeError("Role must be a string.")
@@ -57,6 +60,9 @@ def grant(
             raise ValueError("Schemas must be defined for the permission.")
 
         for schema in self.schemas:
+            if feedback and feedback.is_cancelled():
+                raise PumException("Permission grant cancelled by user")
+
             # Detect if schema exists; if not, warn and continue
             cursor = SqlContent("SELECT 1 FROM pg_namespace WHERE nspname = {schema}").execute(
                 connection=connection,
@@ -121,6 +127,8 @@ def grant(
                 raise ValueError(f"Unknown permission type: {self.type}")
 
         if commit:
+            if feedback:
+                feedback.lock_cancellation()
             connection.commit()
 
     def _grant_existing_types(
@@ -223,14 +231,24 @@ def exists(self, connection: psycopg.Connection) -> bool:
         return bool(cursor._pum_results)
 
     def create(
-        self, connection: psycopg.Connection, grant: bool = False, commit: bool = False
+        self,
+        connection: psycopg.Connection,
+        grant: bool = False,
+        commit: bool = False,
+        feedback: Optional["Feedback"] = None,
     ) -> None:
         """Create the role in the database.
         Args:
             connection: The database connection to execute the SQL statements.
             grant: Whether to grant permissions to the role. Defaults to False.
             commit: Whether to commit the transaction. Defaults to False.
+            feedback: Optional feedback object for progress reporting.
         """
+        if feedback and feedback.is_cancelled():
+            from .exceptions import PumException
+
+            raise PumException("Role creation cancelled by user")
+
         if self.exists(connection):
             logger.debug(f"Role {self.name} already exists, skipping creation.")
         else:
@@ -262,9 +280,13 @@ def create(
                 )
         if grant:
             for permission in self.permissions():
-                permission.grant(role=self.name, connection=connection, commit=commit)
+                permission.grant(
+                    role=self.name, connection=connection, commit=False, feedback=feedback
+                )
 
         if commit:
+            if feedback:
+                feedback.lock_cancellation()
             connection.commit()
 
 
@@ -320,11 +342,17 @@ def create_roles(
         """
         roles_list = list(self.roles.values())
         for role in roles_list:
+            if feedback and feedback.is_cancelled():
+                from .exceptions import PumException
+
+                raise PumException("Role creation cancelled by user")
             if feedback:
                 feedback.increment_step()
                 feedback.report_progress(f"Creating role: {role.name}")
-            role.create(connection=connection, commit=False, grant=grant)
+            role.create(connection=connection, commit=False, grant=grant, feedback=feedback)
         if commit:
+            if feedback:
+                feedback.lock_cancellation()
             connection.commit()
 
     def grant_permissions(
@@ -341,11 +369,19 @@ def grant_permissions(
         """
         roles_list = list(self.roles.values())
         for role in roles_list:
+            if feedback and feedback.is_cancelled():
+                from .exceptions import PumException
+
+                raise PumException("Permission grant cancelled by user")
             if feedback:
                 feedback.increment_step()
                 feedback.report_progress(f"Granting permissions to role: {role.name}")
             for permission in role.permissions():
-                permission.grant(role=role.name, connection=connection, commit=False)
+                permission.grant(
+                    role=role.name, connection=connection, commit=False, feedback=feedback
+                )
         logger.info("All permissions granted to roles.")
         if commit:
+            if feedback:
+                feedback.lock_cancellation()
             connection.commit()

pum/upgrader.py

@@ -176,6 +176,7 @@ def install(
             )
 
         if commit:
+            feedback.lock_cancellation()
             feedback.report_progress("Committing changes...")
             connection.commit()
             logger.info("Changes committed to the database.")
@@ -380,6 +381,7 @@ def upgrade(
                     connection=connection, commit=False, feedback=feedback
                 )
 
+        feedback.lock_cancellation()
         feedback.report_progress("Committing changes...")
         connection.commit()
         logger.info("Upgrade completed and changes committed to the database.")
@@ -432,6 +434,8 @@ def uninstall(
             uninstall_hook.execute(connection=connection, commit=False, parameters=parameters)
 
         if commit:
+            feedback.lock_cancellation()
             feedback.report_progress("Committing changes...")
             connection.commit()
             logger.info("Uninstall completed and changes committed to the database.")
+            logger.info("Uninstall completed and changes committed to the database.")

test/test_feedback.py

@@ -261,6 +261,57 @@ def test_feedback_progression(self) -> None:
             self.assertTrue(any("create_second_table.sql" in msg for msg in feedback.messages))
             self.assertTrue(any("create_third_table.sql" in msg for msg in feedback.messages))
 
+    def test_cancellation_locked_after_commit(self) -> None:
+        """Test that cancellation is locked after commit and cannot be triggered."""
+        feedback = LogFeedback()
+
+        # Initially, cancellation should work
+        self.assertFalse(feedback.is_cancelled())
+        feedback.cancel()
+        self.assertTrue(feedback.is_cancelled())
+
+        # Reset for next test
+        feedback.reset()
+        self.assertFalse(feedback.is_cancelled())
+
+        # Lock cancellation (simulating what happens before commit)
+        feedback.lock_cancellation()
+
+        # Try to cancel - should have no effect
+        feedback.cancel()
+        self.assertFalse(
+            feedback.is_cancelled(), "Cancellation should be locked after lock_cancellation()"
+        )
+
+        # is_cancelled should always return False when locked
+        feedback._is_cancelled = True  # Force internal flag
+        self.assertFalse(feedback.is_cancelled(), "is_cancelled() should return False when locked")
+
+    def test_install_cancellation_not_possible_after_commit(self) -> None:
+        """Test that cancellation during install doesn't work after commit starts."""
+        test_dir = Path("test") / "data" / "single_changelog"
+        cfg = PumConfig(test_dir, pum={"module": "test_cancellation_lock"})
+
+        # Custom feedback that tries to cancel during "Committing changes"
+        class CancelOnCommitFeedback(CustomFeedback):
+            def report_progress(self, message: str, current: int = 0, total: int = 0) -> None:
+                super().report_progress(message, current, total)
+                if "Committing changes" in message:
+                    # Try to cancel during commit - should have no effect
+                    self.cancel()
+
+        feedback = CancelOnCommitFeedback()
+        upgrader = Upgrader(cfg)
+
+        with psycopg.connect(f"service={self.pg_service}") as conn:
+            # Should complete successfully despite cancel attempt during commit
+            upgrader.install(connection=conn, feedback=feedback, commit=True)
+
+            # Verify cancellation was attempted but had no effect
+            self.assertTrue(any("Committing changes" in msg for msg in feedback.messages))
+            # Even though cancel() was called, is_cancelled() should return False
+            self.assertFalse(feedback.is_cancelled())
+
 
 if __name__ == "__main__":
     unittest.main()