mz_adapter/coord/sequencer/inner/

subscribe.rs

// Copyright Materialize, Inc. and contributors. All rights reserved.
//
// Use of this software is governed by the Business Source License
// included in the LICENSE file.
//
// As of the Change Date specified in that file, in accordance with
// the Business Source License, use of this software will be governed
// by the Apache License, Version 2.0.

use maplit::btreemap;
use mz_adapter_types::connection::ConnectionId;
use mz_cluster_client::ReplicaId;
use mz_compute_types::ComputeInstanceId;
use mz_compute_types::dataflows::DataflowDescription;
use mz_compute_types::plan::Plan;
use mz_ore::collections::CollectionExt;
use mz_ore::instrument;
use mz_repr::GlobalId;
use mz_repr::explain::{ExprHumanizerExt, TransientItem};
use mz_repr::optimize::{OptimizerFeatures, OverrideFrom};
use mz_sql::plan::{self, QueryWhen, SubscribeFrom};
use mz_sql::session::metadata::SessionMetadata;
use std::collections::BTreeSet;
use timely::progress::Antichain;
use tokio::sync::mpsc;
use tracing::Span;
use uuid::Uuid;

use crate::active_compute_sink::{ActiveComputeSink, ActiveSubscribe};
use crate::command::ExecuteResponse;
use crate::coord::sequencer::inner::return_if_err;
use crate::coord::sequencer::{check_log_reads, emit_optimizer_notices};
use crate::coord::{
    Coordinator, ExplainContext, ExplainPlanContext, Message, PlanValidity, StageResult, Staged,
    SubscribeExplain, SubscribeFinish, SubscribeOptimizeMir, SubscribeStage,
    SubscribeTimestampOptimizeLir, TargetCluster,
};
use crate::error::AdapterError;
use crate::explain::optimizer_trace::OptimizerTrace;
use crate::optimize::Optimize;
use crate::session::{Session, TransactionOps};
use crate::{
    AdapterNotice, ExecuteContext, ExecuteContextGuard, ReadHolds, TimelineContext, optimize,
};

impl Staged for SubscribeStage {
    type Ctx = ExecuteContext;

    fn validity(&mut self) -> &mut PlanValidity {
        match self {
            SubscribeStage::OptimizeMir(stage) => &mut stage.validity,
            SubscribeStage::TimestampOptimizeLir(stage) => &mut stage.validity,
            SubscribeStage::Finish(stage) => &mut stage.validity,
            SubscribeStage::Explain(stage) => &mut stage.validity,
        }
    }

    async fn stage(
        self,
        coord: &mut Coordinator,
        ctx: &mut ExecuteContext,
    ) -> Result<StageResult<Box<Self>>, AdapterError> {
        match self {
            SubscribeStage::OptimizeMir(stage) => coord.subscribe_optimize_mir(stage),
            SubscribeStage::TimestampOptimizeLir(stage) => {
                coord.subscribe_timestamp_optimize_lir(ctx, stage).await
            }
            SubscribeStage::Finish(stage) => coord.subscribe_finish(ctx, stage).await,
            SubscribeStage::Explain(stage) => coord.subscribe_explain(ctx.session(), stage).await,
        }
    }

    fn message(self, ctx: ExecuteContext, span: Span) -> Message {
        Message::SubscribeStageReady {
            ctx,
            span,
            stage: self,
        }
    }

    fn cancel_enabled(&self) -> bool {
        true
    }
}

impl Coordinator {
    #[instrument]
    pub(crate) async fn sequence_subscribe(
        &mut self,
        mut ctx: ExecuteContext,
        plan: plan::SubscribePlan,
        target_cluster: TargetCluster,
    ) {
        let stage = return_if_err!(
            self.subscribe_validate(
                ctx.session_mut(),
                plan,
                target_cluster,
                ExplainContext::None
            ),
            ctx
        );
        self.sequence_staged(ctx, Span::current(), stage).await;
    }

    #[instrument]
    pub(crate) async fn explain_subscribe(
        &mut self,
        mut ctx: ExecuteContext,
        plan::ExplainPlanPlan {
            stage,
            format,
            config,
            explainee,
        }: plan::ExplainPlanPlan,
        target_cluster: TargetCluster,
    ) {
        let plan::Explainee::Statement(stmt) = explainee else {
            // This is currently asserted in the `sequence_explain_plan` code that
            // calls this method.
            unreachable!()
        };
        let plan::ExplaineeStatement::Subscribe { broken, plan } = stmt else {
            // This is currently asserted in the `sequence_explain_plan` code that
            // calls this method.
            unreachable!()
        };

        let desc = match &plan.from {
            SubscribeFrom::Id(_) => None,
            SubscribeFrom::Query { desc, .. } => Some(desc.clone()),
        };

        // Create an OptimizerTrace instance to collect plans emitted when
        // executing the optimizer pipeline.
        let optimizer_trace = OptimizerTrace::new(stage.paths());

        let explain_ctx = ExplainContext::Plan(ExplainPlanContext {
            broken,
            config,
            format,
            stage,
            replan: None,
            desc,
            optimizer_trace,
        });
        let stage = return_if_err!(
            self.subscribe_validate(ctx.session_mut(), plan, target_cluster, explain_ctx),
            ctx
        );
        self.sequence_staged(ctx, Span::current(), stage).await;
    }

