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.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.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 { optimized_expr, .. }) => {
513                    let view_expr = optimized_expr.as_ref().clone().into_inner();
514
515                    // Inspect global ids that occur in the Gets in view_expr, and collect the ids
516                    // of monotonic dependees.
517                    let mut monotonic_ids = BTreeSet::new();
518                    let recursion_result: Result<(), RecursionLimitError> = view_expr
519                        .try_visit_post(&mut |e| {
520                            if let MirRelationExpr::Get {
521                                id: Id::Global(got_id),
522                                ..
523                            } = e
524                            {
525                                if self.monotonic_object_inner(*got_id, memo, features)? {
526                                    monotonic_ids.insert(*got_id);
527                                }
528                            }
529                            Ok(())
530                        });
531                    if let Err(error) = recursion_result {
532                        // We still might have got some of the IDs, so just log and continue. Now
533                        // the subsequent monotonicity analysis can have false negatives.
534                        warn!(%id, "error inspecting view for monotonicity: {error}");
535                    }
536
537                    let mut builder = DerivedBuilder::new(features);
538                    builder.require(Monotonic::new(monotonic_ids.clone()));
539                    let derived = builder.visit(&view_expr);
540
541                    Ok(*derived
542                        .as_view()
543                        .value::<Monotonic>()
544                        .expect("Expected monotonic result from non empty tree"))
545                }
546                CatalogItem::Index(Index { on, .. }) => {
547                    self.monotonic_object_inner(*on, memo, features)
548                }
549                CatalogItem::Secret(_)
550                | CatalogItem::Type(_)
551                | CatalogItem::Connection(_)
552                | CatalogItem::Log(_)
553                | CatalogItem::MaterializedView(_)
554                | CatalogItem::Sink(_)
555                | CatalogItem::Func(_)
556                | CatalogItem::ContinualTask(_) => Ok(false),
557            }
558        })?;
559
560        memo.insert(id, monotonic);
561
562        Ok(monotonic)
563    }
564}
565
566impl<'a> CheckedRecursion for DataflowBuilder<'a> {
567    fn recursion_guard(&self) -> &RecursionGuard {
568        &self.recursion_guard
569    }
570}
571
572fn eval_unmaterializable_func(
573    state: &CatalogState,
574    f: &UnmaterializableFunc,
575    logical_time: EvalTime,
576    session: &dyn SessionMetadata,
577) -> Result<MirScalarExpr, OptimizerError> {
578    let pack_1d_array = |datums: Vec<Datum>| {
579        let mut row = Row::default();
580        row.packer()
581            .try_push_array(
582                &[ArrayDimension {
583                    lower_bound: 1,
584                    length: datums.len(),
585                }],
586                datums,
587            )
588            .expect("known to be a valid array");
589        Ok(MirScalarExpr::literal_from_single_element_row(
590            row,
591            f.output_type().scalar_type,
592        ))
593    };
594    let pack_dict = |mut datums: Vec<(String, String)>| {
595        datums.sort();
596        let mut row = Row::default();
597        row.packer().push_dict(
598            datums
599                .iter()
600                .map(|(key, value)| (key.as_str(), Datum::from(value.as_str()))),
601        );
602        Ok(MirScalarExpr::literal_from_single_element_row(
603            row,
604            f.output_type().scalar_type,
605        ))
606    };
607    let pack = |datum| {
608        Ok(MirScalarExpr::literal_ok(
609            datum,
610            f.output_type().scalar_type,
611        ))
612    };
613
614    match f {
615        UnmaterializableFunc::CurrentDatabase => pack(Datum::from(session.database())),
616        UnmaterializableFunc::CurrentSchema => {
617            let search_path = state.resolve_search_path(session);
618            let schema = search_path
619                .first()
620                .map(|(db, schema)| &*state.get_schema(db, schema, session.conn_id()).name.schema);
621            pack(Datum::from(schema))
622        }
623        UnmaterializableFunc::CurrentSchemasWithSystem => {
624            let search_path = state.resolve_search_path(session);
625            let search_path = state.effective_search_path(&search_path, false);
626            pack_1d_array(
627                search_path
628                    .into_iter()
629                    .map(|(db, schema)| {
630                        let schema = state.get_schema(&db, &schema, session.conn_id());
631                        Datum::String(&schema.name.schema)
632                    })
633                    .collect(),
634            )
635        }
636        UnmaterializableFunc::CurrentSchemasWithoutSystem => {
637            let search_path = state.resolve_search_path(session);
638            pack_1d_array(
639                search_path
640                    .into_iter()
641                    .map(|(db, schema)| {
642                        let schema = state.get_schema(&db, &schema, session.conn_id());
643                        Datum::String(&schema.name.schema)
644                    })
645                    .collect(),
646            )
647        }
648        UnmaterializableFunc::ViewableVariables => pack_dict(
649            viewable_variables(state, session)
650                .map(|var| (var.name().to_lowercase(), var.value()))
651                .collect(),
652        ),
653        UnmaterializableFunc::CurrentTimestamp => {
654            let t: Datum = session.pcx().wall_time.try_into()?;
655            pack(t)
656        }
657        UnmaterializableFunc::CurrentUser => pack(Datum::from(
658            state.get_role(session.current_role_id()).name(),
659        )),
660        UnmaterializableFunc::SessionUser => pack(Datum::from(
661            state.get_role(session.session_role_id()).name(),
662        )),
663        UnmaterializableFunc::IsRbacEnabled => pack(Datum::from(
664            rbac::is_rbac_enabled_for_session(state.system_config(), session),
665        )),
666        UnmaterializableFunc::MzEnvironmentId => {
667            pack(Datum::from(&*state.config().environment_id.to_string()))
668        }
669        UnmaterializableFunc::MzIsSuperuser => pack(Datum::from(session.is_superuser())),
670        UnmaterializableFunc::MzNow => match logical_time {
671            EvalTime::Time(logical_time) => pack(Datum::MzTimestamp(logical_time)),
672            EvalTime::NotAvailable => Err(OptimizerError::UncallableFunction {
673                func: UnmaterializableFunc::MzNow,
674                context: "this",
675            }),
676        },
677        UnmaterializableFunc::MzRoleOidMemberships => {
678            let role_memberships = role_oid_memberships(state);
679            let mut role_memberships: Vec<(_, Vec<_>)> = role_memberships
680                .into_iter()
681                .map(|(role_id, role_membership)| {
682                    (
683                        role_id.to_string(),
684                        role_membership
685                            .into_iter()
686                            .map(|role_id| role_id.to_string())
687                            .collect(),
688                    )
689                })
690                .collect();
691            role_memberships.sort();
692            let mut row = Row::default();
693            row.packer().push_dict_with(|row| {
694                for (role_id, role_membership) in &role_memberships {
695                    row.push(Datum::from(role_id.as_str()));
696                    row.try_push_array(
697                        &[ArrayDimension {
698                            lower_bound: 1,
699                            length: role_membership.len(),
700                        }],
701                        role_membership.iter().map(|role_id| Datum::from(role_id.as_str())),
702                    ).expect("role_membership is 1 dimensional, and its length is used for the array length");
703                }
704            });
705            Ok(MirScalarExpr::literal_from_single_element_row(
706                row,
707                f.output_type().scalar_type,
708            ))
709        }
710        UnmaterializableFunc::MzSessionId => pack(Datum::from(state.config().session_id)),
711        UnmaterializableFunc::MzUptime => {
712            let uptime = state.config().start_instant.elapsed();
713            let uptime = chrono::Duration::from_std(uptime).map_or(Datum::Null, Datum::from);
714            pack(uptime)
715        }
716        UnmaterializableFunc::MzVersion => pack(Datum::from(
717            &*state
718                .config()
719                .build_info
720                .human_version(state.config().helm_chart_version.clone()),
721        )),
722        UnmaterializableFunc::MzVersionNum => {
723            pack(Datum::Int32(state.config().build_info.version_num()))
724        }
725        UnmaterializableFunc::PgBackendPid => pack(Datum::Int32(i32::reinterpret_cast(
726            session.conn_id().unhandled(),
727        ))),
728        UnmaterializableFunc::PgPostmasterStartTime => {
729            let t: Datum = state.config().start_time.try_into()?;
730            pack(t)
731        }
732        UnmaterializableFunc::Version => {
733            let build_info = state.config().build_info;
734            let version = format!(
735                "PostgreSQL {}.{} on {} (Materialize {})",
736                SERVER_MAJOR_VERSION,
737                SERVER_MINOR_VERSION,
738                mz_build_info::TARGET_TRIPLE,
739                build_info.version,
740            );
741            pack(Datum::from(&*version))
742        }
743    }
744}
745
746fn role_oid_memberships<'a>(catalog: &'a CatalogState) -> BTreeMap<u32, BTreeSet<u32>> {
747    let mut role_memberships = BTreeMap::new();
748    for role_id in catalog.get_roles() {
749        let role = catalog.get_role(role_id);
750        if !role_memberships.contains_key(&role.oid) {
751            role_oid_memberships_inner(catalog, role_id, &mut role_memberships);
752        }
753    }
754    role_memberships
755}
756
757fn role_oid_memberships_inner<'a>(
758    catalog: &'a CatalogState,
759    role_id: &RoleId,
760    role_memberships: &mut BTreeMap<u32, BTreeSet<u32>>,
761) {
762    let role = catalog.get_role(role_id);
763    role_memberships.insert(role.oid, btreeset! {role.oid});
764    for parent_role_id in role.membership.map.keys() {
765        let parent_role = catalog.get_role(parent_role_id);
766        if !role_memberships.contains_key(&parent_role.oid) {
767            role_oid_memberships_inner(catalog, parent_role_id, role_memberships);
768        }
769        let parent_membership: BTreeSet<_> = role_memberships
770            .get(&parent_role.oid)
771            .expect("inserted in recursive call above")
772            .into_iter()
773            .cloned()
774            .collect();
775        role_memberships
776            .get_mut(&role.oid)
777            .expect("inserted above")
778            .extend(parent_membership);
779    }
780}
mz_adapter/optimize/dataflows.rs

mz_adapter/optimize/
dataflows.rs
