adapter,sql: plan and sequence replacement MV application

Author: teskje

src/adapter/src/catalog/transact.rs

@@ -100,6 +100,10 @@ pub enum Op {
         typ: SqlColumnType,
         sql: RawDataType,
     },
+    AlterMaterializedViewApplyReplacement {
+        id: CatalogItemId,
+        replacement_id: CatalogItemId,
+    },
     CreateDatabase {
         name: String,
         owner_id: RoleId,
@@ -769,6 +773,28 @@ impl Catalog {
                 tx.update_item(id, new_entry.into())?;
                 storage_collections_to_register.insert(new_global_id, shard_id);
             }
+            Op::AlterMaterializedViewApplyReplacement { id, replacement_id } => {
+                let mut new_entry = state.get_entry(&id).clone();
+                let replacement = state.get_entry(&replacement_id);
+
+                let CatalogItem::MaterializedView(mv) = &mut new_entry.item else {
+                    return Err(AdapterError::internal(
+                        "ALTER MATERIALIZED VIEW ... APPLY REPLACEMENT",
+                        "id must refer to a materialized view",
+                    ));
+                };
+                let CatalogItem::MaterializedView(replacement_mv) = &replacement.item else {
+                    return Err(AdapterError::internal(
+                        "ALTER MATERIALIZED VIEW ... APPLY REPLACEMENT",
+                        "replacement_id must refer to a materialized view",
+                    ));
+                };
+
+                mv.apply_replacement(replacement_mv.clone());
+
+                tx.remove_item(replacement_id)?;
+                tx.update_item(id, new_entry.into())?;
+            }
             Op::CreateDatabase { name, owner_id } => {
                 let database_owner_privileges = vec![rbac::owner_privilege(
                     mz_sql::catalog::ObjectType::Database,

src/adapter/src/command.rs

@@ -644,6 +644,7 @@ impl ExecuteResponse {
             | AlterSource
             | AlterSink
             | AlterTableAddColumn
+            | AlterMaterializedViewApplyReplacement
             | AlterNetworkPolicy => &[AlteredObject],
             AlterDefaultPrivileges => &[AlteredDefaultPrivileges],
             AlterSetCluster => &[AlteredObject],

src/adapter/src/coord/appends.rs

@@ -1004,6 +1004,7 @@ pub(crate) fn waiting_on_startup_appends(
         | Plan::AlterRole(_)
         | Plan::AlterOwner(_)
         | Plan::AlterTableAddColumn(_)
+        | Plan::AlterMaterializedViewApplyReplacement(_)
         | Plan::Declare(_)
         | Plan::Fetch(_)
         | Plan::Close(_)

src/adapter/src/coord/catalog_serving.rs

@@ -128,6 +128,7 @@ pub fn auto_run_on_catalog_server<'a, 's, 'p>(
         | Plan::AlterRole(_)
         | Plan::AlterOwner(_)
         | Plan::AlterTableAddColumn(_)
+        | Plan::AlterMaterializedViewApplyReplacement(_)
         | Plan::Declare(_)
         | Plan::Fetch(_)
         | Plan::Close(_)

src/adapter/src/coord/ddl.rs

@@ -214,9 +214,7 @@ impl Coordinator {
         let mut webhook_sources_to_restart = BTreeSet::new();
         let mut table_gids_to_drop = vec![];
         let mut storage_sink_gids_to_drop = vec![];
-        let mut indexes_to_drop = vec![];
-        let mut materialized_views_to_drop = vec![];
-        let mut continual_tasks_to_drop = vec![];
+        let mut compute_gids_to_drop = vec![];
         let mut views_to_drop = vec![];
         let mut replication_slots_to_drop: Vec<(PostgresConnection, String)> = vec![];
         let mut secrets_to_drop = vec![];
@@ -275,21 +273,21 @@ impl Coordinator {
                                         storage_sink_gids_to_drop.push(sink.global_id());
                                     }
                                     CatalogItem::Index(index) => {
-                                        indexes_to_drop.push((index.cluster_id, index.global_id()));
+                                        compute_gids_to_drop
+                                            .push((index.cluster_id, index.global_id()));
                                     }
                                     CatalogItem::MaterializedView(mv) => {
-                                        materialized_views_to_drop
+                                        compute_gids_to_drop
                                             .push((mv.cluster_id, mv.global_id_writes()));
+                                        sources_to_drop
+                                            .extend(mv.global_ids().map(|gid| (*id, gid)));
                                     }
                                     CatalogItem::View(view) => {
                                         views_to_drop.push((*id, view.clone()))
                                     }
                                     CatalogItem::ContinualTask(ct) => {
-                                        continual_tasks_to_drop.push((
-                                            *id,
-                                            ct.cluster_id,
-                                            ct.global_id(),
-                                        ));
+                                        compute_gids_to_drop.push((ct.cluster_id, ct.global_id()));
+                                        sources_to_drop.push((*id, ct.global_id()));
                                     }
                                     CatalogItem::Secret(_) => {
                                         secrets_to_drop.push(*id);
@@ -422,9 +420,7 @@ impl Coordinator {
             .map(|(_, gid)| *gid)
             .chain(table_gids_to_drop.iter().map(|(_, gid)| *gid))
             .chain(storage_sink_gids_to_drop.iter().copied())
-            .chain(indexes_to_drop.iter().map(|(_, gid)| *gid))
-            .chain(materialized_views_to_drop.iter().map(|(_, gid)| *gid))
-            .chain(continual_tasks_to_drop.iter().map(|(_, _, gid)| *gid))
+            .chain(compute_gids_to_drop.iter().map(|(_, gid)| *gid))
             .chain(views_to_drop.iter().map(|(_id, view)| view.global_id()))
             .collect();
 
@@ -494,30 +490,19 @@ impl Coordinator {
             }
         }
 
-        let storage_ids_to_drop = sources_to_drop
+        let storage_gids_to_drop = sources_to_drop
             .iter()
             .map(|(_, gid)| *gid)
             .chain(storage_sink_gids_to_drop.iter().copied())
-            .chain(table_gids_to_drop.iter().map(|(_, gid)| *gid))
-            .chain(materialized_views_to_drop.iter().map(|(_, gid)| *gid))
-            .chain(continual_tasks_to_drop.iter().map(|(_, _, gid)| *gid));
-        let compute_ids_to_drop = indexes_to_drop
-            .iter()
-            .copied()
-            .chain(materialized_views_to_drop.iter().copied())
-            .chain(
-                continual_tasks_to_drop
-                    .iter()
-                    .map(|(_, cluster_id, gid)| (*cluster_id, *gid)),
-            );
+            .chain(table_gids_to_drop.iter().map(|(_, gid)| *gid));
 
         // Check if any Timelines would become empty, if we dropped the specified storage or
         // compute resources.
         //
         // Note: only after a Transaction succeeds do we actually drop the timeline
         let collection_id_bundle = self.build_collection_id_bundle(
-            storage_ids_to_drop,
-            compute_ids_to_drop,
+            storage_gids_to_drop,
+            compute_gids_to_drop.clone(),
             clusters_to_drop.clone(),
         );
         let timeline_associations: BTreeMap<_, _> = self
@@ -655,14 +640,8 @@ impl Coordinator {
                     self.cancel_pending_copy(&conn_id);
                 }
             }
-            if !indexes_to_drop.is_empty() {
-                self.drop_indexes(indexes_to_drop);
-            }
-            if !materialized_views_to_drop.is_empty() {
-                self.drop_materialized_views(materialized_views_to_drop);
-            }
-            if !continual_tasks_to_drop.is_empty() {
-                self.drop_continual_tasks(continual_tasks_to_drop);
+            if !compute_gids_to_drop.is_empty() {
+                self.drop_compute_collections(compute_gids_to_drop);
             }
             if !vpc_endpoints_to_drop.is_empty() {
                 self.drop_vpc_endpoints_in_background(vpc_endpoints_to_drop)
@@ -1030,7 +1009,7 @@ impl Coordinator {
             .unwrap_or_terminate("cannot fail to drop sinks");
     }
 
-    pub(crate) fn drop_indexes(&mut self, indexes: Vec<(ClusterId, GlobalId)>) {
+    pub(crate) fn drop_compute_collections(&mut self, indexes: Vec<(ClusterId, GlobalId)>) {
         let mut by_cluster: BTreeMap<_, Vec<_>> = BTreeMap::new();
         for (cluster_id, gid) in indexes {
             by_cluster.entry(cluster_id).or_default().push(gid);
@@ -1046,58 +1025,6 @@ impl Coordinator {
         }
     }
 
-    /// A convenience method for dropping materialized views.
-    fn drop_materialized_views(&mut self, mviews: Vec<(ClusterId, GlobalId)>) {
-        let mut by_cluster: BTreeMap<_, Vec<_>> = BTreeMap::new();
-        let mut mv_gids = Vec::new();
-        for (cluster_id, gid) in mviews {
-            by_cluster.entry(cluster_id).or_default().push(gid);
-            mv_gids.push(gid);
-        }
-
-        // Drop compute sinks.
-        for (cluster_id, ids) in by_cluster {
-            let compute = &mut self.controller.compute;
-            // A cluster could have been dropped, so verify it exists.
-            if compute.instance_exists(cluster_id) {
-                compute
-                    .drop_collections(cluster_id, ids)
-                    .unwrap_or_terminate("cannot fail to drop collections");
-            }
-        }
-
-        // Drop storage resources.
-        let storage_metadata = self.catalog.state().storage_metadata();
-        self.controller
-            .storage
-            .drop_sources(storage_metadata, mv_gids)
-            .unwrap_or_terminate("cannot fail to drop sources");
-    }
-
-    /// A convenience method for dropping continual tasks.
-    fn drop_continual_tasks(&mut self, cts: Vec<(CatalogItemId, ClusterId, GlobalId)>) {
-        let mut by_cluster: BTreeMap<_, Vec<_>> = BTreeMap::new();
-        let mut source_ids = Vec::new();
-        for (item_id, cluster_id, gid) in cts {
-            by_cluster.entry(cluster_id).or_default().push(gid);
-            source_ids.push((item_id, gid));
-        }
-
-        // Drop compute sinks.
-        for (cluster_id, ids) in by_cluster {
-            let compute = &mut self.controller.compute;
-            // A cluster could have been dropped, so verify it exists.
-            if compute.instance_exists(cluster_id) {
-                compute
-                    .drop_collections(cluster_id, ids)
-                    .unwrap_or_terminate("cannot fail to drop collections");
-            }
-        }
-
-        // Drop storage sources.
-        self.drop_sources(source_ids)
-    }
-
     fn drop_vpc_endpoints_in_background(&self, vpc_endpoints: Vec<CatalogItemId>) {
         let cloud_resource_controller = Arc::clone(self.cloud_resource_controller
             .as_ref()
@@ -1553,6 +1480,7 @@ impl Coordinator {
                 | Op::AlterRetainHistory { .. }
                 | Op::AlterNetworkPolicy { .. }
                 | Op::AlterAddColumn { .. }
+                | Op::AlterMaterializedViewApplyReplacement { .. }
                 | Op::UpdatePrivilege { .. }
                 | Op::UpdateDefaultPrivilege { .. }
                 | Op::GrantRole { .. }

src/adapter/src/coord/peek.rs

@@ -854,7 +854,7 @@ impl crate::coord::Coordinator {
         // If a dataflow was created, drop it once the peek command is sent.
         if let Some(index_id) = drop_dataflow {
             self.remove_compute_ids_from_timeline(vec![(compute_instance, index_id)]);
-            self.drop_indexes(vec![(compute_instance, index_id)]);
+            self.drop_compute_collections(vec![(compute_instance, index_id)]);
         }
 
         let persist_client = self.persist_client.clone();

src/adapter/src/coord/sequencer.rs

@@ -512,6 +512,12 @@ impl Coordinator {
                     let result = self.sequence_alter_table(&mut ctx, plan).await;
                     ctx.retire(result);
                 }
+                Plan::AlterMaterializedViewApplyReplacement(plan) => {
+                    let result = self
+                        .sequence_alter_materialized_view_apply_replacement(&mut ctx, plan)
+                        .await;
+                    ctx.retire(result);
+                }
                 Plan::AlterNetworkPolicy(plan) => {
                     let res = self
                         .sequence_alter_network_policy(ctx.session(), plan)

src/adapter/src/coord/sequencer/inner.rs

@@ -61,7 +61,10 @@ use mz_sql::names::{
     Aug, ObjectId, QualifiedItemName, ResolvedDatabaseSpecifier, ResolvedIds, ResolvedItemName,
     SchemaSpecifier, SystemObjectId,
 };
-use mz_sql::plan::{ConnectionDetails, NetworkPolicyRule, StatementContext};
+use mz_sql::plan::{
+    AlterMaterializedViewApplyReplacementPlan, ConnectionDetails, NetworkPolicyRule,
+    StatementContext,
+};
 use mz_sql::pure::{PurifiedSourceExport, generate_subsource_statements};
 use mz_storage_types::sinks::StorageSinkDesc;
 use mz_storage_types::sources::{GenericSourceConnection, IngestionDescription, SourceExport};
@@ -5204,6 +5207,57 @@ impl Coordinator {
 
         Ok(ExecuteResponse::AlteredObject(ObjectType::Table))
     }
+
+    #[instrument]
+    pub(super) async fn sequence_alter_materialized_view_apply_replacement(
+        &mut self,
+        ctx: &mut ExecuteContext,
+        plan: AlterMaterializedViewApplyReplacementPlan,
+    ) -> Result<ExecuteResponse, AdapterError> {
+        const ERROR_CONTEXT: &str = "ALTER MATERIALIZED VIEW ... APPLY REPLACEMENT";
+
+        let AlterMaterializedViewApplyReplacementPlan { id, replacement_id } = plan;
+
+        // TODO(alter-mv): Wait until there is overlap between the old MV's write frontier and the
+        // new MV's as-of, to ensure no times are skipped.
+
+        let Some(old) = self.catalog().get_entry(&id).materialized_view() else {
+            return Err(AdapterError::internal(
+                ERROR_CONTEXT,
+                "id must refer to a materialized view",
+            ));
+        };
+        let Some(new) = self
+            .catalog()
+            .get_entry(&replacement_id)
+            .materialized_view()
+        else {
+            return Err(AdapterError::internal(
+                ERROR_CONTEXT,
+                "replacement_id must refer to a materialized view",
+            ));
+        };
+
+        let old_cluster_id = old.cluster_id;
+        let new_cluster_id = new.cluster_id;
+
+        let old_gid = old.global_id_writes();
+        let new_gid = new.global_id_writes();
+
+        let ops = vec![catalog::Op::AlterMaterializedViewApplyReplacement { id, replacement_id }];
+
+        self.catalog_transact_with_side_effects(Some(ctx), ops, move |coord, _ctx| {
+            Box::pin(async move {
+                // Cut over the MV computation, by shutting down the old dataflow and allowing the
+                // new dataflow to start writing.
+                coord.drop_compute_collections(vec![(old_cluster_id, old_gid)]);
+                coord.allow_writes(new_cluster_id, new_gid);
+            })
+        })
+        .await?;
+
+        Ok(ExecuteResponse::AlteredObject(ObjectType::MaterializedView))
+    }
 }
 
 #[derive(Debug)]