    #[instrument]
    fn subscribe_validate(
        &self,
        session: &mut Session,
        plan: plan::SubscribePlan,
        target_cluster: TargetCluster,
        explain_ctx: ExplainContext,
    ) -> Result<SubscribeStage, AdapterError> {
        let plan::SubscribePlan { from, when, .. } = &plan;

        let cluster = self
            .catalog()
            .resolve_target_cluster(target_cluster, session)?;
        let cluster_id = cluster.id;

        // Only check cluster replicas if we're not in explain mode.
        if explain_ctx.needs_cluster() && cluster.replicas().next().is_none() {
            return Err(AdapterError::NoClusterReplicasAvailable {
                name: cluster.name.clone(),
                is_managed: cluster.is_managed(),
            });
        }

        let mut replica_id = session
            .vars()
            .cluster_replica()
            .map(|name| {
                cluster
                    .replica_id(name)
                    .ok_or(AdapterError::UnknownClusterReplica {
                        cluster_name: cluster.name.clone(),
                        replica_name: name.to_string(),
                    })
            })
            .transpose()?;

        // SUBSCRIBE AS OF, similar to peeks, doesn't need to worry about transaction
        // timestamp semantics.
        if when == &QueryWhen::Immediately {
            // If this isn't a SUBSCRIBE AS OF, the SUBSCRIBE can be in a transaction if it's the
            // only operation.
            session.add_transaction_ops(TransactionOps::Subscribe)?;
        }

        let depends_on = from.depends_on();

        // Run `check_log_reads` and emit notices.
        let notices = check_log_reads(
            self.catalog(),
            cluster,
            &depends_on,
            &mut replica_id,
            session.vars(),
        )?;
        session.add_notices(notices);

        // Determine timeline.
        let mut timeline = self
            .catalog()
            .validate_timeline_context(depends_on.iter().copied())?;
        if matches!(timeline, TimelineContext::TimestampIndependent) && from.contains_temporal() {
            // If the from IDs are timestamp independent but the query contains temporal functions
            // then the timeline context needs to be upgraded to timestamp dependent.
            timeline = TimelineContext::TimestampDependent;
        }

        let dependencies = depends_on
            .iter()
            .map(|id| self.catalog().resolve_item_id(id))
            .collect();
        let validity = PlanValidity::new(
            self.catalog().transient_revision(),
            dependencies,
            Some(cluster_id),
            replica_id,
            session.role_metadata().clone(),
        );

        Ok(SubscribeStage::OptimizeMir(SubscribeOptimizeMir {
            validity,
            plan,
            timeline,
            dependency_ids: depends_on,
            cluster_id,
            replica_id,
            explain_ctx,
        }))
    }

    #[instrument]
    fn subscribe_optimize_mir(
        &self,
        SubscribeOptimizeMir {
            mut validity,
            plan,
            timeline,
            dependency_ids,
            cluster_id,
            replica_id,
            explain_ctx,
        }: SubscribeOptimizeMir,
    ) -> Result<StageResult<Box<SubscribeStage>>, AdapterError> {
        let plan::SubscribePlan {
            with_snapshot,
            up_to,
            ..
        } = &plan;

        // Collect optimizer parameters.
        let compute_instance = self
            .instance_snapshot(cluster_id)
            .expect("compute instance does not exist");
        let (_, view_id) = self.allocate_transient_id();
        let (_, sink_id) = self.allocate_transient_id();
        let debug_name = format!("subscribe-{}", sink_id);
        let optimizer_config = optimize::OptimizerConfig::from(self.catalog().system_config())
            .override_from(&self.catalog.get_cluster(cluster_id).config.features())
            .override_from(&explain_ctx);

        // Build an optimizer for this SUBSCRIBE.
        let mut optimizer = optimize::subscribe::Optimizer::new(
            self.owned_catalog(),
            compute_instance,
            view_id,
            sink_id,
            *with_snapshot,
            *up_to,
            debug_name,
            optimizer_config,
            self.optimizer_metrics(),
        );
        let catalog = self.owned_catalog();

        let span = Span::current();
        Ok(StageResult::Handle(mz_ore::task::spawn_blocking(
            || "optimize subscribe (mir)",
            move || {
                span.in_scope(|| {
                    let _dispatch_guard = explain_ctx.dispatch_guard();

                    // MIR ⇒ MIR optimization (global)
                    let global_mir_plan = optimizer.catch_unwind_optimize(plan.clone())?;
                    // Add introduced indexes as validity dependencies.
                    validity.extend_dependencies(
                        global_mir_plan
                            .id_bundle(optimizer.cluster_id())
                            .iter()
                            .map(|id| catalog.resolve_item_id(&id)),
                    );

                    let stage =
                        SubscribeStage::TimestampOptimizeLir(SubscribeTimestampOptimizeLir {
                            validity,
                            plan,
                            timeline,
                            optimizer,
                            global_mir_plan,
                            dependency_ids,
                            replica_id,
                            explain_ctx,
                        });
                    Ok(Box::new(stage))
                })
            },
        )))
    }

