Merge pull request #34262 from ggevay/frontend-peek-rbac-and-statistics

ggevay · web-flow · commit 37e1c7b44880 · 2025-11-25T15:47:57.000+01:00
Frontend peek sequencing -- rbac and statistics
diff --git a/src/adapter/src/coord/sequencer.rs b/src/adapter/src/coord/sequencer.rs
@@ -12,7 +12,7 @@
 
 //! Logic for executing a planned SQL query.
 
-use std::collections::BTreeSet;
+use std::collections::{BTreeMap, BTreeSet};
 use std::str::FromStr;
 use std::sync::Arc;
 
@@ -40,17 +40,21 @@ use mz_sql::plan::{
 };
 use mz_sql::rbac;
 use mz_sql::session::metadata::SessionMetadata;
+use mz_sql::session::vars;
 use mz_sql::session::vars::SessionVars;
 use mz_sql_parser::ast::{Raw, Statement};
 use mz_storage_client::client::TableData;
 use mz_storage_client::storage_collections::StorageCollections;
 use mz_storage_types::connections::inline::IntoInlineConnection;
+use mz_storage_types::controller::StorageError;
 use mz_storage_types::stats::RelationPartStats;
 use mz_transform::dataflow::DataflowMetainfo;
 use mz_transform::notice::{OptimizerNoticeApi, OptimizerNoticeKind, RawOptimizerNotice};
+use mz_transform::{EmptyStatisticsOracle, StatisticsOracle};
 use timely::progress::Antichain;
+use timely::progress::Timestamp as TimelyTimestamp;
 use tokio::sync::oneshot;
-use tracing::{Instrument, Level, Span, event};
+use tracing::{Instrument, Level, Span, event, warn};
 
 use crate::ExecuteContext;
 use crate::catalog::{Catalog, CatalogState};
@@ -183,14 +187,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,
@@ -1205,3 +1209,88 @@ pub(crate) async fn explain_plan_inner(
 
     Ok(rows)
 }
