grant roles after upgrade

Author: 3nids

docs/docs/cli/upgrade.md

@@ -1,9 +1,10 @@
-usage: pum upgrade [-h] [-p PARAMETER PARAMETER] [-u MAX_VERSION] [--beta-testing] [--skip-drop-app]
+usage: pum upgrade [-h] [-p PARAMETER PARAMETER] [-u MAX_VERSION] [-g] [--beta-testing] [--skip-drop-app]
 [--skip-create-app]
 ### options:
 - `-h, --help`: show this help message and exit
 - `-p PARAMETER PARAMETER, --parameter PARAMETER PARAMETER`: Assign variable for running SQL deltas. Format is name value.
 - `-u MAX_VERSION, --max-version MAX_VERSION`: maximum version to upgrade
+- `-g, --grant`: Grant permissions to roles
 - `--beta-testing`: Install in beta testing mode.
 - `--skip-drop-app`: Skip drop app handlers during upgrade.
 - `--skip-create-app`: Skip create app handlers during upgrade.

pum/cli.py

@@ -240,6 +240,9 @@ def create_parser() -> argparse.ArgumentParser:
         action="append",
     )
     parser_upgrade.add_argument("-u", "--max-version", help="maximum version to upgrade")
+    parser_upgrade.add_argument(
+        "-g", "--grant", help="Grant permissions to roles", action="store_true"
+    )
     parser_upgrade.add_argument(
         "--beta-testing", help="Install in beta testing mode.", action="store_true"
     )
@@ -419,6 +422,7 @@ def cli() -> int:  # noqa: PLR0912
                 connection=conn,
                 parameters=parameters,
                 max_version=args.max_version,
+                grant=args.grant,
                 beta_testing=args.beta_testing,
                 skip_drop_app=args.skip_drop_app,
                 skip_create_app=args.skip_create_app,

pum/upgrader.py

@@ -193,6 +193,7 @@ def upgrade(
         beta_testing: bool = False,
         skip_drop_app: bool = False,
         skip_create_app: bool = False,
+        grant: bool = True,
     ) -> None:
         """Upgrades the given module
         The changelogs are applied in the order they are found in the directory.
@@ -212,6 +213,8 @@ def upgrade(
                 If True, drop app handlers will be skipped.
             skip_create_app:
                 If True, create app handlers will be skipped.
+            grant:
+                If True, permissions will be granted to the roles.
         """
         if not self.schema_migrations.exists(connection):
             msg = (
@@ -251,5 +254,8 @@ def upgrade(
             for create_app_hook in self.config.create_app_handlers():
                 create_app_hook.execute(connection=connection, commit=False, parameters=parameters)
 
+        if grant:
+            self.config.role_manager().grant_permissions(connection=connection, commit=False)
+
         connection.commit()
         logger.info("Upgrade completed and changes committed to the database.")

test/test_upgrader.py

@@ -25,7 +25,11 @@ def tearDown(self) -> None:
             cur.execute("DROP SCHEMA IF EXISTS pum_test_data CASCADE;")
             cur.execute("DROP SCHEMA IF EXISTS pum_custom_migrations_schema CASCADE;")
             cur.execute("DROP SCHEMA IF EXISTS pum_test_app CASCADE;")
+            cur.execute("DROP SCHEMA IF EXISTS pum_test_data_schema_1 CASCADE;")
+            cur.execute("DROP SCHEMA IF EXISTS pum_test_data_schema_2 CASCADE;")
             cur.execute("DROP TABLE IF EXISTS public.pum_migrations;")
+            cur.execute("DROP ROLE IF EXISTS pum_test_user;")
+            cur.execute("DROP ROLE IF EXISTS pum_test_viewer;")
 
         self.tmpdir.cleanup()
         self.tmp = None
@@ -43,7 +47,11 @@ def setUp(self) -> None:
             cur.execute("DROP SCHEMA IF EXISTS pum_test_data CASCADE;")
             cur.execute("DROP SCHEMA IF EXISTS pum_custom_migrations_schema CASCADE;")
             cur.execute("DROP SCHEMA IF EXISTS pum_test_app CASCADE;")
+            cur.execute("DROP SCHEMA IF EXISTS pum_test_data_schema_1 CASCADE;")
+            cur.execute("DROP SCHEMA IF EXISTS pum_test_data_schema_2 CASCADE;")
             cur.execute("DROP TABLE IF EXISTS public.pum_migrations;")
+            cur.execute("DROP ROLE IF EXISTS pum_test_user;")
+            cur.execute("DROP ROLE IF EXISTS pum_test_viewer;")
 
         self.tmpdir = tempfile.TemporaryDirectory()
         self.tmp = self.tmpdir.name
@@ -465,6 +473,35 @@ def test_upgrade(self) -> None:
             upgrader.upgrade(connection=conn)
             self.assertEqual(sm.baseline(conn), Version("2.0.0"))
 
+    def test_upgrade_with_grant(self) -> None:
+        """Test that permissions are granted correctly after upgrade."""
+        test_dir = Path("test") / "data" / "roles"
+        cfg = PumConfig.from_yaml(test_dir / ".pum.yaml")
+        with psycopg.connect(f"service={self.pg_service}") as conn:
+            # Install with roles but without granting permissions
+            Upgrader(cfg).install(connection=conn, roles=True, grant=False, commit=True)
+
+            cur = conn.cursor()
+            # Verify viewer role doesn't have SELECT permission initially
+            cur.execute(
+                "SELECT has_table_privilege('pum_test_viewer', 'pum_test_data_schema_1.some_table_1', 'SELECT');"
+            )
+            self.assertFalse(cur.fetchone()[0])
+
+            # Now upgrade with grant=True (even though there are no new changelogs, it should grant permissions)
+            Upgrader(cfg).upgrade(connection=conn, grant=True)
+
+            # Verify permissions were granted
+            cur.execute(
+                "SELECT has_table_privilege('pum_test_viewer', 'pum_test_data_schema_1.some_table_1', 'SELECT');"
+            )
+            self.assertTrue(cur.fetchone()[0])
+
+            cur.execute(
+                "SELECT has_table_privilege('pum_test_user', 'pum_test_data_schema_2.some_table_2', 'INSERT');"
+            )
+            self.assertTrue(cur.fetchone()[0])
+
 
 if __name__ == "__main__":
     unittest.main()