adapter: Frontend peek sequencing -- refactor rbac::check_plan

Author: ggevay

src/adapter/src/coord/sequencer.rs

@@ -183,14 +183,14 @@ impl Coordinator {
 
             if let Err(e) = rbac::check_plan(
                 &session_catalog,
-                |id| {
+                Some(|id| {
                     // We use linear search through active connections if needed, which is fine
                     // because the RBAC check will call the closure at most once.
                     self.active_conns()
                         .into_iter()
                         .find(|(conn_id, _)| conn_id.unhandled() == id)
                         .map(|(_, conn_meta)| *conn_meta.authenticated_role_id())
-                },
+                }),
                 ctx.session(),
                 &plan,
                 target_cluster_id,

src/adapter/src/frontend_peek.rs

@@ -19,6 +19,7 @@ use mz_expr::{CollectionPlan, ResultSpec};
 use mz_ore::cast::{CastFrom, CastLossy};
 use mz_ore::now::EpochMillis;
 use mz_repr::optimize::{OptimizerFeatures, OverrideFrom};
+use mz_repr::role_id::RoleId;
 use mz_repr::{Datum, GlobalId, IntoRowIterator, Timestamp};
 use mz_sql::catalog::CatalogCluster;
 use mz_sql::plan::{self, Plan, QueryWhen};
@@ -300,11 +301,7 @@ impl PeekClient {
 
         rbac::check_plan(
             &conn_catalog,
-            |_id| {
-                // This is only used by `Plan::SideEffectingFunc`, so it is irrelevant for us here
-                // TODO(peek-seq): refactor `check_plan` to make this nicer
-                unreachable!()
-            },
+            None::<fn(u32) -> Option<RoleId>>,
             session,
             &plan,
             Some(target_cluster_id),

src/sql/src/rbac.rs

@@ -337,7 +337,8 @@ pub fn check_usage(
 pub fn check_plan(
     catalog: &impl SessionCatalog,
     // Function mapping a connection ID to an authenticated role. The roles may have been dropped concurrently.
-    active_conns: impl FnOnce(u32) -> Option<RoleId>,
+    // Only required for Plan::SideEffectingFunc; can be None for other plan types.
+    active_conns: Option<impl FnOnce(u32) -> Option<RoleId>>,
     session: &dyn SessionMetadata,
     plan: &Plan,
     target_cluster_id: Option<ClusterId>,
@@ -377,7 +378,7 @@ pub fn is_rbac_enabled_for_session(
 fn generate_rbac_requirements(
     catalog: &impl SessionCatalog,
     plan: &Plan,
-    active_conns: impl FnOnce(u32) -> Option<RoleId>,
+    active_conns: Option<impl FnOnce(u32) -> Option<RoleId>>,
     target_cluster_id: Option<ClusterId>,
     role_id: RoleId,
 ) -> RbacRequirements {
@@ -1464,11 +1465,12 @@ fn generate_rbac_requirements(
         },
         Plan::SideEffectingFunc(func) => {
             let role_membership = match func {
-                SideEffectingFunc::PgCancelBackend { connection_id } => {
-                    active_conns(*connection_id)
-                        .map(|x| [x].into())
-                        .unwrap_or_default()
-                }
+                SideEffectingFunc::PgCancelBackend { connection_id } => active_conns
+                    .expect("active_conns is required for Plan::SideEffectingFunc")(
+                    *connection_id
+                )
+                .map(|x| [x].into())
+                .unwrap_or_default(),
             };
             RbacRequirements {
                 role_membership,