+
+/// Creates a statistics oracle for query optimization.
+///
+/// This is a free-standing function that can be called from both the old peek sequencing
+/// and the new frontend peek sequencing.
+pub(crate) async fn statistics_oracle(
+    session: &Session,
+    source_ids: &BTreeSet<GlobalId>,
+    query_as_of: &Antichain<Timestamp>,
+    is_oneshot: bool,
+    system_config: &vars::SystemVars,
+    storage_collections: &dyn StorageCollections<Timestamp = Timestamp>,
+) -> Result<Box<dyn StatisticsOracle>, AdapterError> {
+    if !session.vars().enable_session_cardinality_estimates() {
+        return Ok(Box::new(EmptyStatisticsOracle));
+    }
+
+    let timeout = if is_oneshot {
+        // TODO(mgree): ideally, we would shorten the timeout even more if we think the query could take the fast path
+        system_config.optimizer_oneshot_stats_timeout()
+    } else {
+        system_config.optimizer_stats_timeout()
+    };
+
+    let cached_stats = mz_ore::future::timeout(
+        timeout,
+        CachedStatisticsOracle::new(source_ids, query_as_of, storage_collections),
+    )
+    .await;
+
+    match cached_stats {
+        Ok(stats) => Ok(Box::new(stats)),
+        Err(mz_ore::future::TimeoutError::DeadlineElapsed) => {
+            warn!(
+                is_oneshot = is_oneshot,
+                "optimizer statistics collection timed out after {}ms",
+                timeout.as_millis()
+            );
+
+            Ok(Box::new(EmptyStatisticsOracle))
+        }
+        Err(mz_ore::future::TimeoutError::Inner(e)) => Err(AdapterError::Storage(e)),
+    }
+}
+
+#[derive(Debug)]
+struct CachedStatisticsOracle {
+    cache: BTreeMap<GlobalId, usize>,
+}
+
+impl CachedStatisticsOracle {
+    pub async fn new<T: TimelyTimestamp>(
+        ids: &BTreeSet<GlobalId>,
+        as_of: &Antichain<T>,
+        storage_collections: &dyn mz_storage_client::storage_collections::StorageCollections<Timestamp = T>,
+    ) -> Result<Self, StorageError<T>> {
+        let mut cache = BTreeMap::new();
+
+        for id in ids {
+            let stats = storage_collections.snapshot_stats(*id, as_of.clone()).await;
+
+            match stats {
+                Ok(stats) => {
+                    cache.insert(*id, stats.num_updates);
+                }
+                Err(StorageError::IdentifierMissing(id)) => {
+                    ::tracing::debug!("no statistics for {id}")
+                }
+                Err(e) => return Err(e),
+            }
+        }
+
+        Ok(Self { cache })
+    }
+}
+
+impl StatisticsOracle for CachedStatisticsOracle {
+    fn cardinality_estimate(&self, id: GlobalId) -> Option<usize> {
+        self.cache.get(&id).map(|estimate| *estimate)
+    }
+
+    fn as_map(&self) -> BTreeMap<GlobalId, usize> {
+        self.cache.clone()
+    }
+}
diff --git a/src/adapter/src/coord/sequencer/inner.rs b/src/adapter/src/coord/sequencer/inner.rs
@@ -86,11 +86,9 @@ use mz_storage_client::controller::{CollectionDescription, DataSource, ExportDes
 use mz_storage_types::AlterCompatible;
 use mz_storage_types::connections::inline::IntoInlineConnection;
 use mz_storage_types::controller::StorageError;
-use mz_transform::EmptyStatisticsOracle;
 use mz_transform::dataflow::DataflowMetainfo;
 use smallvec::SmallVec;
 use timely::progress::Antichain;
-use timely::progress::Timestamp as TimelyTimestamp;
 use tokio::sync::{oneshot, watch};
 use tracing::{Instrument, Span, info, warn};
 
@@ -5124,93 +5122,23 @@ impl Coordinator {
 
         Ok(ExecuteResponse::AlteredObject(ObjectType::Table))
     }
-}
-
-#[derive(Debug)]
-struct CachedStatisticsOracle {
-    cache: BTreeMap<GlobalId, usize>,
-}
-
-impl CachedStatisticsOracle {
-    pub async fn new<T: TimelyTimestamp>(
-        ids: &BTreeSet<GlobalId>,
-        as_of: &Antichain<T>,
-        storage_collections: &dyn mz_storage_client::storage_collections::StorageCollections<Timestamp = T>,
-    ) -> Result<Self, StorageError<T>> {
-        let mut cache = BTreeMap::new();
-
-        for id in ids {
-            let stats = storage_collections.snapshot_stats(*id, as_of.clone()).await;
 
-            match stats {
-                Ok(stats) => {
-                    cache.insert(*id, stats.num_updates);
-                }
-                Err(StorageError::IdentifierMissing(id)) => {
-                    ::tracing::debug!("no statistics for {id}")
-                }
-                Err(e) => return Err(e),
-            }
-        }
-
-        Ok(Self { cache })
-    }
-}
-
-impl mz_transform::StatisticsOracle for CachedStatisticsOracle {
-    fn cardinality_estimate(&self, id: GlobalId) -> Option<usize> {
-        self.cache.get(&id).map(|estimate| *estimate)
-    }
-
-    fn as_map(&self) -> BTreeMap<GlobalId, usize> {
-        self.cache.clone()
-    }
-}
-
-impl Coordinator {
     pub(super) async fn statistics_oracle(
         &self,
         session: &Session,
         source_ids: &BTreeSet<GlobalId>,
         query_as_of: &Antichain<Timestamp>,
         is_oneshot: bool,
     ) -> Result<Box<dyn mz_transform::StatisticsOracle>, AdapterError> {
-        if !session.vars().enable_session_cardinality_estimates() {
-            return Ok(Box::new(EmptyStatisticsOracle));
-        }
-
-        let timeout = if is_oneshot {
-            // TODO(mgree): ideally, we would shorten the timeout even more if we think the query could take the fast path
-            self.catalog()
-                .system_config()
-                .optimizer_oneshot_stats_timeout()
-        } else {
-            self.catalog().system_config().optimizer_stats_timeout()
-        };
-
-        let cached_stats = mz_ore::future::timeout(
-            timeout,
-            CachedStatisticsOracle::new(
-                source_ids,
-                query_as_of,
-                self.controller.storage_collections.as_ref(),
-            ),
+        super::statistics_oracle(
+            session,
+            source_ids,
+            query_as_of,
+            is_oneshot,
+            self.catalog().system_config(),
+            self.controller.storage_collections.as_ref(),
         )
-        .await;
-
-        match cached_stats {
-            Ok(stats) => Ok(Box::new(stats)),
-            Err(mz_ore::future::TimeoutError::DeadlineElapsed) => {
-                warn!(
-                    is_oneshot = is_oneshot,
-                    "optimizer statistics collection timed out after {}ms",
-                    timeout.as_millis()
-                );
-
-                Ok(Box::new(EmptyStatisticsOracle))
-            }
-            Err(mz_ore::future::TimeoutError::Inner(e)) => Err(AdapterError::Storage(e)),
-        }
+        .await
     }
 }
 
diff --git a/src/adapter/src/frontend_peek.rs b/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};
@@ -35,7 +36,7 @@ use tracing_opentelemetry::OpenTelemetrySpanExt;
 use crate::catalog::CatalogState;
 use crate::command::Command;
 use crate::coord::peek::PeekPlan;
-use crate::coord::sequencer::eval_copy_to_uri;
+use crate::coord::sequencer::{eval_copy_to_uri, statistics_oracle};
 use crate::coord::timestamp_selection::TimestampDetermination;
 use crate::coord::{Coordinator, CopyToContext, ExplainContext, ExplainPlanContext, TargetCluster};
 use crate::explain::insights::PlanInsightsContext;
@@ -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),
@@ -518,14 +515,16 @@ impl PeekClient {
 
         // # From peek_optimize
 
-        if session.vars().enable_session_cardinality_estimates() {
-            debug!(
-                "Bailing out from try_frontend_peek_inner, because of enable_session_cardinality_estimates"
-            );
-            return Ok(None);
-        }
-        // TODO(peek-seq): wire up statistics
-        let stats = Box::new(EmptyStatisticsOracle);
+        let stats = statistics_oracle(
+            session,
+            &source_ids,
+            &determination.timestamp_context.antichain(),
+            true,
+            catalog.system_config(),
+            &*self.storage_collections,
+        )
+        .await
+        .unwrap_or_else(|_| Box::new(EmptyStatisticsOracle));
 
         // Generate data structures that can be moved to another task where we will perform possibly
         // expensive optimizations.
diff --git a/src/sql/src/rbac.rs b/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,
diff --git a/src/storage-client/src/storage_collections.rs b/src/storage-client/src/storage_collections.rs
@@ -88,7 +88,7 @@ mod metrics;
 /// - Hands out [ReadHolds](ReadHold) that prevent a collection's since from
 /// advancing while it needs to be read at a specific time.
 #[async_trait]
-pub trait StorageCollections: Debug {
+pub trait StorageCollections: Debug + Sync {
     type Timestamp: TimelyTimestamp;
 
     /// On boot, reconcile this [StorageCollections] with outside state. We get