    #[instrument]
    async fn subscribe_timestamp_optimize_lir(
        &mut self,
        ctx: &ExecuteContext,
        SubscribeTimestampOptimizeLir {
            validity,
            plan,
            timeline,
            mut optimizer,
            global_mir_plan,
            dependency_ids,
            replica_id,
            explain_ctx,
        }: SubscribeTimestampOptimizeLir,
    ) -> Result<StageResult<Box<SubscribeStage>>, AdapterError> {
        let plan::SubscribePlan { when, .. } = &plan;

        // Timestamp selection
        let oracle_read_ts = self.oracle_read_ts(ctx.session(), &timeline, when).await;
        let bundle = &global_mir_plan.id_bundle(optimizer.cluster_id());
        let (determination, read_holds) = self.determine_timestamp(
            ctx.session(),
            bundle,
            when,
            optimizer.cluster_id(),
            &timeline,
            oracle_read_ts,
            None,
        )?;

        let as_of = determination.timestamp_context.timestamp_or_default();

        if let Some(id) = ctx.extra().contents() {
            self.set_statement_execution_timestamp(id, as_of);
        }
        if let Some(up_to) = optimizer.up_to() {
            if as_of == up_to {
                ctx.session()
                    .add_notice(AdapterNotice::EqualSubscribeBounds { bound: up_to });
            } else if as_of > up_to {
                return Err(AdapterError::AbsurdSubscribeBounds { as_of, up_to });
            }
        }

        self.store_transaction_read_holds(ctx.session().conn_id().clone(), read_holds);

        let global_mir_plan = global_mir_plan.resolve(Antichain::from_elem(as_of));

        // Optimize LIR
        let span = Span::current();
        Ok(StageResult::Handle(mz_ore::task::spawn_blocking(
            || "optimize subscribe (lir)",
            move || {
                span.in_scope(|| {
                    let _dispatch_guard = explain_ctx.dispatch_guard();

                    let cluster_id = optimizer.cluster_id();

                    let mut pipeline = || -> Result<_, AdapterError> {
                        // MIR ⇒ LIR lowering and LIR ⇒ LIR optimization (global)
                        let global_lir_plan =
                            optimizer.catch_unwind_optimize(global_mir_plan.clone())?;
                        Ok(global_lir_plan)
                    };

                    let stage = match pipeline() {
                        Ok(global_lir_plan) => {
                            if let ExplainContext::Plan(explain_ctx) = explain_ctx {
                                let (_, df_meta) = global_lir_plan.unapply();
                                SubscribeStage::Explain(SubscribeExplain {
                                    validity,
                                    optimizer,
                                    df_meta,
                                    cluster_id,
                                    explain_ctx,
                                })
                            } else {
                                SubscribeStage::Finish(SubscribeFinish {
                                    validity,
                                    cluster_id,
                                    plan,
                                    global_lir_plan,
                                    dependency_ids,
                                    replica_id,
                                })
                            }
                        }
                        Err(err) => {
                            let ExplainContext::Plan(explain_ctx) = explain_ctx else {
                                return Err(err);
                            };

                            if explain_ctx.broken {
                                tracing::error!("error while handling EXPLAIN statement: {}", err);
                                SubscribeStage::Explain(SubscribeExplain {
                                    validity,
                                    optimizer,
                                    df_meta: Default::default(),
                                    cluster_id,
                                    explain_ctx,
                                })
                            } else {
                                return Err(err);
                            }
                        }
                    };

                    Ok(Box::new(stage))
                })
            },
        )))
    }

