Fix dropping role with table privileges granted by non-neon_superuser

We were previously only revoking privileges granted by neon_superuser.
However, we need to do it for all grantors.

Signed-off-by: Tristan Partin <tristan@neon.tech>

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,30 @@ BEGIN
         -- ii) it's easy to add more schemas to the list if needed.
         WHERE schema_name IN ('public')
     LOOP
+        SET LOCAL ROLE neon_superuser;
+
         revoke_query := format(
-            'REVOKE ALL PRIVILEGES ON ALL TABLES IN SCHEMA %I FROM {role_name} GRANTED BY neon_superuser;',
+            'REVOKE ALL PRIVILEGES ON ALL TABLES IN SCHEMA %I FROM {role_name} GRANTED BY neon_superuser',
             schema
         );
 
         EXECUTE revoke_query;
+
+        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);
+
+            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

@@ -370,3 +370,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