mz_adapter/optimize/
dataflows.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10//! Types and methods for building and shipping dataflow descriptions.
11//!
12//! Dataflows are buildable from the coordinator's `catalog` and `indexes`
13//! members, which respectively describe the collection backing identifiers
14//! and indicate which identifiers have arrangements available. This module
15//! isolates that logic from the rest of the somewhat complicated coordinator.
16
17use std::collections::{BTreeMap, BTreeSet};
18
19use chrono::{DateTime, Utc};
20use maplit::{btreemap, btreeset};
21use tracing::warn;
22
23use mz_catalog::memory::objects::{CatalogItem, DataSourceDesc, Index, TableDataSource, View};
24use mz_compute_client::controller::error::InstanceMissing;
25use mz_compute_types::ComputeInstanceId;
26use mz_compute_types::dataflows::{DataflowDesc, DataflowDescription, IndexDesc};
27use mz_controller::Controller;
28use mz_expr::visit::Visit;
29use mz_expr::{
30    CollectionPlan, Id, MapFilterProject, MirRelationExpr, MirScalarExpr, OptimizedMirRelationExpr,
31    RECURSION_LIMIT, UnmaterializableFunc,
32};
33use mz_ore::cast::ReinterpretCast;
34use mz_ore::stack::{CheckedRecursion, RecursionGuard, RecursionLimitError, maybe_grow};
35use mz_repr::adt::array::ArrayDimension;
36use mz_repr::explain::trace_plan;
37use mz_repr::optimize::OptimizerFeatures;
38use mz_repr::role_id::RoleId;
39use mz_repr::{Datum, GlobalId, ReprRelationType, Row};
40use mz_sql::catalog::CatalogRole;
41use mz_sql::rbac;
42use mz_sql::session::metadata::SessionMetadata;
43use mz_transform::analysis::DerivedBuilder;
44use mz_transform::analysis::monotonic::Monotonic;
45
46use crate::catalog::CatalogState;
47use crate::coord::id_bundle::CollectionIdBundle;
48use crate::optimize::{Optimize, OptimizerCatalog, OptimizerConfig, OptimizerError, view};
49use crate::session::{SERVER_MAJOR_VERSION, SERVER_MINOR_VERSION};
50use crate::util::viewable_variables;
51
52/// A reference-less snapshot of a compute instance. There is no guarantee `instance_id` continues
53/// to exist after this has been made.
54#[derive(Debug, Clone)]
55pub struct ComputeInstanceSnapshot {
56    instance_id: ComputeInstanceId,
57    /// The collections that exist on this compute instance. If it's None, then any collection that
58    /// a caller asks us about is considered to exist.
59    /// TODO(peek-seq): Remove this completely once all callers are able to handle suddenly missing
60    /// collections, in which case we won't need a `ComputeInstanceSnapshot` at all.
61    collections: Option<BTreeSet<GlobalId>>,
62}
63
64impl ComputeInstanceSnapshot {
65    pub fn new(controller: &Controller, id: ComputeInstanceId) -> Result<Self, InstanceMissing> {
66        controller
67            .compute
68            .collection_ids(id)
69            .map(|collection_ids| Self {
70                instance_id: id,
71                collections: Some(collection_ids.collect()),
72            })
73    }
74
75    pub fn new_from_parts(instance_id: ComputeInstanceId, collections: BTreeSet<GlobalId>) -> Self {
76        Self {
77            instance_id,
78            collections: Some(collections),
79        }
80    }
81
82    pub fn new_without_collections(instance_id: ComputeInstanceId) -> Self {
83        Self {
84            instance_id,
85            collections: None,
86        }
87    }
88
89    /// Return the ID of this compute instance.
90    pub fn instance_id(&self) -> ComputeInstanceId {
91        self.instance_id
92    }
93
94    /// Reports whether the instance contains the indicated collection. If the snapshot doesn't
95    /// track collections, then it returns true.
96    pub fn contains_collection(&self, id: &GlobalId) -> bool {
97        self.collections
98            .as_ref()
99            .map_or(true, |collections| collections.contains(id))
100    }
101
102    /// Inserts the given collection into the snapshot.
103    pub fn insert_collection(&mut self, id: GlobalId) {
104        self.collections
105            .as_mut()
106            .expect("insert_collection called on snapshot with None collections")
107            .insert(id);
108    }
109}
110
111/// Borrows of catalog and indexes sufficient to build dataflow descriptions.
112#[derive(Debug)]
113pub struct DataflowBuilder<'a> {
114    pub catalog: &'a dyn OptimizerCatalog,
115    /// A handle to the compute abstraction, which describes indexes by identifier.
116    ///
117    /// This can also be used to grab a handle to the storage abstraction, through
118    /// its `storage_mut()` method.
119    pub compute: ComputeInstanceSnapshot,
120    /// If set, indicates that the `DataflowBuilder` operates in "replan" mode
121    /// and should consider only catalog items that are strictly less than the
122    /// given [`GlobalId`].
123    ///
124    /// In particular, indexes with higher [`GlobalId`] that are present in the
125    /// catalog will be ignored.
126    ///
127    /// Bound from [`OptimizerConfig::replan`].
128    pub replan: Option<GlobalId>,
129    /// A guard for recursive operations in this [`DataflowBuilder`] instance.
130    recursion_guard: RecursionGuard,
131}
132
133/// Behavior to prepare relation and scalar expressions for use in a dataflow.
134pub trait ExprPrep {
135    /// Prepare a relation expression.
136    fn prep_relation_expr(&self, expr: &mut OptimizedMirRelationExpr)
137    -> Result<(), OptimizerError>;
138
139    /// Prepare a scalar expression.
140    fn prep_scalar_expr(&self, expr: &mut MirScalarExpr) -> Result<(), OptimizerError>;
141}
142
143/// A no-op expression preparer.
144pub struct ExprPrepNoop;
145impl ExprPrep for ExprPrepNoop {
146    fn prep_relation_expr(&self, _: &mut OptimizedMirRelationExpr) -> Result<(), OptimizerError> {
147        Ok(())
148    }
149    fn prep_scalar_expr(&self, _expr: &mut MirScalarExpr) -> Result<(), OptimizerError> {
150        Ok(())
151    }
152}
153
154/// Preparing an expression for maintained dataflow, e.g., index, materialized view, or subscribe.
155/// Produces errors for calls to unmaterializable functions.
156pub struct ExprPrepMaintained;
157
158impl ExprPrep for ExprPrepMaintained {
159    fn prep_relation_expr(
160        &self,
161        expr: &mut OptimizedMirRelationExpr,
162    ) -> Result<(), OptimizerError> {
163        expr.0.try_visit_mut_post(&mut |e| {
164            // Carefully test filter expressions, which may represent temporal filters.
165            if let MirRelationExpr::Filter { input, predicates } = &*e {
166                let mfp = MapFilterProject::new(input.arity()).filter(predicates.iter().cloned());
167                match mfp.into_plan() {
168                    Err(e) => Err(OptimizerError::UnsupportedTemporalExpression(e)),
169                    Ok(mut mfp) => {
170                        for s in mfp.iter_nontemporal_exprs() {
171                            self.prep_scalar_expr(s)?;
172                        }
173                        Ok(())
174                    }
175                }
176            } else {
177                e.try_visit_scalars_mut1(&mut |s| self.prep_scalar_expr(s))
178            }
179        })
180    }
181
182    fn prep_scalar_expr(&self, expr: &mut MirScalarExpr) -> Result<(), OptimizerError> {
183        // Reject the query if it contains any unmaterializable function calls.
184        let mut last_observed_unmaterializable_func = None;
185        expr.visit_mut_post(&mut |e| {
186            if let MirScalarExpr::CallUnmaterializable(f) = e {
187                last_observed_unmaterializable_func = Some(f.clone());
188            }
189        })?;
190
191        if let Some(f) = last_observed_unmaterializable_func {
192            Err(OptimizerError::UnmaterializableFunction(f))
193        } else {
194            Ok(())
195        }
196    }
197}
198
199/// Prepare an expression to run once at a logical time in a session.
200/// Calls to all unmaterializable functions are replaced with constants.
201pub struct ExprPrepOneShot<'a> {
202    pub logical_time: EvalTime,
203    pub session: &'a dyn SessionMetadata,
204    pub catalog_state: &'a CatalogState,
205}
206
207impl ExprPrep for ExprPrepOneShot<'_> {
208    fn prep_relation_expr(
209        &self,
210        expr: &mut OptimizedMirRelationExpr,
211    ) -> Result<(), OptimizerError> {
212        expr.0
213            .try_visit_scalars_mut(&mut |s| self.prep_scalar_expr(s))
214    }
215
216    fn prep_scalar_expr(&self, expr: &mut MirScalarExpr) -> Result<(), OptimizerError> {
217        // Evaluate each unmaterializable function and replace the
218        // invocation with the result.
219        expr.try_visit_mut_post(&mut |e| {
220            if let MirScalarExpr::CallUnmaterializable(f) = e {
221                *e = eval_unmaterializable_func(
222                    self.catalog_state,
223                    f,
224                    self.logical_time,
225                    self.session,
226                )?;
227            }
228            Ok(())
229        })
230    }
231}
232
233/// Prepare an expression for evaluation in a CHECK expression of a webhook source.
234/// Replaces calls to `UnmaterializableFunc::CurrentTimestamp`, others are left untouched.
235pub struct ExprPrepWebhookValidation {
236    /// Time at which this expression is being evaluated.
237    pub now: DateTime<Utc>,
238}
239
240impl ExprPrep for ExprPrepWebhookValidation {
241    fn prep_relation_expr(
242        &self,
243        expr: &mut OptimizedMirRelationExpr,
244    ) -> Result<(), OptimizerError> {
245        expr.0
246            .try_visit_scalars_mut(&mut |s| self.prep_scalar_expr(s))
247    }
248
249    fn prep_scalar_expr(&self, expr: &mut MirScalarExpr) -> Result<(), OptimizerError> {
250        let now = self.now;
251        expr.try_visit_mut_post(&mut |e| {
252            if let MirScalarExpr::CallUnmaterializable(f @ UnmaterializableFunc::CurrentTimestamp) =
253                e
254            {
255                let now: Datum = now.try_into()?;
256                let const_expr = MirScalarExpr::literal_ok(now, f.output_type().scalar_type);
257                *e = const_expr;
258            }
259            Ok(())
260        })
261    }
262}
263
264#[derive(Clone, Copy, Debug)]
265pub enum EvalTime {
266    Time(mz_repr::Timestamp),
267    /// Errors on mz_now() calls.
268    NotAvailable,
269}
270
271/// Returns an ID bundle with the given dataflows imports.
272pub fn dataflow_import_id_bundle<P>(
273    dataflow: &DataflowDescription<P>,
274    compute_instance: ComputeInstanceId,
275) -> CollectionIdBundle {
276    let storage_ids = dataflow.source_imports.keys().copied().collect();
277    let compute_ids = dataflow.index_imports.keys().copied().collect();
278    CollectionIdBundle {
279        storage_ids,
280        compute_ids: btreemap! {compute_instance => compute_ids},
281    }
282}
283
284impl<'a> DataflowBuilder<'a> {
285    pub fn new(catalog: &'a dyn OptimizerCatalog, compute: ComputeInstanceSnapshot) -> Self {
286        Self {
287            catalog,
288            compute,
289            replan: None,
290            recursion_guard: RecursionGuard::with_limit(RECURSION_LIMIT),
291        }
292    }
293
294    // TODO(aalexandrov): strictly speaking it should be better if we can make
295    // `config: &OptimizerConfig` a field in the enclosing builder. However,
296    // before we can do that we should make sure that nobody outside of the
297    // optimizer is using a DataflowBuilder instance.
298    pub(super) fn with_config(mut self, config: &OptimizerConfig) -> Self {
299        self.replan = config.replan;
300        self
301    }
302
303    /// Imports the view, source, or table with `id` into the provided
304    /// dataflow description. [`OptimizerFeatures`] is used while running
305    /// the [`Monotonic`] analysis.
306    ///
307    /// Panics if `id` refers to a non-importable item, such as an index or sink.
308    pub fn import_into_dataflow(
309        &mut self,
310        id: &GlobalId,
311        dataflow: &mut DataflowDesc,
312        features: &OptimizerFeatures,
313    ) -> Result<(), OptimizerError> {
314        maybe_grow(|| {
315            // Avoid importing the item redundantly.
316            if dataflow.is_imported(id) {
317                return Ok(());
318            }
319
320            let monotonic = self.monotonic_object(*id, features);
321
322            // A valid index is any index on `id` that is known to index oracle.
323            // Here, we import all indexes that belong to all imported collections. Later,
324            // `prune_and_annotate_dataflow_index_imports` runs at the end of the MIR
325            // pipeline, and removes unneeded index imports based on the optimized plan.
326            let mut valid_indexes = self.indexes_on(*id).peekable();
327            if valid_indexes.peek().is_some() {
328                for (index_id, idx) in valid_indexes {
329                    let index_desc = IndexDesc {
330                        on_id: *id,
331                        key: idx.keys.to_vec(),
332                    };
333                    let entry = self.catalog.get_entry(id);
334                    let desc = entry
335                        .relation_desc()
336                        .expect("indexes can only be built on items with descs");
337                    dataflow.import_index(
338                        index_id,
339                        index_desc,
340                        ReprRelationType::from(desc.typ()),
341                        monotonic,
342                    );
343                }
344            } else {
345                drop(valid_indexes);
346                let entry = self.catalog.get_entry(id);
347                // Note that the following match should be kept in sync with `sufficient_collections`.
348                match entry.item() {
349                    CatalogItem::Table(table) => {
350                        dataflow.import_source(*id, table.desc_for(id).into_typ(), monotonic);
351                    }
352                    CatalogItem::Source(source) => {
353                        dataflow.import_source(*id, source.desc.typ().clone(), monotonic);
354                    }
355                    CatalogItem::View(view) => {
356                        let expr = view.locally_optimized_expr.as_ref();
357                        self.import_view_into_dataflow(id, expr, dataflow, features)?;
358                    }
359                    CatalogItem::MaterializedView(mview) if mview.replacement_target.is_some() => {
360                        // Can't read from replacements, use the view definition directly.
361                        let expr = mview.locally_optimized_expr.as_ref();
362                        self.import_view_into_dataflow(id, expr, dataflow, features)?;
363                    }
364                    CatalogItem::MaterializedView(mview) => {
365                        dataflow.import_source(*id, mview.desc_for(id).into_typ(), monotonic);
366                    }
367                    CatalogItem::Log(log) => {
368                        dataflow.import_source(*id, log.variant.desc().typ().clone(), monotonic);
369                    }
370                    CatalogItem::ContinualTask(ct) => {
371                        dataflow.import_source(*id, ct.desc.typ().clone(), monotonic);
372                    }
373                    CatalogItem::Sink(_)
374                    | CatalogItem::Index(_)
375                    | CatalogItem::Type(_)
376                    | CatalogItem::Func(_)
377                    | CatalogItem::Secret(_)
378                    | CatalogItem::Connection(_) => {
379                        // Non-importable thing; can't get here.
380                        unreachable!()
381                    }
382                }
383            }
384            Ok(())
385        })
386    }
387
388    /// Imports the view with the specified ID and expression into the provided
389    /// dataflow description. [`OptimizerFeatures`] is used while running
390    /// expression [`mz_transform::analysis::Analysis`].
391    ///
392    /// You should generally prefer calling
393    /// [`DataflowBuilder::import_into_dataflow`], which can handle objects of
394    /// any type as long as they exist in the catalog. This method exists for
395    /// when the view does not exist in the catalog, e.g., because it is
396    /// identified by a [`GlobalId::Transient`].
397    pub fn import_view_into_dataflow(
398        &mut self,
399        view_id: &GlobalId,
400        view: &OptimizedMirRelationExpr,
401        dataflow: &mut DataflowDesc,
402        features: &OptimizerFeatures,
403    ) -> Result<(), OptimizerError> {
404        for get_id in view.depends_on() {
405            self.import_into_dataflow(&get_id, dataflow, features)?;
406        }
407        dataflow.insert_plan(*view_id, view.clone());
408        Ok(())
409    }
410
411    // Re-optimize the imported view plans using the current optimizer
412    // configuration if reoptimization is requested.
413    pub fn maybe_reoptimize_imported_views(
414        &self,
415        df_desc: &mut DataflowDesc,
416        config: &OptimizerConfig,
417    ) -> Result<(), OptimizerError> {
418        if !config.features.reoptimize_imported_views {
419            return Ok(()); // Do nothing if not explicitly requested.
420        }
421
422        let mut view_optimizer = view::Optimizer::new(config.clone(), None);
423        for desc in df_desc.objects_to_build.iter_mut().rev() {
424            if matches!(desc.id, GlobalId::Explain | GlobalId::Transient(_)) {
425                continue; // Skip descriptions that do not reference proper views.
426            }
427            if let CatalogItem::View(view) = &self.catalog.get_entry(&desc.id).item {
428                let _span = tracing::span!(
429                    target: "optimizer",
430                    tracing::Level::DEBUG,
431                    "view",
432                    path.segment = desc.id.to_string()
433                )
434                .entered();
435
436                // Reoptimize the view and update the resulting `desc.plan`.
437                desc.plan = view_optimizer.optimize(view.raw_expr.as_ref().clone())?;
438
439                // Report the optimized plan under this span.
440                trace_plan(desc.plan.as_inner());
441            }
442        }
443
444        Ok(())
445    }
446
447    /// Determine the given source's monotonicity.
448    fn monotonic_source(&self, data_source: &DataSourceDesc) -> bool {
449        match data_source {
450            DataSourceDesc::Ingestion { .. } => false,
451            DataSourceDesc::OldSyntaxIngestion {
452                desc, data_config, ..
453            } => data_config.monotonic(&desc.connection),
454            DataSourceDesc::Webhook { .. } => true,
455            DataSourceDesc::IngestionExport {
456                ingestion_id,
457                data_config,
458                ..
459            } => {
460                let source_desc = self
461                    .catalog
462                    .get_entry_by_item_id(ingestion_id)
463                    .source_desc()
464                    .expect("ingestion export must reference a source")
465                    .expect("ingestion export must reference a source");
466                data_config.monotonic(&source_desc.connection)
467            }
468            DataSourceDesc::Introspection(_)
469            | DataSourceDesc::Progress
470            | DataSourceDesc::Catalog => false,
471        }
472    }
473
474    /// Determine the given objects's monotonicity.
475    ///
476    /// This recursively traverses the expressions of all views depended on by the given object.
477    /// If this becomes a performance problem, we could add the monotonicity information of views
478    /// into the catalog instead.
479    ///
480    /// Note that materialized views are never monotonic, no matter their definition, because the
481    /// self-correcting persist_sink may insert retractions to correct the contents of its output
482    /// collection.
483    fn monotonic_object(&self, id: GlobalId, features: &OptimizerFeatures) -> bool {
484        self.monotonic_object_inner(id, &mut BTreeMap::new(), features)
485            .unwrap_or_else(|e| {
486                warn!(%id, "error inspecting object for monotonicity: {e}");
487                false
488            })
489    }
490
491    fn monotonic_object_inner(
492        &self,
493        id: GlobalId,
494        memo: &mut BTreeMap<GlobalId, bool>,
495        features: &OptimizerFeatures,
496    ) -> Result<bool, RecursionLimitError> {
497        // An object might be reached multiple times. If we already computed the monotonicity of
498        // the given ID, use that. If not, then compute it and remember the result.
499        if let Some(monotonic) = memo.get(&id) {
500            return Ok(*monotonic);
501        }
502
503        let monotonic = self.checked_recur(|_| {
504            match self.catalog.get_entry(&id).item() {
505                CatalogItem::Source(source) => Ok(self.monotonic_source(&source.data_source)),
506                CatalogItem::Table(table) => match &table.data_source {
507                    TableDataSource::TableWrites { .. } => Ok(false),
508                    TableDataSource::DataSource { desc, timeline: _ } => {
509                        Ok(self.monotonic_source(desc))
510                    }
511                },
512                CatalogItem::View(View {
513                    locally_optimized_expr: optimized_expr,
514                    ..
515                }) => {
516                    let view_expr = optimized_expr.as_ref().clone().into_inner();
517
518                    // Inspect global ids that occur in the Gets in view_expr, and collect the ids
519                    // of monotonic dependees.
520                    let mut monotonic_ids = BTreeSet::new();
521                    let recursion_result: Result<(), RecursionLimitError> = view_expr
522                        .try_visit_post(&mut |e| {
523                            if let MirRelationExpr::Get {
524                                id: Id::Global(got_id),
525                                ..
526                            } = e
527                            {
528                                if self.monotonic_object_inner(*got_id, memo, features)? {
529                                    monotonic_ids.insert(*got_id);
530                                }
531                            }
532                            Ok(())
533                        });
534                    if let Err(error) = recursion_result {
535                        // We still might have got some of the IDs, so just log and continue. Now
536                        // the subsequent monotonicity analysis can have false negatives.
537                        warn!(%id, "error inspecting view for monotonicity: {error}");
538                    }
539
540                    let mut builder = DerivedBuilder::new(features);
541                    builder.require(Monotonic::new(monotonic_ids.clone()));
542                    let derived = builder.visit(&view_expr);
543
544                    Ok(*derived
545                        .as_view()
546                        .value::<Monotonic>()
547                        .expect("Expected monotonic result from non empty tree"))
548                }
549                CatalogItem::Index(Index { on, .. }) => {
550                    self.monotonic_object_inner(*on, memo, features)
551                }
552                CatalogItem::Secret(_)
553                | CatalogItem::Type(_)
554                | CatalogItem::Connection(_)
555                | CatalogItem::Log(_)
556                | CatalogItem::MaterializedView(_)
557                | CatalogItem::Sink(_)
558                | CatalogItem::Func(_)
559                | CatalogItem::ContinualTask(_) => Ok(false),
560            }
561        })?;
562
563        memo.insert(id, monotonic);
564
565        Ok(monotonic)
566    }
567}
568
569impl<'a> CheckedRecursion for DataflowBuilder<'a> {
570    fn recursion_guard(&self) -> &RecursionGuard {
571        &self.recursion_guard
572    }
573}
574
575fn eval_unmaterializable_func(
576    state: &CatalogState,
577    f: &UnmaterializableFunc,
578    logical_time: EvalTime,
579    session: &dyn SessionMetadata,
580) -> Result<MirScalarExpr, OptimizerError> {
581    let pack_1d_array = |datums: Vec<Datum>| {
582        let mut row = Row::default();
583        row.packer()
584            .try_push_array(
585                &[ArrayDimension {
586                    lower_bound: 1,
587                    length: datums.len(),
588                }],
589                datums,
590            )
591            .expect("known to be a valid array");
592        Ok(MirScalarExpr::literal_from_single_element_row(
593            row,
594            f.output_type().scalar_type,
595        ))
596    };
597    let pack_dict = |mut datums: Vec<(String, String)>| {
598        datums.sort();
599        let mut row = Row::default();
600        row.packer().push_dict(
601            datums
602                .iter()
603                .map(|(key, value)| (key.as_str(), Datum::from(value.as_str()))),
604        );
605        Ok(MirScalarExpr::literal_from_single_element_row(
606            row,
607            f.output_type().scalar_type,
608        ))
609    };
610    let pack = |datum| {
611        Ok(MirScalarExpr::literal_ok(
612            datum,
613            f.output_type().scalar_type,
614        ))
615    };
616
617    match f {
618        UnmaterializableFunc::CurrentDatabase => pack(Datum::from(session.database())),
619        UnmaterializableFunc::CurrentSchema => {
620            let search_path = state.resolve_search_path(session);
621            let schema = search_path
622                .first()
623                .map(|(db, schema)| &*state.get_schema(db, schema, session.conn_id()).name.schema);
624            pack(Datum::from(schema))
625        }
626        UnmaterializableFunc::CurrentSchemasWithSystem => {
627            let search_path = state.resolve_search_path(session);
628            let search_path = state.effective_search_path(&search_path, false);
629            pack_1d_array(
630                search_path
631                    .into_iter()
632                    .map(|(db, schema)| {
633                        let schema = state.get_schema(&db, &schema, session.conn_id());
634                        Datum::String(&schema.name.schema)
635                    })
636                    .collect(),
637            )
638        }
639        UnmaterializableFunc::CurrentSchemasWithoutSystem => {
640            let search_path = state.resolve_search_path(session);
641            pack_1d_array(
642                search_path
643                    .into_iter()
644                    .map(|(db, schema)| {
645                        let schema = state.get_schema(&db, &schema, session.conn_id());
646                        Datum::String(&schema.name.schema)
647                    })
648                    .collect(),
649            )
650        }
651        UnmaterializableFunc::ViewableVariables => pack_dict(
652            viewable_variables(state, session)
653                .map(|var| (var.name().to_lowercase(), var.value()))
654                .collect(),
655        ),
656        UnmaterializableFunc::CurrentTimestamp => {
657            let t: Datum = session.pcx().wall_time.try_into()?;
658            pack(t)
659        }
660        UnmaterializableFunc::CurrentUser => pack(Datum::from(
661            state.get_role(session.current_role_id()).name(),
662        )),
663        UnmaterializableFunc::SessionUser => pack(Datum::from(
664            state.get_role(session.session_role_id()).name(),
665        )),
666        UnmaterializableFunc::IsRbacEnabled => pack(Datum::from(
667            rbac::is_rbac_enabled_for_session(state.system_config(), session),
668        )),
669        UnmaterializableFunc::MzEnvironmentId => {
670            pack(Datum::from(&*state.config().environment_id.to_string()))
671        }
672        UnmaterializableFunc::MzIsSuperuser => pack(Datum::from(session.is_superuser())),
673        UnmaterializableFunc::MzNow => match logical_time {
674            EvalTime::Time(logical_time) => pack(Datum::MzTimestamp(logical_time)),
675            EvalTime::NotAvailable => Err(OptimizerError::UncallableFunction {
676                func: UnmaterializableFunc::MzNow,
677                context: "this",
678            }),
679        },
680        UnmaterializableFunc::MzRoleOidMemberships => {
681            let role_memberships = role_oid_memberships(state);
682            let mut role_memberships: Vec<(_, Vec<_>)> = role_memberships
683                .into_iter()
684                .map(|(role_id, role_membership)| {
685                    (
686                        role_id.to_string(),
687                        role_membership
688                            .into_iter()
689                            .map(|role_id| role_id.to_string())
690                            .collect(),
691                    )
692                })
693                .collect();
694            role_memberships.sort();
695            let mut row = Row::default();
696            row.packer().push_dict_with(|row| {
697                for (role_id, role_membership) in &role_memberships {
698                    row.push(Datum::from(role_id.as_str()));
699                    row.try_push_array(
700                        &[ArrayDimension {
701                            lower_bound: 1,
702                            length: role_membership.len(),
703                        }],
704                        role_membership.iter().map(|role_id| Datum::from(role_id.as_str())),
705                    ).expect("role_membership is 1 dimensional, and its length is used for the array length");
706                }
707            });
708            Ok(MirScalarExpr::literal_from_single_element_row(
709                row,
710                f.output_type().scalar_type,
711            ))
712        }
713        UnmaterializableFunc::MzSessionId => pack(Datum::from(state.config().session_id)),
714        UnmaterializableFunc::MzUptime => {
715            let uptime = state.config().start_instant.elapsed();
716            let uptime = chrono::Duration::from_std(uptime).map_or(Datum::Null, Datum::from);
717            pack(uptime)
718        }
719        UnmaterializableFunc::MzVersion => pack(Datum::from(
720            &*state
721                .config()
722                .build_info
723                .human_version(state.config().helm_chart_version.clone()),
724        )),
725        UnmaterializableFunc::MzVersionNum => {
726            pack(Datum::Int32(state.config().build_info.version_num()))
727        }
728        UnmaterializableFunc::PgBackendPid => pack(Datum::Int32(i32::reinterpret_cast(
729            session.conn_id().unhandled(),
730        ))),
731        UnmaterializableFunc::PgPostmasterStartTime => {
732            let t: Datum = state.config().start_time.try_into()?;
733            pack(t)
734        }
735        UnmaterializableFunc::Version => {
736            let build_info = state.config().build_info;
737            let version = format!(
738                "PostgreSQL {}.{} on {} (Materialize {})",
739                SERVER_MAJOR_VERSION,
740                SERVER_MINOR_VERSION,
741                mz_build_info::TARGET_TRIPLE,
742                build_info.version,
743            );
744            pack(Datum::from(&*version))
745        }
746    }
747}
748
749fn role_oid_memberships<'a>(catalog: &'a CatalogState) -> BTreeMap<u32, BTreeSet<u32>> {
750    let mut role_memberships = BTreeMap::new();
751    for role_id in catalog.get_roles() {
752        let role = catalog.get_role(role_id);
753        if !role_memberships.contains_key(&role.oid) {
754            role_oid_memberships_inner(catalog, role_id, &mut role_memberships);
755        }
756    }
757    role_memberships
758}
759
760fn role_oid_memberships_inner<'a>(
761    catalog: &'a CatalogState,
762    role_id: &RoleId,
763    role_memberships: &mut BTreeMap<u32, BTreeSet<u32>>,
764) {
765    let role = catalog.get_role(role_id);
766    role_memberships.insert(role.oid, btreeset! {role.oid});
767    for parent_role_id in role.membership.map.keys() {
768        let parent_role = catalog.get_role(parent_role_id);
769        if !role_memberships.contains_key(&parent_role.oid) {
770            role_oid_memberships_inner(catalog, parent_role_id, role_memberships);
771        }
772        let parent_membership: BTreeSet<_> = role_memberships
773            .get(&parent_role.oid)
774            .expect("inserted in recursive call above")
775            .into_iter()
776            .cloned()
777            .collect();
778        role_memberships
779            .get_mut(&role.oid)
780            .expect("inserted above")
781            .extend(parent_membership);
782    }
783}
mz_adapter/optimize/dataflows.rs

mz_adapter/optimize/
dataflows.rs
