1use 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#[derive(Debug, Clone)]
55pub struct ComputeInstanceSnapshot {
56 instance_id: ComputeInstanceId,
57 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 pub fn instance_id(&self) -> ComputeInstanceId {
91 self.instance_id
92 }
93
94 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 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#[derive(Debug)]
113pub struct DataflowBuilder<'a> {
114 pub catalog: &'a dyn OptimizerCatalog,
115 pub compute: ComputeInstanceSnapshot,
120 pub replan: Option<GlobalId>,
129 recursion_guard: RecursionGuard,
131}
132
133pub trait ExprPrep {
135 fn prep_relation_expr(&self, expr: &mut OptimizedMirRelationExpr)
137 -> Result<(), OptimizerError>;
138
139 fn prep_scalar_expr(&self, expr: &mut MirScalarExpr) -> Result<(), OptimizerError>;
141}
142
143pub 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
154pub 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 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 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
199pub 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 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
233pub struct ExprPrepWebhookValidation {
236 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 NotAvailable,
269}
270
271pub 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 pub(super) fn with_config(mut self, config: &OptimizerConfig) -> Self {
299 self.replan = config.replan;
300 self
301 }
302
303 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 if dataflow.is_imported(id) {
317 return Ok(());
318 }
319
320 let monotonic = self.monotonic_object(*id, features);
321
322 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 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 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::Sink(_)
371 | CatalogItem::Index(_)
372 | CatalogItem::Type(_)
373 | CatalogItem::Func(_)
374 | CatalogItem::Secret(_)
375 | CatalogItem::Connection(_) => {
376 unreachable!()
378 }
379 }
380 }
381 Ok(())
382 })
383 }
384
385 pub fn import_view_into_dataflow(
395 &mut self,
396 view_id: &GlobalId,
397 view: &OptimizedMirRelationExpr,
398 dataflow: &mut DataflowDesc,
399 features: &OptimizerFeatures,
400 ) -> Result<(), OptimizerError> {
401 for get_id in view.depends_on() {
402 self.import_into_dataflow(&get_id, dataflow, features)?;
403 }
404 dataflow.insert_plan(*view_id, view.clone());
405 Ok(())
406 }
407
408 pub fn maybe_reoptimize_imported_views(
411 &self,
412 df_desc: &mut DataflowDesc,
413 config: &OptimizerConfig,
414 ) -> Result<(), OptimizerError> {
415 if !config.features.reoptimize_imported_views {
416 return Ok(()); }
418
419 let mut view_optimizer = view::Optimizer::new(config.clone(), None);
420 for desc in df_desc.objects_to_build.iter_mut().rev() {
421 if matches!(desc.id, GlobalId::Explain | GlobalId::Transient(_)) {
422 continue; }
424 if let CatalogItem::View(view) = &self.catalog.get_entry(&desc.id).item {
425 let _span = tracing::span!(
426 target: "optimizer",
427 tracing::Level::DEBUG,
428 "view",
429 path.segment = desc.id.to_string()
430 )
431 .entered();
432
433 desc.plan = view_optimizer.optimize(view.raw_expr.as_ref().clone())?;
435
436 trace_plan(desc.plan.as_inner());
438 }
439 }
440
441 Ok(())
442 }
443
444 fn monotonic_source(&self, data_source: &DataSourceDesc) -> bool {
446 match data_source {
447 DataSourceDesc::Ingestion { .. } => false,
448 DataSourceDesc::OldSyntaxIngestion {
449 desc, data_config, ..
450 } => data_config.monotonic(&desc.connection),
451 DataSourceDesc::Webhook { .. } => true,
452 DataSourceDesc::IngestionExport {
453 ingestion_id,
454 data_config,
455 ..
456 } => {
457 let source_desc = self
458 .catalog
459 .get_entry_by_item_id(ingestion_id)
460 .source_desc()
461 .expect("ingestion export must reference a source")
462 .expect("ingestion export must reference a source");
463 data_config.monotonic(&source_desc.connection)
464 }
465 DataSourceDesc::Introspection(_)
466 | DataSourceDesc::Progress
467 | DataSourceDesc::Catalog => false,
468 }
469 }
470
471 fn monotonic_object(&self, id: GlobalId, features: &OptimizerFeatures) -> bool {
481 self.monotonic_object_inner(id, &mut BTreeMap::new(), features)
482 .unwrap_or_else(|e| {
483 warn!(%id, "error inspecting object for monotonicity: {e}");
484 false
485 })
486 }
487
488 fn monotonic_object_inner(
489 &self,
490 id: GlobalId,
491 memo: &mut BTreeMap<GlobalId, bool>,
492 features: &OptimizerFeatures,
493 ) -> Result<bool, RecursionLimitError> {
494 if let Some(monotonic) = memo.get(&id) {
497 return Ok(*monotonic);
498 }
499
500 let monotonic = self.checked_recur(|_| {
501 match self.catalog.get_entry(&id).item() {
502 CatalogItem::Source(source) => Ok(self.monotonic_source(&source.data_source)),
503 CatalogItem::Table(table) => match &table.data_source {
504 TableDataSource::TableWrites { .. } => Ok(false),
505 TableDataSource::DataSource { desc, timeline: _ } => {
506 Ok(self.monotonic_source(desc))
507 }
508 },
509 CatalogItem::View(View {
510 locally_optimized_expr: optimized_expr,
511 ..
512 }) => {
513 let view_expr = optimized_expr.as_ref().clone().into_inner();
514
515 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 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(_) => Ok(false),
556 }
557 })?;
558
559 memo.insert(id, monotonic);
560
561 Ok(monotonic)
562 }
563}
564
565impl<'a> CheckedRecursion for DataflowBuilder<'a> {
566 fn recursion_guard(&self) -> &RecursionGuard {
567 &self.recursion_guard
568 }
569}
570
571fn eval_unmaterializable_func(
572 state: &CatalogState,
573 f: &UnmaterializableFunc,
574 logical_time: EvalTime,
575 session: &dyn SessionMetadata,
576) -> Result<MirScalarExpr, OptimizerError> {
577 let pack_1d_array = |datums: Vec<Datum>| {
578 let mut row = Row::default();
579 row.packer()
580 .try_push_array(
581 &[ArrayDimension {
582 lower_bound: 1,
583 length: datums.len(),
584 }],
585 datums,
586 )
587 .expect("known to be a valid array");
588 Ok(MirScalarExpr::literal_from_single_element_row(
589 row,
590 f.output_type().scalar_type,
591 ))
592 };
593 let pack_dict = |mut datums: Vec<(String, String)>| {
594 datums.sort();
595 let mut row = Row::default();
596 row.packer().push_dict(
597 datums
598 .iter()
599 .map(|(key, value)| (key.as_str(), Datum::from(value.as_str()))),
600 );
601 Ok(MirScalarExpr::literal_from_single_element_row(
602 row,
603 f.output_type().scalar_type,
604 ))
605 };
606 let pack = |datum| {
607 Ok(MirScalarExpr::literal_ok(
608 datum,
609 f.output_type().scalar_type,
610 ))
611 };
612
613 match f {
614 UnmaterializableFunc::CurrentDatabase => pack(Datum::from(session.database())),
615 UnmaterializableFunc::CurrentSchema => {
616 let search_path = state.resolve_search_path(session);
617 let schema = search_path
618 .first()
619 .map(|(db, schema)| &*state.get_schema(db, schema, session.conn_id()).name.schema);
620 pack(Datum::from(schema))
621 }
622 UnmaterializableFunc::CurrentSchemasWithSystem => {
623 let search_path = state.resolve_search_path(session);
624 let search_path = state.effective_search_path(&search_path, false);
625 pack_1d_array(
626 search_path
627 .into_iter()
628 .map(|(db, schema)| {
629 let schema = state.get_schema(&db, &schema, session.conn_id());
630 Datum::String(&schema.name.schema)
631 })
632 .collect(),
633 )
634 }
635 UnmaterializableFunc::CurrentSchemasWithoutSystem => {
636 let search_path = state.resolve_search_path(session);
637 pack_1d_array(
638 search_path
639 .into_iter()
640 .map(|(db, schema)| {
641 let schema = state.get_schema(&db, &schema, session.conn_id());
642 Datum::String(&schema.name.schema)
643 })
644 .collect(),
645 )
646 }
647 UnmaterializableFunc::ViewableVariables => pack_dict(
648 viewable_variables(state, session)
649 .map(|var| (var.name().to_lowercase(), var.value()))
650 .collect(),
651 ),
652 UnmaterializableFunc::CurrentTimestamp => {
653 let t: Datum = session.pcx().wall_time.try_into()?;
654 pack(t)
655 }
656 UnmaterializableFunc::CurrentUser => pack(Datum::from(
657 state.get_role(session.current_role_id()).name(),
658 )),
659 UnmaterializableFunc::SessionUser => pack(Datum::from(
660 state.get_role(session.session_role_id()).name(),
661 )),
662 UnmaterializableFunc::IsRbacEnabled => pack(Datum::from(
663 rbac::is_rbac_enabled_for_session(state.system_config(), session),
664 )),
665 UnmaterializableFunc::MzEnvironmentId => {
666 pack(Datum::from(&*state.config().environment_id.to_string()))
667 }
668 UnmaterializableFunc::MzIsSuperuser => pack(Datum::from(session.is_superuser())),
669 UnmaterializableFunc::MzNow => match logical_time {
670 EvalTime::Time(logical_time) => pack(Datum::MzTimestamp(logical_time)),
671 EvalTime::NotAvailable => Err(OptimizerError::UncallableFunction {
672 func: UnmaterializableFunc::MzNow,
673 context: "this",
674 }),
675 },
676 UnmaterializableFunc::MzRoleOidMemberships => {
677 let role_memberships = role_oid_memberships(state);
678 let mut role_memberships: Vec<(_, Vec<_>)> = role_memberships
679 .into_iter()
680 .map(|(role_id, role_membership)| {
681 (
682 role_id.to_string(),
683 role_membership
684 .into_iter()
685 .map(|role_id| role_id.to_string())
686 .collect(),
687 )
688 })
689 .collect();
690 role_memberships.sort();
691 let mut row = Row::default();
692 row.packer().push_dict_with(|row| {
693 for (role_id, role_membership) in &role_memberships {
694 row.push(Datum::from(role_id.as_str()));
695 row.try_push_array(
696 &[ArrayDimension {
697 lower_bound: 1,
698 length: role_membership.len(),
699 }],
700 role_membership.iter().map(|role_id| Datum::from(role_id.as_str())),
701 ).expect("role_membership is 1 dimensional, and its length is used for the array length");
702 }
703 });
704 Ok(MirScalarExpr::literal_from_single_element_row(
705 row,
706 f.output_type().scalar_type,
707 ))
708 }
709 UnmaterializableFunc::MzSessionId => pack(Datum::from(state.config().session_id)),
710 UnmaterializableFunc::MzUptime => {
711 let uptime = state.config().start_instant.elapsed();
712 let uptime = chrono::Duration::from_std(uptime).map_or(Datum::Null, Datum::from);
713 pack(uptime)
714 }
715 UnmaterializableFunc::MzVersion => pack(Datum::from(
716 &*state
717 .config()
718 .build_info
719 .human_version(state.config().helm_chart_version.clone()),
720 )),
721 UnmaterializableFunc::MzVersionNum => {
722 pack(Datum::Int32(state.config().build_info.version_num()))
723 }
724 UnmaterializableFunc::PgBackendPid => pack(Datum::Int32(i32::reinterpret_cast(
725 session.conn_id().unhandled(),
726 ))),
727 UnmaterializableFunc::PgPostmasterStartTime => {
728 let t: Datum = state.config().start_time.try_into()?;
729 pack(t)
730 }
731 UnmaterializableFunc::Version => {
732 let build_info = state.config().build_info;
733 let version = format!(
734 "PostgreSQL {}.{} on {} (Materialize {})",
735 SERVER_MAJOR_VERSION,
736 SERVER_MINOR_VERSION,
737 mz_build_info::TARGET_TRIPLE,
738 build_info.version,
739 );
740 pack(Datum::from(&*version))
741 }
742 }
743}
744
745fn role_oid_memberships<'a>(catalog: &'a CatalogState) -> BTreeMap<u32, BTreeSet<u32>> {
746 let mut role_memberships = BTreeMap::new();
747 for role_id in catalog.get_roles() {
748 let role = catalog.get_role(role_id);
749 if !role_memberships.contains_key(&role.oid) {
750 role_oid_memberships_inner(catalog, role_id, &mut role_memberships);
751 }
752 }
753 role_memberships
754}
755
756fn role_oid_memberships_inner<'a>(
757 catalog: &'a CatalogState,
758 role_id: &RoleId,
759 role_memberships: &mut BTreeMap<u32, BTreeSet<u32>>,
760) {
761 let role = catalog.get_role(role_id);
762 role_memberships.insert(role.oid, btreeset! {role.oid});
763 for parent_role_id in role.membership.map.keys() {
764 let parent_role = catalog.get_role(parent_role_id);
765 if !role_memberships.contains_key(&parent_role.oid) {
766 role_oid_memberships_inner(catalog, parent_role_id, role_memberships);
767 }
768 let parent_membership: BTreeSet<_> = role_memberships
769 .get(&parent_role.oid)
770 .expect("inserted in recursive call above")
771 .into_iter()
772 .cloned()
773 .collect();
774 role_memberships
775 .get_mut(&role.oid)
776 .expect("inserted above")
777 .extend(parent_membership);
778 }
779}