Fix dropping role with table privileges granted by non-neon_superuser

compute_tools/src/spec_apply.rs

@@ -550,7 +550,7 @@ async fn get_operations<'a>(
                                     Operation {
                                         query: format!(
                                             include_str!("sql/pre_drop_role_revoke_privileges.sql"),
-                                            role_name = quoted,
+                                            role_name = op.name,
                                         ),
                                         comment: None,
                                     },

compute_tools/src/sql/pre_drop_role_revoke_privileges.sql

@@ -1,8 +1,11 @@
-SET SESSION ROLE neon_superuser;
+-- Any time you see {role_name}, remember that this SQL script gets embedded
+-- into compute_ctl, and is run with a format!(), so {role_name} is a Rust
+-- variable.
 
 DO $$
 DECLARE
     schema TEXT;
+    grantor TEXT;
     revoke_query TEXT;
 BEGIN
     FOR schema IN
@@ -15,12 +18,21 @@ BEGIN
         -- ii) it's easy to add more schemas to the list if needed.
         WHERE schema_name IN ('public')
     LOOP
-        revoke_query := format(
-            'REVOKE ALL PRIVILEGES ON ALL TABLES IN SCHEMA %I FROM {role_name} GRANTED BY neon_superuser;',
-            schema
-        );
+        FOR grantor IN
+            SELECT DISTINCT rtg.grantor
+            FROM information_schema.role_table_grants AS rtg
+            WHERE grantee = '{role_name}'
+        LOOP
+            EXECUTE format('SET LOCAL ROLE %I', grantor);
 
-        EXECUTE revoke_query;
+            revoke_query := format(
+                'REVOKE ALL PRIVILEGES ON ALL TABLES IN SCHEMA %I FROM "{role_name}" GRANTED BY %I',
+                schema,
+                grantor
+            );
+
+            EXECUTE revoke_query;
+        END LOOP;
     END LOOP;
 END;
 $$;

test_runner/regress/test_compute_catalog.py

@@ -257,10 +257,12 @@ def test_dropdb_with_subscription(neon_simple_env: NeonEnv):
         assert curr_db[0] == "publisher_db"
 
 
-def test_compute_drop_role(neon_simple_env: NeonEnv):
+def test_drop_role_with_table_privileges_from_neon_superuser(neon_simple_env: NeonEnv):
     """
     Test that compute_ctl can drop a role even if it has some depending objects
-    like permissions in one of the databases.
+    like permissions in one of the databases that were granted by
+    neon_superuser.
+
     Reproduction test for /~https://github.com/neondatabase/cloud/issues/13582
     """
     env = neon_simple_env
@@ -370,3 +372,66 @@ def test_compute_drop_role(neon_simple_env: NeonEnv):
         cursor.execute("SELECT rolname FROM pg_roles WHERE rolname = 'readonly2'")
         role = cursor.fetchone()
         assert role is None
+
+
+def test_drop_role_with_table_privileges_from_non_neon_superuser(neon_simple_env: NeonEnv):
+    """
+    Test that compute_ctl can drop a role if the role has previously been
+    granted table privileges by a role other than neon_superuser.
+    """
+    TEST_DB_NAME = "neondb"
+    TEST_GRANTOR = "my_grantor"
+    TEST_GRANTEE = "my_grantee"
+
+    env = neon_simple_env
+
+    endpoint = env.endpoints.create_start("main")
+    endpoint.respec_deep(
+        **{
+            "skip_pg_catalog_updates": False,
+            "cluster": {
+                "roles": [
+                    {
+                        # We need to create role via compute_ctl, because in this case it will receive
+                        # additional grants equivalent to our real environment, so we can repro some
+                        # issues.
+                        "name": TEST_GRANTOR,
+                        # Some autocomplete-suggested hash, no specific meaning.
+                        "encrypted_password": "SCRAM-SHA-256$4096:hBT22QjqpydQWqEulorfXA==$miBogcoj68JWYdsNB5PW1X6PjSLBEcNuctuhtGkb4PY=:hxk2gxkwxGo6P7GCtfpMlhA9zwHvPMsCz+NQf2HfvWk=",
+                        "options": [],
+                    },
+                ],
+                "databases": [
+                    {
+                        "name": TEST_DB_NAME,
+                        "owner": TEST_GRANTOR,
+                    },
+                ],
+            },
+        }
+    )
+
+    endpoint.reconfigure()
+
+    with endpoint.cursor(dbname=TEST_DB_NAME, user=TEST_GRANTOR) as cursor:
+        cursor.execute(f'CREATE USER "{TEST_GRANTEE}"')
+        cursor.execute("CREATE TABLE test_table(id bigint)")
+        cursor.execute(f'GRANT ALL ON TABLE test_table TO "{TEST_GRANTEE}"')
+
+    endpoint.respec_deep(
+        **{
+            "skip_pg_catalog_updates": False,
+            "delta_operations": [
+                {
+                    "action": "delete_role",
+                    "name": TEST_GRANTEE,
+                },
+            ],
+        }
+    )
+    endpoint.reconfigure()
+
+    with endpoint.cursor() as cursor:
+        cursor.execute(f"SELECT rolname FROM pg_roles WHERE rolname = '{TEST_GRANTEE}'")
+        role = cursor.fetchone()
+        assert role is None