    #[instrument]
    async fn subscribe_finish(
        &mut self,
        ctx: &mut ExecuteContext,
        SubscribeFinish {
            validity: _,
            cluster_id,
            plan,
            global_lir_plan,
            dependency_ids,
            replica_id,
        }: SubscribeFinish,
    ) -> Result<StageResult<Box<SubscribeStage>>, AdapterError> {
        let (df_desc, df_meta) = global_lir_plan.unapply();
        emit_optimizer_notices(&*self.catalog, ctx.session(), &df_meta.optimizer_notices);
        let conn_id = ctx.session.conn_id().clone();
        let session_uuid = ctx.session().uuid();
        let txn_read_holds = self
            .txn_read_holds
            .remove(&conn_id)
            .expect("must have previously installed read holds");
        let resp = self
            .implement_subscribe(
                ctx.extra_mut(),
                df_desc,
                dependency_ids,
                cluster_id,
                replica_id,
                conn_id,
                session_uuid,
                txn_read_holds,
                plan,
            )
            .await?;
        Ok(StageResult::Response(resp))
    }

    #[instrument]
    pub(crate) async fn implement_subscribe(
        &mut self,
        ctx_extra: &mut ExecuteContextGuard,
        df_desc: DataflowDescription<Plan>,
        dependency_ids: BTreeSet<GlobalId>,
        cluster_id: ComputeInstanceId,
        replica_id: Option<ReplicaId>,
        conn_id: ConnectionId,
        session_uuid: Uuid,
        read_holds: ReadHolds<mz_repr::Timestamp>,
        plan: plan::SubscribePlan,
    ) -> Result<ExecuteResponse, AdapterError> {
        let sink_id = df_desc.sink_id();

        let (tx, rx) = mpsc::unbounded_channel();
        let active_subscribe = ActiveSubscribe {
            conn_id: conn_id.clone(),
            session_uuid,
            channel: tx,
            emit_progress: plan.emit_progress,
            as_of: df_desc
                .as_of
                .as_ref()
                .and_then(|t| t.as_option())
                .copied()
                .expect("set to Some in an earlier stage"),
            arity: df_desc
                .sink_exports
                .values()
                .into_element()
                .from_desc
                .arity(),
            cluster_id,
            depends_on: dependency_ids,
            start_time: self.now(),
            output: plan.output,
        };
        active_subscribe.initialize();

        // Add metadata for the new SUBSCRIBE.
        let write_notify_fut = self
            .add_active_compute_sink(sink_id, ActiveComputeSink::Subscribe(active_subscribe))
            .await;
        // Ship dataflow.
        let ship_dataflow_fut = self.ship_dataflow(df_desc, cluster_id, replica_id);

        // Both adding metadata for the new SUBSCRIBE and shipping the underlying dataflow, send
        // requests to external services, which can take time, so we run them concurrently.
        let ((), ()) = futures::future::join(write_notify_fut, ship_dataflow_fut).await;

        // Explicitly drop read holds, just to make it obvious what's happening.
        drop(read_holds);

        let resp = ExecuteResponse::Subscribing {
            rx,
            ctx_extra: std::mem::take(ctx_extra),
            instance_id: cluster_id,
        };
        let resp = match plan.copy_to {
            None => resp,
            Some(format) => ExecuteResponse::CopyTo {
                format,
                resp: Box::new(resp),
            },
        };
        Ok(resp)
    }

    #[instrument]
    async fn subscribe_explain(
        &self,
        session: &Session,
        SubscribeExplain {
            optimizer,
            df_meta,
            cluster_id,
            explain_ctx:
                ExplainPlanContext {
                    config,
                    format,
                    stage,
                    optimizer_trace,
                    desc,
                    ..
                },
            ..
        }: SubscribeExplain,
    ) -> Result<StageResult<Box<SubscribeStage>>, AdapterError> {
        let session_catalog = self.catalog().for_session(session);

        let expr_humanizer = {
            let transient_items = btreemap! {
                optimizer.sink_id() => TransientItem::new(
                    Some(vec![GlobalId::Explain.to_string()]),
                    desc.map(|d| d.iter_names().map(|c| c.to_string()).collect()),
                )
            };
            ExprHumanizerExt::new(transient_items, &session_catalog)
        };

        let target_cluster = self.catalog().get_cluster(cluster_id);

        let features = OptimizerFeatures::from(self.catalog().system_config())
            .override_from(&target_cluster.config.features())
            .override_from(&config.features);

        let rows = optimizer_trace
            .into_rows(
                format,
                &config,
                &features,
                &expr_humanizer,
                None,
                Some(target_cluster),
                df_meta,
                stage,
                plan::ExplaineeStatementKind::Subscribe,
                None,
            )
            .await?;

        Ok(StageResult::Response(Self::send_immediate_rows(rows)))
    }
}