sql: Ensure role member grantor is always valid (#18781)

This commit prevents dropping a role, if that role is a grantor for some
other role membership. This helps ensure that the grantor references in
the catalog always remain valid.
    
As of this commit, the only valid grantor is mz_system which is
impossible to delete. So it's impossible to hit this error scenario.
However, this is more future-proof for when we implement ADMIN OPTION
and allow other grantors.
    
This is based off of the equivalent commit in PostgreSQL here:

postgres/postgres@6566133
That commit will be part of PostgreSQL v16 and was not included in v15.
    
Part of #11579

Author: jkosh44

src/adapter/src/catalog.rs

@@ -622,10 +622,10 @@ impl CatalogState {
                 membership.insert(cur_id.clone());
                 let role = self.get_role(cur_id);
                 soft_assert!(
-                    !role.membership().contains(id),
+                    !role.membership().keys().contains(id),
                     "circular membership exists in the catalog"
                 );
-                queue.extend(role.membership().into_iter());
+                queue.extend(role.membership().keys());
             }
         }
         membership
@@ -7015,6 +7015,16 @@ impl SessionCatalog for ConnCatalog<'_> {
         self.state.get_role(id)
     }
 
+    fn get_roles(&self) -> Vec<&dyn CatalogRole> {
+        // `as` is ok to use to cast to a trait object.
+        #[allow(clippy::as_conversions)]
+        self.state
+            .roles_by_id
+            .values()
+            .map(|role| role as &dyn CatalogRole)
+            .collect()
+    }
+
     fn collect_role_membership(&self, id: &RoleId) -> BTreeSet<RoleId> {
         self.state.collect_role_membership(id)
     }
@@ -7222,12 +7232,8 @@ impl mz_sql::catalog::CatalogRole for Role {
         self.attributes.create_cluster
     }
 
-    fn membership(&self) -> BTreeSet<&RoleId> {
-        self.membership
-            .map
-            .iter()
-            .map(|(role_id, _grantor_id)| role_id)
-            .collect()
+    fn membership(&self) -> &BTreeMap<RoleId, RoleId> {
+        &self.membership.map
     }
 }
 

src/adapter/src/coord/sequencer/inner.rs

@@ -3672,7 +3672,8 @@ impl Coordinator {
         let catalog = self.catalog();
         let mut ops = Vec::new();
         for member_id in member_ids {
-            let member_membership: BTreeSet<_> = catalog.get_role(&member_id).membership();
+            let member_membership: BTreeSet<_> =
+                catalog.get_role(&member_id).membership().keys().collect();
             if member_membership.contains(&role_id) {
                 let role_name = catalog.get_role(&role_id).name().to_string();
                 let member_name = catalog.get_role(&member_id).name().to_string();
@@ -3713,7 +3714,8 @@ impl Coordinator {
         let catalog = self.catalog();
         let mut ops = Vec::new();
         for member_id in member_ids {
-            let member_membership: BTreeSet<_> = catalog.get_role(&member_id).membership();
+            let member_membership: BTreeSet<_> =
+                catalog.get_role(&member_id).membership().keys().collect();
             if !member_membership.contains(&role_id) {
                 let role_name = catalog.get_role(&role_id).name().to_string();
                 let member_name = catalog.get_role(&member_id).name().to_string();

src/sql/src/catalog.rs

@@ -155,6 +155,9 @@ pub trait SessionCatalog: fmt::Debug + ExprHumanizer + Send + Sync {
     /// Panics if `id` does not specify a valid role.
     fn get_role(&self, id: &RoleId) -> &dyn CatalogRole;
 
+    /// Gets all roles.
+    fn get_roles(&self) -> Vec<&dyn CatalogRole>;
+
     /// Collects all role IDs that `id` is transitively a member of.
     fn collect_role_membership(&self, id: &RoleId) -> BTreeSet<RoleId>;
 
@@ -458,8 +461,11 @@ pub trait CatalogRole {
     /// Indicates whether the role has the cluster creation attribute.
     fn create_cluster(&self) -> bool;
 
-    /// Returns all role IDs that this role is an immediate a member of.
-    fn membership(&self) -> BTreeSet<&RoleId>;
+    /// Returns all role IDs that this role is an immediate a member of, and the grantor of that
+    /// membership.
+    ///
+    /// Key is the role that some role is a member of, value is the grantor role ID.
+    fn membership(&self) -> &BTreeMap<RoleId, RoleId>;
 }
 
 /// A cluster in a [`SessionCatalog`].

src/sql/src/plan/statement/ddl.rs

@@ -3443,6 +3443,17 @@ fn plan_drop_role(
             if &id == scx.catalog.active_role_id() {
                 sql_bail!("current role cannot be dropped");
             }
+            for role in scx.catalog.get_roles() {
+                for (member_id, grantor_id) in role.membership() {
+                    if &id == grantor_id {
+                        let member_role = scx.catalog.get_role(member_id);
+                        sql_bail!(
+                            "cannot drop role {}: still depended up by membership of role {} in role {}",
+                            name.as_str(), role.name(), member_role.name()
+                        );
+                    }
+                }
+            }
             Ok(Some(role.id()))
         }
         Err(_) if if_exists => {

test/upgrade/check-from-current_source-role.td

@@ -0,0 +1,19 @@
+# Copyright Materialize, Inc. and contributors. All rights reserved.
+#
+# Use of this software is governed by the Business Source License
+# included in the LICENSE file at the root of this repository.
+#
+# As of the Change Date specified in that file, in accordance with
+# the Business Source License, use of this software will be governed
+# by the Apache License, Version 2.0.
+
+> SELECT role.name AS role, member.name AS member, grantor.name AS grantor FROM mz_role_members membership LEFT JOIN mz_roles role ON membership.role_id = role.id LEFT JOIN mz_roles member ON membership.member = member.id LEFT JOIN mz_roles grantor ON membership.grantor = grantor.id;
+group  joe  mz_system
+
+> DROP ROLE superuser_login;
+
+> DROP ROLE "space role";
+
+> DROP ROLE joe;
+
+> DROP ROLE group;

test/upgrade/check-from-v0.47.0-role.td

@@ -27,6 +27,9 @@ SELECT 1;
 joe
 group
 
+$ skip-if
+SELECT mz_version_num() >= 5200;
+
 > SELECT role.name AS role, member.name AS member, grantor.name AS grantor FROM mz_role_members membership LEFT JOIN mz_roles role ON membership.role_id = role.id LEFT JOIN mz_roles member ON membership.member = member.id LEFT JOIN mz_roles grantor ON membership.grantor = grantor.id;
 group  joe  mz_system