Skip to main content

mz_adapter/
frontend_peek.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
10use std::collections::BTreeMap;
11use std::collections::BTreeSet;
12use std::sync::Arc;
13use std::time::Duration;
14
15use futures::StreamExt;
16use itertools::Itertools;
17use mz_adapter_types::dyncfgs::ENABLE_FRONTEND_SUBSCRIBES;
18use mz_compute_types::ComputeInstanceId;
19use mz_controller_types::{ClusterId, ReplicaId};
20use mz_expr::{CollectionPlan, ResultSpec, RowSetFinishing};
21use mz_ore::cast::{CastFrom, CastLossy};
22use mz_ore::collections::CollectionExt;
23use mz_ore::now::EpochMillis;
24use mz_ore::task::JoinHandle;
25use mz_ore::{soft_assert_eq_or_log, soft_assert_or_log, soft_panic_or_log};
26use mz_repr::optimize::{OptimizerFeatures, OverrideFrom};
27use mz_repr::{Datum, GlobalId, IntoRowIterator, Row, RowIterator, Timestamp};
28use mz_sql::ast::Raw;
29use mz_sql::catalog::CatalogCluster;
30use mz_sql::plan::Params;
31use mz_sql::plan::{
32    self, Explainee, ExplaineeStatement, Plan, QueryWhen, SelectPlan, SideEffectingFunc,
33    SubscribePlan,
34};
35use mz_sql::rbac;
36use mz_sql::session::metadata::SessionMetadata;
37use mz_sql::session::user::MZ_SYSTEM_ROLE_ID;
38use mz_sql::session::vars::{
39    CLUSTER, CLUSTER_REPLICA, IsolationLevel, TRANSACTION_ISOLATION, Var, VarInput,
40};
41use mz_sql_parser::ast::{CopyDirection, ExplainStage, ShowStatement, Statement};
42use mz_transform::EmptyStatisticsOracle;
43use mz_transform::dataflow::DataflowMetainfo;
44use opentelemetry::trace::TraceContextExt;
45use timely::progress::Antichain;
46use tracing::{Span, debug, warn};
47use tracing_opentelemetry::OpenTelemetrySpanExt;
48
49use crate::catalog::Catalog;
50use crate::command::{CatalogSnapshot, Command};
51use crate::coord;
52use crate::coord::peek::{FastPathPlan, PeekPlan, PeekResponseUnary};
53use crate::coord::sequencer::{eval_copy_to_uri, statistics_oracle};
54use crate::coord::timeline::timedomain_for;
55use crate::coord::timestamp_selection::TimestampDetermination;
56use crate::coord::{Coordinator, CopyToContext, ExplainContext, ExplainPlanContext, TargetCluster};
57use crate::explain::insights::PlanInsightsContext;
58use crate::explain::optimizer_trace::OptimizerTrace;
59use crate::optimize::Optimize;
60use crate::optimize::dataflows::{ComputeInstanceSnapshot, DataflowBuilder};
61use crate::peek_client::{ExecutionLogging, TakeOver};
62use crate::session::{Session, TransactionOps, TransactionStatus};
63use crate::statement_logging::StatementLifecycleEvent;
64use crate::statement_logging::WatchSetCreation;
65use crate::{
66    AdapterError, AdapterNotice, CollectionIdBundle, ExecuteResponse, PeekClient, ReadHolds,
67    TimelineContext, TimestampContext, TimestampProvider, optimize,
68};
69
70impl PeekClient {
71    /// Attempt to sequence a peek from the session task.
72    ///
73    /// Returns `Ok(Some(response))` if we handled the peek, or `Ok(None)` to fall back to the
74    /// Coordinator's sequencing. If it returns an error, it should be returned to the user.
75    ///
76    /// `logging` holds the end-of-execution obligation for this statement. The
77    /// caller retires it, this function only takes the statement over and, at
78    /// the dispatch sites that hand execution off, defuses the slot.
79    pub(crate) async fn try_frontend_peek(
80        &mut self,
81        portal_name: &str,
82        session: &mut Session,
83        logging: &mut ExecutionLogging,
84    ) -> Result<Option<ExecuteResponse>, AdapterError> {
85        // # From handle_execute
86
87        if session.vars().emit_trace_id_notice() {
88            let span_context = tracing::Span::current()
89                .context()
90                .span()
91                .span_context()
92                .clone();
93            if span_context.is_valid() {
94                session.add_notice(AdapterNotice::QueryTrace {
95                    trace_id: span_context.trace_id(),
96                });
97            }
98        }
99
100        let catalog = self.catalog_snapshot("try_frontend_peek").await;
101
102        // Extract things from the portal. A failed verification does not begin
103        // an entry, mirroring the coordinator: the portal is what statement
104        // logging draws its record from. An adopted outer entry still gets its
105        // `Errored` end from the retirement site in `SessionClient::execute`.
106        let (stmt, params, logging_info, lifecycle_timestamps) = {
107            Coordinator::verify_portal(&*catalog, session, portal_name)?;
108            let portal = session
109                .get_portal_unverified(portal_name)
110                // The portal is a session-level thing, so it couldn't have concurrently disappeared
111                // since the above verification.
112                .expect("called verify_portal above");
113            let params = portal.parameters.clone();
114            let stmt = portal.stmt.clone();
115            let logging_info = Arc::clone(&portal.logging);
116            let lifecycle_timestamps = portal.lifecycle_timestamps.clone();
117            (stmt, params, logging_info, lifecycle_timestamps)
118        };
119
120        // Before planning, check if this is a statement type we can handle.
121        // This must happen BEFORE statement logging setup to avoid orphaned execution records.
122        if let Some(ref stmt) = stmt {
123            match &**stmt {
124                Statement::Select(_)
125                | Statement::ExplainAnalyzeObject(_)
126                | Statement::ExplainAnalyzeCluster(_)
127                | Statement::Show(ShowStatement::ShowObjects(_))
128                | Statement::Show(ShowStatement::ShowColumns(_)) => {
129                    // These are always fine, just continue.
130                    // Note: EXPLAIN ANALYZE will `plan` to `Plan::Select`.
131                    // Note: ShowObjects plans to `Plan::Select`, ShowColumns plans to `Plan::ShowColumns`.
132                    // We handle `Plan::ShowColumns` specially in `try_frontend_peek_inner`.
133                }
134                Statement::ExplainPlan(explain_stmt) => {
135                    // Only handle ExplainPlan for SELECT statements.
136                    // We don't want to handle e.g. EXPLAIN CREATE MATERIALIZED VIEW here, because that
137                    // requires purification before planning, which the frontend peek sequencing doesn't
138                    // do.
139                    match &explain_stmt.explainee {
140                        mz_sql_parser::ast::Explainee::Select(..) => {
141                            // This is a SELECT, continue
142                        }
143                        _ => {
144                            debug!(
145                                "Bailing out from try_frontend_peek, because EXPLAIN is not for a SELECT query"
146                            );
147                            return Ok(None);
148                        }
149                    }
150                }
151                Statement::ExplainPushdown(explain_stmt) => {
152                    // Only handle EXPLAIN FILTER PUSHDOWN for non-BROKEN SELECT statements
153                    match &explain_stmt.explainee {
154                        mz_sql_parser::ast::Explainee::Select(_, false) => {}
155                        _ => {
156                            debug!(
157                                "Bailing out from try_frontend_peek, because EXPLAIN FILTER PUSHDOWN is not for a SELECT query or is for EXPLAIN BROKEN"
158                            );
159                            return Ok(None);
160                        }
161                    }
162                }
163                Statement::Copy(copy_stmt) => {
164                    match &copy_stmt.direction {
165                        CopyDirection::To => {
166                            // This is COPY TO (...), continue
167                        }
168                        CopyDirection::From => {
169                            debug!(
170                                "Bailing out from try_frontend_peek, because COPY FROM is not supported"
171                            );
172                            return Ok(None);
173                        }
174                    }
175                }
176
177                Statement::Subscribe(_)
178                    if ENABLE_FRONTEND_SUBSCRIBES.get(catalog.system_config().dyncfgs()) =>
179                {
180                    // We have a subscribe statement to process; continue.
181                }
182                _ => {
183                    debug!(
184                        "Bailing out from try_frontend_peek, because statement type is not supported"
185                    );
186                    return Ok(None);
187                }
188            }
189        }
190
191        // Past this point the coordinator never sees this statement, so every
192        // exit has to produce an outcome for it. The `Ok(None)` bailouts are all
193        // above.
194        logging.take_over(
195            self,
196            session,
197            stmt.as_deref(),
198            &params,
199            &logging_info,
200            &catalog,
201            lifecycle_timestamps,
202            TakeOver::StatementToRun,
203        );
204
205        self.try_frontend_peek_inner(session, catalog, stmt, params, logging)
206            .await
207    }
208
209    /// Executes a coordinator-owned `SELECT`, pinned to `replica_id`, and
210    /// returns its rows.
211    ///
212    /// The caller must build `sql` itself rather than accept it from a user.
213    /// It runs as the system role, in a session of its own, and is not
214    /// statement-logged. This calls into the coordinator, so awaiting it on the
215    /// coordinator's main loop deadlocks.
216    ///
217    /// NOTE: Dropping the returned future does not cancel the peek, since
218    /// there is no connection to cancel it with. A peek dropped after its issue
219    /// keeps its read holds until the replica answers or goes away. One dropped
220    /// between its registration with the coordinator and its issue leaves the
221    /// registration behind, since only a response, a cancel of its connection,
222    /// or a drop of its inputs retires it.
223    pub(crate) async fn background_peek(
224        &mut self,
225        sql: &str,
226        cluster_id: ClusterId,
227        replica_id: ReplicaId,
228    ) -> Result<Vec<Row>, AdapterError> {
229        // Not `catalog_snapshot`, which expects the coordinator to outlive its
230        // caller. A background client may race the coordinator's shutdown.
231        let CatalogSnapshot { catalog } = self
232            .call_coordinator(|tx| Command::CatalogSnapshot { tx })
233            .await?;
234        let (cluster_name, replica_name) = {
235            let cluster = catalog
236                .try_get_cluster(cluster_id)
237                .ok_or_else(|| AdapterError::Internal(format!("unknown cluster {cluster_id}")))?;
238            let replica = cluster
239                .replica(replica_id)
240                .ok_or_else(|| AdapterError::Internal(format!("unknown replica {replica_id}")))?;
241            (cluster.name.clone(), replica.name.clone())
242        };
243        let stmt = mz_sql::parse::parse(sql)?.into_element().ast;
244
245        let mut session = Session::dummy();
246        session.initialize_role_metadata(MZ_SYSTEM_ROLE_ID);
247        let system_vars = catalog.system_config();
248        // Serializable reads at the freshest time the inputs can serve. A
249        // strict-serializable read picks the stalest valid time at or past the
250        // oracle's read timestamp instead. In a read-only environment whose
251        // leader is not running, nothing advances the oracle, and such a read
252        // sees introspection logs as of a time before the replicas started.
253        for (var, value) in [
254            (&CLUSTER, cluster_name.as_str()),
255            (&CLUSTER_REPLICA, replica_name.as_str()),
256            (&TRANSACTION_ISOLATION, "serializable"),
257        ] {
258            session
259                .vars_mut()
260                .set(system_vars, var.name(), VarInput::Flat(value), false)?;
261        }
262        session.start_transaction_single_stmt(mz_ore::now::to_datetime((catalog.config().now)()));
263
264        let mut logging = ExecutionLogging::adopt(None, self);
265        let response = self
266            .try_frontend_peek_inner(
267                &mut session,
268                catalog,
269                Some(Arc::new(stmt)),
270                Params::empty(),
271                &mut logging,
272            )
273            .await?;
274
275        let mut rows = Vec::new();
276        let mut collect = |mut iter: Box<dyn RowIterator + Send + Sync>| {
277            while let Some(row) = iter.next() {
278                rows.push(row.to_owned());
279            }
280        };
281        match response {
282            Some(ExecuteResponse::SendingRowsImmediate { rows }) => collect(rows),
283            Some(ExecuteResponse::SendingRowsStreaming { mut rows, .. }) => {
284                while let Some(response) = rows.next().await {
285                    match response {
286                        PeekResponseUnary::Rows(iter) => collect(iter),
287                        PeekResponseUnary::Error(error) => return Err(error),
288                        PeekResponseUnary::Canceled => return Err(AdapterError::Canceled),
289                        PeekResponseUnary::DependencyDropped(dependency) => {
290                            return Err(dependency.to_concurrent_dependency_drop());
291                        }
292                    }
293                }
294            }
295            response => {
296                return Err(AdapterError::Internal(format!(
297                    "background peek returned an unexpected response: {response:?}"
298                )));
299            }
300        }
301        Ok(rows)
302    }
303
304    /// This is encapsulated in an inner function so that the outer function can still do statement
305    /// logging after the `?` returns of the inner function.
306    ///
307    /// Dispatch sites that hand the statement to the coordinator for
308    /// asynchronous completion (registered peeks, subscribes) `defuse` the
309    /// logging slot at the point where the coordinator takes over. Everywhere
310    /// else the slot stays armed and `SessionClient::execute` logs the end from
311    /// the returned result.
312    #[mz_ore::instrument(level = "debug")]
313    async fn try_frontend_peek_inner(
314        &mut self,
315        session: &mut Session,
316        catalog: Arc<Catalog>,
317        stmt: Option<Arc<Statement<Raw>>>,
318        params: Params,
319        logging: &mut ExecutionLogging,
320    ) -> Result<Option<ExecuteResponse>, AdapterError> {
321        let stmt = match stmt {
322            Some(stmt) => stmt,
323            None => {
324                debug!("try_frontend_peek_inner succeeded on an empty query");
325                return Ok(Some(ExecuteResponse::EmptyQuery));
326            }
327        };
328
329        // # From handle_execute_inner
330
331        let conn_catalog = catalog.for_session(session);
332        // (`resolved_ids` should be derivable from `stmt`. If `stmt` is later transformed to
333        // remove/add IDs, then `resolved_ids` should be updated to also remove/add those IDs.)
334        let (stmt, resolved_ids) = mz_sql::names::resolve(&conn_catalog, (*stmt).clone())?;
335
336        let pcx = session.pcx();
337        let (plan, sql_impl_ids) =
338            mz_sql::plan::plan(Some(pcx), &conn_catalog, stmt, &params, &resolved_ids)?;
339
340        /// What do we do with the result of the select?
341        enum QueryPlan<'a> {
342            Select(&'a SelectPlan),
343            CopyTo(&'a SelectPlan, CopyToContext),
344            Subscribe(&'a SubscribePlan),
345        }
346
347        let (query_plan, explain_ctx) = match &plan {
348            Plan::Select(select_plan) => {
349                let explain_ctx = if session.vars().emit_plan_insights_notice() {
350                    let optimizer_trace = OptimizerTrace::new(ExplainStage::PlanInsights.paths());
351                    ExplainContext::PlanInsightsNotice(optimizer_trace)
352                } else {
353                    ExplainContext::None
354                };
355                (QueryPlan::Select(select_plan), explain_ctx)
356            }
357            Plan::ShowColumns(show_columns_plan) => {
358                // ShowColumns wraps a SelectPlan, extract it and proceed as normal.
359                (
360                    QueryPlan::Select(&show_columns_plan.select_plan),
361                    ExplainContext::None,
362                )
363            }
364            Plan::ExplainPlan(plan::ExplainPlanPlan {
365                stage,
366                format,
367                config,
368                explainee: Explainee::Statement(ExplaineeStatement::Select { broken, plan, desc }),
369            }) => {
370                // Create OptimizerTrace to collect optimizer plans
371                let optimizer_trace = OptimizerTrace::new(stage.paths());
372                let explain_ctx = ExplainContext::Plan(ExplainPlanContext {
373                    broken: *broken,
374                    config: config.clone(),
375                    format: *format,
376                    stage: *stage,
377                    replan: None,
378                    desc: Some(desc.clone()),
379                    optimizer_trace,
380                });
381                (QueryPlan::Select(plan), explain_ctx)
382            }
383            // COPY TO S3
384            Plan::CopyTo(plan::CopyToPlan {
385                select_plan,
386                desc,
387                to,
388                connection,
389                connection_id,
390                format,
391                max_file_size,
392            }) => {
393                let uri = eval_copy_to_uri(to.clone(), session, catalog.state())?;
394
395                // (output_batch_count will be set later)
396                let copy_to_ctx = CopyToContext {
397                    desc: desc.clone(),
398                    uri,
399                    connection: connection.clone(),
400                    connection_id: *connection_id,
401                    format: format.clone(),
402                    max_file_size: *max_file_size,
403                    output_batch_count: None,
404                };
405
406                (
407                    QueryPlan::CopyTo(select_plan, copy_to_ctx),
408                    ExplainContext::None,
409                )
410            }
411            Plan::ExplainPushdown(plan::ExplainPushdownPlan { explainee }) => {
412                // Only handle EXPLAIN FILTER PUSHDOWN for SELECT statements
413                match explainee {
414                    plan::Explainee::Statement(plan::ExplaineeStatement::Select {
415                        broken: false,
416                        plan,
417                        desc: _,
418                    }) => {
419                        let explain_ctx = ExplainContext::Pushdown;
420                        (QueryPlan::Select(plan), explain_ctx)
421                    }
422                    _ => {
423                        // This shouldn't happen because we already checked for this at the AST
424                        // level before calling `try_frontend_peek_inner`. The logging takeover
425                        // has already happened, so falling back to the coordinator would count
426                        // the statement twice. Report the inconsistency instead.
427                        soft_panic_or_log!(
428                            "unexpected EXPLAIN FILTER PUSHDOWN plan kind in frontend peek sequencing: {:?}",
429                            explainee
430                        );
431                        return Err(AdapterError::Internal(
432                            "unexpected EXPLAIN FILTER PUSHDOWN plan kind in frontend peek sequencing"
433                                .into(),
434                        ));
435                    }
436                }
437            }
438            Plan::SideEffectingFunc(sef_plan) => {
439                // Look up the target connection's authenticated role, so that
440                // check_plan can perform RBAC for side-effecting functions.
441                //
442                // The RBAC check reflects the state at this point in time. A
443                // concurrent change to the issuer's role membership does not
444                // affect the already in-flight execution, similarly to how
445                // privilege changes don't affect other kinds of in-flight
446                // statements.
447                let target_conn = match sef_plan {
448                    SideEffectingFunc::PgCancelBackend {
449                        connection_id: Some(connection_id),
450                    } => {
451                        self.call_coordinator(|tx| Command::LookupConnection {
452                            connection_id: *connection_id,
453                            tx,
454                        })
455                        .await?
456                    }
457                    SideEffectingFunc::PgCancelBackend {
458                        connection_id: None,
459                    } => None,
460                };
461                let target_conn_role = target_conn.as_ref().map(|(_, role)| *role);
462
463                rbac::check_plan(
464                    &conn_catalog,
465                    target_conn_role,
466                    session,
467                    &plan,
468                    None,
469                    &resolved_ids,
470                    &sql_impl_ids,
471                )?;
472
473                // RBAC passed. Delegate execution to the Coordinator.
474                let response = self
475                    .call_coordinator(|tx| Command::ExecuteSideEffectingFunc {
476                        plan: sef_plan.clone(),
477                        conn_id: session.conn_id().clone(),
478                        tx,
479                    })
480                    .await??;
481
482                // We held the target's `ConnectionId` handle from the RBAC
483                // check until the Coordinator executed the function, which
484                // prevented the raw connection ID from being reused by a new
485                // connection. So the connection the Coordinator acted on (if
486                // it found one) is the one whose role we checked above.
487                drop(target_conn);
488
489                return Ok(Some(response));
490            }
491            Plan::Subscribe(subscribe) => (QueryPlan::Subscribe(subscribe), ExplainContext::None),
492            _ => {
493                // This shouldn't happen because we already checked for this at the AST
494                // level before calling `try_frontend_peek_inner`. The logging takeover has
495                // already happened, so falling back to the coordinator would count the
496                // statement twice. Report the inconsistency instead.
497                soft_panic_or_log!(
498                    "Unexpected plan kind in frontend peek sequencing: {:?}",
499                    plan
500                );
501                return Err(AdapterError::Internal(
502                    "unexpected plan kind in frontend peek sequencing".into(),
503                ));
504            }
505        };
506
507        let when = match query_plan {
508            QueryPlan::Select(s) => &s.when,
509            QueryPlan::CopyTo(s, _) => &s.when,
510            QueryPlan::Subscribe(s) => &s.when,
511        };
512
513        let depends_on = match query_plan {
514            QueryPlan::Select(s) => s.source.depends_on(),
515            QueryPlan::CopyTo(s, _) => s.source.depends_on(),
516            QueryPlan::Subscribe(s) => s.from.depends_on(),
517        };
518
519        let contains_temporal = match query_plan {
520            QueryPlan::Select(s) => s.source.contains_temporal(),
521            QueryPlan::CopyTo(s, _) => s.source.contains_temporal(),
522            QueryPlan::Subscribe(s) => s.from.contains_temporal(),
523        };
524
525        // # From sequence_plan
526
527        // We have checked the plan kind above.
528        assert!(plan.allowed_in_read_only());
529
530        let (cluster, target_cluster_id, target_cluster_name) = {
531            let target_cluster = match session.transaction().cluster() {
532                // Use the current transaction's cluster.
533                Some(cluster_id) => TargetCluster::Transaction(cluster_id),
534                // If there isn't a current cluster set for a transaction, then try to auto route.
535                None => coord::catalog_serving::auto_run_on_catalog_server(
536                    &conn_catalog,
537                    session,
538                    &plan,
539                ),
540            };
541            let cluster = catalog.resolve_target_cluster(target_cluster, session)?;
542            (cluster, cluster.id, &cluster.name)
543        };
544
545        // Log cluster selection
546        if let Some(logging_id) = logging.id() {
547            self.log_set_cluster(logging_id, target_cluster_id, target_cluster_name.clone());
548        }
549
550        coord::catalog_serving::check_cluster_restrictions(
551            target_cluster_name.as_str(),
552            &conn_catalog,
553            &plan,
554        )?;
555
556        rbac::check_plan(
557            &conn_catalog,
558            // SideEffectingFunc is handled above (with its own check_plan call) and returns
559            // early, so no target connection role is needed for the remaining plan types here.
560            None,
561            session,
562            &plan,
563            Some(target_cluster_id),
564            &resolved_ids,
565            &sql_impl_ids,
566        )?;
567
568        if let Some((_, wait_future)) =
569            coord::appends::waiting_on_startup_appends(&*catalog, session, &plan)
570        {
571            wait_future.await;
572        }
573
574        let max_query_result_size = Some(session.vars().max_query_result_size());
575
576        let compute_instance_snapshot =
577            ComputeInstanceSnapshot::new_without_collections(cluster.id());
578
579        let optimizer_config = optimize::OptimizerConfig::from(catalog.system_config())
580            .override_from(&catalog.get_cluster(cluster.id()).config.features())
581            // A cluster-scoped LaunchDarkly rule beats a manual `FEATURES` pin.
582            .override_from(
583                &catalog
584                    .state()
585                    .cluster_scoped_optimizer_overrides(cluster.id()),
586            )
587            .override_from(&explain_ctx);
588
589        if cluster.replicas().next().is_none() && explain_ctx.needs_cluster() {
590            return Err(AdapterError::NoClusterReplicasAvailable {
591                name: cluster.name.clone(),
592                is_managed: cluster.is_managed(),
593            });
594        }
595
596        let (_, view_id) = self.transient_id_gen.allocate_id();
597        let (_, index_id) = self.transient_id_gen.allocate_id();
598
599        let target_replica_name = session.vars().cluster_replica();
600        let mut target_replica = target_replica_name
601            .map(|name| {
602                cluster
603                    .replica_id(name)
604                    .ok_or(AdapterError::UnknownClusterReplica {
605                        cluster_name: cluster.name.clone(),
606                        replica_name: name.to_string(),
607                    })
608            })
609            .transpose()?;
610
611        let source_ids = depends_on;
612        // TODO(peek-seq): validate_timeline_context can be expensive in real scenarios (not in
613        // simple benchmarks), because it traverses transitive dependencies even of indexed views and
614        // materialized views (also traversing their MIR plans).
615        let mut timeline_context = catalog.validate_timeline_context(source_ids.iter().copied())?;
616        if matches!(timeline_context, TimelineContext::TimestampIndependent) && contains_temporal {
617            // If the source IDs are timestamp independent but the query contains temporal functions,
618            // then the timeline context needs to be upgraded to timestamp dependent. This is
619            // required because `source_ids` doesn't contain functions.
620            timeline_context = TimelineContext::TimestampDependent;
621        }
622
623        let notices = coord::sequencer::check_log_reads(
624            &catalog,
625            cluster,
626            &source_ids,
627            &mut target_replica,
628            session.vars(),
629        )?;
630        session.add_notices(notices);
631
632        let isolation_level = session.vars().transaction_isolation().clone();
633        let timeline = Coordinator::get_timeline(&timeline_context);
634        let needs_linearized_read_ts =
635            Coordinator::needs_linearized_read_ts(&isolation_level, when);
636
637        let oracle_read_ts = match timeline {
638            Some(timeline) if needs_linearized_read_ts => {
639                let oracle = self.ensure_oracle(timeline).await?;
640                let oracle_read_ts = oracle.read_ts().await;
641                Some(oracle_read_ts)
642            }
643            Some(_) | None => None,
644        };
645
646        let vars = session.vars();
647        let real_time_recency_ts: Option<Timestamp> = if vars.real_time_recency()
648            && vars.transaction_isolation() == &IsolationLevel::StrictSerializable
649            && !session.contains_read_timestamp()
650        {
651            // Only call the coordinator when we actually need real-time recency
652            self.call_coordinator(|tx| Command::DetermineRealTimeRecentTimestamp {
653                source_ids: source_ids.clone(),
654                real_time_recency_timeout: *vars.real_time_recency_timeout(),
655                tx,
656            })
657            .await??
658        } else {
659            None
660        };
661
662        let dataflow_builder =
663            DataflowBuilder::new(catalog.state(), compute_instance_snapshot.clone());
664        let input_id_bundle = dataflow_builder.sufficient_collections(source_ids.clone());
665
666        // ## From sequence_peek_timestamp
667
668        // Warning: This will be false for AS OF queries, even if we are otherwise inside a
669        // multi-statement transaction. (It's also false for FreshestTableWrite, which is currently
670        // only read-then-write queries, which can't be part of multi-statement transactions, so
671        // FreshestTableWrite doesn't matter.)
672        //
673        // TODO(peek-seq): It's not totally clear to me what the intended semantics are for AS OF
674        // queries inside a transaction: We clearly can't use the transaction timestamp, and we don't
675        // do timedomain validation for AS OF queries, which seems natural. But I'm thinking that it
676        // would be the cleanest to just simply disallow AS OF queries inside transactions.
677        let in_immediate_multi_stmt_txn = session.transaction().in_immediate_multi_stmt_txn(when)
678            && !matches!(query_plan, QueryPlan::Subscribe { .. });
679
680        // Fetch or generate a timestamp for this query and fetch or acquire read holds.
681        let (determination, read_holds) = match session.get_transaction_timestamp_determination() {
682            // Use the transaction's timestamp if it exists and this isn't an AS OF query.
683            // (`in_immediate_multi_stmt_txn` is false for AS OF queries.)
684            Some(
685                determination @ TimestampDetermination {
686                    timestamp_context: TimestampContext::TimelineTimestamp { .. },
687                    ..
688                },
689            ) if in_immediate_multi_stmt_txn => {
690                // This is a subsequent (non-AS OF, non-constant) query in a multi-statement
691                // transaction. We now:
692                // - Validate that the query only accesses collections within the transaction's
693                //   timedomain (which we know from the stored read holds).
694                // - Use the transaction's stored timestamp determination.
695                // - Use the (relevant subset of the) transaction's read holds.
696
697                let txn_read_holds_opt = self
698                    .call_coordinator(|tx| Command::GetTransactionReadHoldsBundle {
699                        conn_id: session.conn_id().clone(),
700                        tx,
701                    })
702                    .await?;
703
704                if let Some(txn_read_holds) = txn_read_holds_opt {
705                    let allowed_id_bundle = txn_read_holds.id_bundle();
706                    let outside = input_id_bundle.difference(&allowed_id_bundle);
707
708                    // Queries without a timestamp and timeline can belong to any existing timedomain.
709                    if determination.timestamp_context.contains_timestamp() && !outside.is_empty() {
710                        let valid_names =
711                            allowed_id_bundle.resolve_names(&*catalog, session.conn_id());
712                        let invalid_names = outside.resolve_names(&*catalog, session.conn_id());
713                        return Err(AdapterError::RelationOutsideTimeDomain {
714                            relations: invalid_names,
715                            names: valid_names,
716                        });
717                    }
718
719                    // Extract the subset of read holds for the collections this query accesses.
720                    let read_holds = txn_read_holds.subset(&input_id_bundle);
721
722                    (determination, read_holds)
723                } else {
724                    // This should never happen: we're in a subsequent query of a multi-statement
725                    // transaction (we have a transaction timestamp), but the coordinator has no
726                    // transaction read holds stored. This indicates a bug in the transaction
727                    // handling.
728                    return Err(AdapterError::Internal(
729                        "Missing transaction read holds for multi-statement transaction"
730                            .to_string(),
731                    ));
732                }
733            }
734            _ => {
735                // There is no timestamp determination yet for this transaction. Either:
736                // - We are not in a multi-statement transaction.
737                // - This is the first (non-AS OF) query in a multi-statement transaction.
738                // - This is an AS OF query.
739                // - This is a constant query (`TimestampContext::NoTimestamp`).
740
741                let timedomain_bundle;
742                let determine_bundle = if in_immediate_multi_stmt_txn {
743                    // This is the first (non-AS OF) query in a multi-statement transaction.
744                    // Determine a timestamp that will be valid for anything in any schema
745                    // referenced by the first query.
746                    timedomain_bundle = timedomain_for(
747                        &*catalog,
748                        &dataflow_builder,
749                        &source_ids,
750                        &timeline_context,
751                        session.conn_id(),
752                        target_cluster_id,
753                    )?;
754                    &timedomain_bundle
755                } else {
756                    // Simply use the inputs of the current query.
757                    &input_id_bundle
758                };
759                let (determination, read_holds) = self
760                    .frontend_determine_timestamp(
761                        session,
762                        determine_bundle,
763                        when,
764                        target_cluster_id,
765                        &timeline_context,
766                        oracle_read_ts,
767                        real_time_recency_ts,
768                    )
769                    .await?;
770
771                // If this query pins the timestamp of a multi-statement transaction, store
772                // the read holds in the coordinator, so subsequent queries can validate
773                // against them. The stored holds define the transaction's timedomain, so
774                // only the statement that determines the transaction timestamp may establish
775                // them, over the same id bundle. A timestamp-less determination (e.g. a
776                // constant query) doesn't pin the transaction timestamp, so its holds must
777                // not be stored.
778                if in_immediate_multi_stmt_txn
779                    && determination.timestamp_context.contains_timestamp()
780                {
781                    self.call_coordinator(|tx| Command::StoreTransactionReadHolds {
782                        conn_id: session.conn_id().clone(),
783                        read_holds: read_holds.clone(),
784                        tx,
785                    })
786                    .await?;
787                }
788
789                (determination, read_holds)
790            }
791        };
792
793        {
794            // Assert that we have a read hold for all the collections in our `input_id_bundle`.
795            for id in input_id_bundle.iter() {
796                let s = read_holds.storage_holds.contains_key(&id);
797                let c = read_holds
798                    .compute_ids()
799                    .map(|(_instance, coll)| coll)
800                    .contains(&id);
801                soft_assert_or_log!(
802                    s || c,
803                    "missing read hold for collection {} in `input_id_bundle`; (in_immediate_multi_stmt_txn: {})",
804                    id,
805                    in_immediate_multi_stmt_txn,
806                );
807            }
808
809            // Assert that each part of the `input_id_bundle` corresponds to the right part of
810            // `read_holds`.
811            for id in input_id_bundle.storage_ids.iter() {
812                soft_assert_or_log!(
813                    read_holds.storage_holds.contains_key(id),
814                    "missing storage read hold for collection {} in `input_id_bundle`; (in_immediate_multi_stmt_txn: {})",
815                    id,
816                    in_immediate_multi_stmt_txn,
817                );
818            }
819            for id in input_id_bundle
820                .compute_ids
821                .iter()
822                .flat_map(|(_instance, colls)| colls)
823            {
824                soft_assert_or_log!(
825                    read_holds
826                        .compute_ids()
827                        .map(|(_instance, coll)| coll)
828                        .contains(id),
829                    "missing compute read hold for collection {} in `input_id_bundle`; (in_immediate_multi_stmt_txn: {})",
830                    id,
831                    in_immediate_multi_stmt_txn,
832                );
833            }
834        }
835
836        // (TODO(peek-seq): The below TODO is copied from the old peek sequencing. We should resolve
837        // this when we decide what to with `AS OF` in transactions.)
838        // TODO: Checking for only `InTransaction` and not `Implied` (also `Started`?) seems
839        // arbitrary and we don't recall why we did it (possibly an error!). Change this to always
840        // set the transaction ops. Decide and document what our policy should be on AS OF queries.
841        // Maybe they shouldn't be allowed in transactions at all because it's hard to explain
842        // what's going on there. This should probably get a small design document.
843
844        // We only track the peeks in the session if the query doesn't use AS
845        // OF or we're inside an explicit transaction. The latter case is
846        // necessary to support PG's `BEGIN` semantics, whose behavior can
847        // depend on whether or not reads have occurred in the txn.
848        let requires_linearization = (&explain_ctx).into();
849        let mut transaction_determination = determination.clone();
850        match query_plan {
851            QueryPlan::Subscribe { .. } => {
852                if when.is_transactional() {
853                    session.add_transaction_ops(TransactionOps::Subscribe)?;
854                }
855            }
856            QueryPlan::Select(..) | QueryPlan::CopyTo(..) => {
857                if when.is_transactional() {
858                    session.add_transaction_ops(TransactionOps::Peeks {
859                        determination: transaction_determination,
860                        cluster_id: target_cluster_id,
861                        requires_linearization,
862                    })?;
863                } else if matches!(session.transaction(), &TransactionStatus::InTransaction(_)) {
864                    // If the query uses AS OF, then ignore the timestamp.
865                    transaction_determination.timestamp_context = TimestampContext::NoTimestamp;
866                    session.add_transaction_ops(TransactionOps::Peeks {
867                        determination: transaction_determination,
868                        cluster_id: target_cluster_id,
869                        requires_linearization,
870                    })?;
871                }
872            }
873        }
874
875        let stats = statistics_oracle(
876            session,
877            &source_ids,
878            &determination.timestamp_context.antichain(),
879            true,
880            catalog.system_config(),
881            &*self.storage_collections,
882        )
883        .await
884        .unwrap_or_else(|_| Box::new(EmptyStatisticsOracle));
885
886        // Generate data structures that can be moved to another task where we will perform possibly
887        // expensive optimizations.
888        let timestamp_context = determination.timestamp_context.clone();
889        let session_meta = session.meta();
890        let now = catalog.config().now.clone();
891        let target_cluster_name = target_cluster_name.clone();
892        let needs_plan_insights = explain_ctx.needs_plan_insights();
893        let determination_for_pushdown = if matches!(explain_ctx, ExplainContext::Pushdown) {
894            // This is a hairy data structure, so avoid this clone if we are not in
895            // EXPLAIN FILTER PUSHDOWN.
896            Some(determination.clone())
897        } else {
898            None
899        };
900
901        let span = Span::current();
902
903        // Prepare data for plan insights if needed
904        let catalog_for_insights = if needs_plan_insights {
905            Some(Arc::clone(&catalog))
906        } else {
907            None
908        };
909        let mut compute_instances = BTreeMap::new();
910        if needs_plan_insights {
911            for user_cluster in catalog.user_clusters() {
912                let snapshot = ComputeInstanceSnapshot::new_without_collections(user_cluster.id);
913                compute_instances.insert(user_cluster.name.clone(), snapshot);
914            }
915        }
916
917        let source_ids_for_closure = source_ids.clone();
918
919        let optimization_future: JoinHandle<Result<_, AdapterError>> = match query_plan {
920            QueryPlan::CopyTo(select_plan, mut copy_to_ctx) => {
921                let raw_expr = select_plan.source.clone();
922
923                // COPY TO path: calculate output_batch_count and create copy_to optimizer
924                let worker_counts = cluster.replicas().map(|r| {
925                    let loc = &r.config.location;
926                    loc.workers().unwrap_or_else(|| loc.num_processes())
927                });
928                let max_worker_count = match worker_counts.max() {
929                    Some(count) => u64::cast_from(count),
930                    None => {
931                        return Err(AdapterError::NoClusterReplicasAvailable {
932                            name: cluster.name.clone(),
933                            is_managed: cluster.is_managed(),
934                        });
935                    }
936                };
937                copy_to_ctx.output_batch_count = Some(max_worker_count);
938
939                let mut optimizer = optimize::copy_to::Optimizer::new(
940                    Arc::clone(&catalog),
941                    compute_instance_snapshot,
942                    view_id,
943                    copy_to_ctx,
944                    optimizer_config,
945                    self.optimizer_metrics.clone(),
946                );
947
948                mz_ore::task::spawn_blocking(
949                    || "optimize copy-to",
950                    move || {
951                        span.in_scope(|| {
952                            let _dispatch_guard = explain_ctx.dispatch_guard();
953
954                            // COPY TO path: HIR ⇒ local MIR ⇒ resolve ⇒ global LIR.
955                            let global_lir_plan = optimize::optimize_oneshot(
956                                &mut optimizer,
957                                raw_expr.clone(),
958                                |local_mir_plan| {
959                                    local_mir_plan.resolve(
960                                        timestamp_context.clone(),
961                                        &session_meta,
962                                        stats,
963                                    )
964                                },
965                            )?;
966                            Ok(Execution::CopyToS3 {
967                                global_lir_plan,
968                                source_ids: source_ids_for_closure,
969                            })
970                        })
971                    },
972                )
973            }
974            QueryPlan::Select(select_plan) => {
975                let select_plan = select_plan.clone();
976                let raw_expr = select_plan.source.clone();
977
978                // SELECT/EXPLAIN path: create peek optimizer
979                let mut optimizer = optimize::peek::Optimizer::new(
980                    Arc::clone(&catalog),
981                    compute_instance_snapshot,
982                    select_plan.finishing.clone(),
983                    view_id,
984                    index_id,
985                    optimizer_config,
986                    self.optimizer_metrics.clone(),
987                );
988
989                mz_ore::task::spawn_blocking(
990                    || "optimize peek",
991                    move || {
992                        span.in_scope(|| {
993                            let _dispatch_guard = explain_ctx.dispatch_guard();
994
995                            // SELECT/EXPLAIN path: HIR ⇒ local MIR ⇒ resolve ⇒
996                            // global LIR. We capture the result (rather than
997                            // propagating with `?`) so that a failure can still be
998                            // routed to `EXPLAIN BROKEN` below.
999                            let global_lir_plan_result = optimize::optimize_oneshot(
1000                                &mut optimizer,
1001                                raw_expr.clone(),
1002                                |local_mir_plan| {
1003                                    local_mir_plan.resolve(
1004                                        timestamp_context.clone(),
1005                                        &session_meta,
1006                                        stats,
1007                                    )
1008                                },
1009                            )
1010                            .map_err(AdapterError::from);
1011                            let optimization_finished_at = now();
1012
1013                            let create_insights_ctx =
1014                                |optimizer: &optimize::peek::Optimizer,
1015                                 is_notice: bool|
1016                                 -> Option<Box<PlanInsightsContext>> {
1017                                    if !needs_plan_insights {
1018                                        return None;
1019                                    }
1020
1021                                    let catalog = catalog_for_insights.as_ref()?;
1022
1023                                    let enable_re_optimize = if needs_plan_insights {
1024                                        // Disable any plan insights that use the optimizer if we only want the
1025                                        // notice and plan optimization took longer than the threshold. This is
1026                                        // to prevent a situation where optimizing takes a while and there are
1027                                        // lots of clusters, which would delay peek execution by the product of
1028                                        // those.
1029                                        //
1030                                        // (This heuristic doesn't work well, see #9492.)
1031                                        let dyncfgs = catalog.system_config().dyncfgs();
1032                                        let opt_limit = mz_adapter_types::dyncfgs
1033                                        ::PLAN_INSIGHTS_NOTICE_FAST_PATH_CLUSTERS_OPTIMIZE_DURATION
1034                                            .get(dyncfgs);
1035                                        !(is_notice && optimizer.duration() > opt_limit)
1036                                    } else {
1037                                        false
1038                                    };
1039
1040                                    Some(Box::new(PlanInsightsContext {
1041                                        stmt: select_plan
1042                                            .select
1043                                            .as_deref()
1044                                            .map(Clone::clone)
1045                                            .map(Statement::Select),
1046                                        raw_expr: raw_expr.clone(),
1047                                        catalog: Arc::clone(catalog),
1048                                        compute_instances,
1049                                        target_instance: target_cluster_name,
1050                                        metrics: optimizer.metrics().clone(),
1051                                        finishing: optimizer.finishing().clone(),
1052                                        optimizer_config: optimizer.config().clone(),
1053                                        session: session_meta,
1054                                        timestamp_context,
1055                                        view_id: optimizer.select_id(),
1056                                        index_id: optimizer.index_id(),
1057                                        enable_re_optimize,
1058                                    }))
1059                                };
1060
1061                            let global_lir_plan = match global_lir_plan_result {
1062                                Ok(plan) => plan,
1063                                Err(err) => {
1064                                    let result = if let ExplainContext::Plan(explain_ctx) =
1065                                        explain_ctx
1066                                        && explain_ctx.broken
1067                                    {
1068                                        // EXPLAIN BROKEN: log error and continue with defaults
1069                                        tracing::error!(
1070                                            "error while handling EXPLAIN statement: {}",
1071                                            err
1072                                        );
1073                                        Ok(Execution::ExplainPlan {
1074                                            df_meta: Default::default(),
1075                                            explain_ctx,
1076                                            optimizer,
1077                                            insights_ctx: None,
1078                                        })
1079                                    } else {
1080                                        Err(err)
1081                                    };
1082                                    return result;
1083                                }
1084                            };
1085
1086                            match explain_ctx {
1087                                ExplainContext::Plan(explain_ctx) => {
1088                                    let (_, df_meta, _) = global_lir_plan.unapply();
1089                                    let insights_ctx = create_insights_ctx(&optimizer, false);
1090                                    Ok(Execution::ExplainPlan {
1091                                        df_meta,
1092                                        explain_ctx,
1093                                        optimizer,
1094                                        insights_ctx,
1095                                    })
1096                                }
1097                                ExplainContext::None => Ok(Execution::Peek {
1098                                    global_lir_plan,
1099                                    optimization_finished_at,
1100                                    plan_insights_optimizer_trace: None,
1101                                    finishing: select_plan.finishing,
1102                                    copy_to: select_plan.copy_to,
1103                                    insights_ctx: None,
1104                                }),
1105                                ExplainContext::PlanInsightsNotice(optimizer_trace) => {
1106                                    let insights_ctx = create_insights_ctx(&optimizer, true);
1107                                    Ok(Execution::Peek {
1108                                        global_lir_plan,
1109                                        optimization_finished_at,
1110                                        plan_insights_optimizer_trace: Some(optimizer_trace),
1111                                        finishing: select_plan.finishing,
1112                                        copy_to: select_plan.copy_to,
1113                                        insights_ctx,
1114                                    })
1115                                }
1116                                ExplainContext::Pushdown => {
1117                                    let (plan, _, _) = global_lir_plan.unapply();
1118                                    let imports = match plan {
1119                                        PeekPlan::SlowPath(plan) => plan
1120                                            .desc
1121                                            .source_imports
1122                                            .into_iter()
1123                                            .filter_map(|(id, import)| {
1124                                                import.desc.arguments.operators.map(|plan| {
1125                                                    (id, plan.into_map_filter_project())
1126                                                })
1127                                            })
1128                                            .collect(),
1129                                        PeekPlan::FastPath(_) => {
1130                                            std::collections::BTreeMap::default()
1131                                        }
1132                                    };
1133                                    Ok(Execution::ExplainPushdown {
1134                                        imports,
1135                                        determination: determination_for_pushdown
1136                                            .expect("it's present for the ExplainPushdown case"),
1137                                    })
1138                                }
1139                            }
1140                        })
1141                    },
1142                )
1143            }
1144            QueryPlan::Subscribe(plan) => {
1145                let plan = plan.clone();
1146                let catalog: Arc<Catalog> = Arc::clone(&catalog);
1147                let debug_name = format!("subscribe-{}", index_id);
1148                let mut optimizer = optimize::subscribe::Optimizer::new(
1149                    catalog,
1150                    compute_instance_snapshot.clone(),
1151                    view_id,
1152                    index_id,
1153                    plan.with_snapshot,
1154                    plan.up_to,
1155                    debug_name,
1156                    optimizer_config,
1157                    self.optimizer_metrics.clone(),
1158                );
1159                mz_ore::task::spawn_blocking(
1160                    || "optimize subscribe",
1161                    move || {
1162                        span.in_scope(|| {
1163                            let _dispatch_guard = explain_ctx.dispatch_guard();
1164
1165                            let global_mir_plan = optimizer.catch_unwind_optimize(plan.clone())?;
1166                            let as_of = timestamp_context.timestamp_or_default();
1167
1168                            if let Some(up_to) = optimizer.up_to() {
1169                                if as_of > up_to {
1170                                    return Err(AdapterError::AbsurdSubscribeBounds {
1171                                        as_of,
1172                                        up_to,
1173                                    });
1174                                }
1175                            }
1176                            let local_mir_plan =
1177                                global_mir_plan.resolve(Antichain::from_elem(as_of));
1178
1179                            let global_lir_plan =
1180                                optimizer.catch_unwind_optimize(local_mir_plan)?;
1181                            let optimization_finished_at = now();
1182
1183                            let (df_desc, df_meta) = global_lir_plan.unapply();
1184                            Ok(Execution::Subscribe {
1185                                subscribe_plan: plan,
1186                                df_desc,
1187                                df_meta,
1188                                optimization_finished_at,
1189                            })
1190                        })
1191                    },
1192                )
1193            }
1194        };
1195
1196        let mut optimization_timeout = *session.vars().statement_timeout();
1197        // Timeout of 0 is equivalent to "off", meaning we will wait "forever."
1198        if optimization_timeout == Duration::ZERO {
1199            optimization_timeout = Duration::MAX;
1200        }
1201        let optimization_result =
1202            // Note: spawn_blocking tasks cannot be cancelled, so on timeout we stop waiting but the
1203            // optimization task continues running in the background until completion. See
1204            // https://github.com/MaterializeInc/database-issues/issues/8644 for properly cancelling
1205            // optimizer runs.
1206            match tokio::time::timeout(optimization_timeout, optimization_future).await {
1207                Ok(Ok(result)) => result,
1208                Ok(Err(AdapterError::Optimizer(err))) => {
1209                    return Err(AdapterError::Internal(format!(
1210                        "internal error in optimizer: {}",
1211                        err
1212                    )));
1213                }
1214                Ok(Err(err)) => {
1215                    return Err(err);
1216                }
1217                Err(_elapsed) => {
1218                    warn!("optimize peek timed out after {:?}", optimization_timeout);
1219                    return Err(AdapterError::StatementTimeout);
1220                }
1221            };
1222
1223        // Log optimization finished
1224        if let Some(logging_id) = logging.id() {
1225            self.log_lifecycle_event(logging_id, StatementLifecycleEvent::OptimizationFinished);
1226        }
1227
1228        // Assert that read holds are correct for the execution plan
1229        Self::assert_read_holds_correct(
1230            &read_holds,
1231            &optimization_result,
1232            &determination,
1233            target_cluster_id,
1234            in_immediate_multi_stmt_txn,
1235        );
1236
1237        // Handle the optimization result: either generate EXPLAIN output or continue with execution
1238        match optimization_result {
1239            Execution::ExplainPlan {
1240                df_meta,
1241                explain_ctx,
1242                optimizer,
1243                insights_ctx,
1244            } => {
1245                let rows = coord::sequencer::explain_plan_inner(
1246                    session,
1247                    &catalog,
1248                    df_meta,
1249                    explain_ctx,
1250                    optimizer,
1251                    insights_ctx,
1252                )
1253                .await?;
1254
1255                Ok(Some(ExecuteResponse::SendingRowsImmediate {
1256                    rows: Box::new(rows.into_row_iter()),
1257                }))
1258            }
1259            Execution::ExplainPushdown {
1260                imports,
1261                determination,
1262            } => {
1263                let as_of = determination.timestamp_context.antichain();
1264                let mz_now = determination
1265                    .timestamp_context
1266                    .timestamp()
1267                    .map(|t| ResultSpec::value(Datum::MzTimestamp(*t)))
1268                    .unwrap_or_else(ResultSpec::value_all);
1269
1270                Ok(Some(
1271                    coord::sequencer::explain_pushdown_future_inner(
1272                        session,
1273                        &*catalog,
1274                        &self.storage_collections,
1275                        as_of,
1276                        mz_now,
1277                        imports,
1278                    )
1279                    .await
1280                    .await?,
1281                ))
1282            }
1283            Execution::Peek {
1284                global_lir_plan,
1285                optimization_finished_at: _optimization_finished_at,
1286                plan_insights_optimizer_trace,
1287                finishing,
1288                copy_to,
1289                insights_ctx,
1290            } => {
1291                // Continue with normal execution
1292
1293                // The typ here was generated from the HIR SQL type and simply stored in LIR.
1294                let (peek_plan, df_meta, typ) = global_lir_plan.unapply();
1295
1296                coord::sequencer::emit_optimizer_notices(
1297                    &*catalog,
1298                    session,
1299                    &df_meta.optimizer_notices,
1300                );
1301
1302                // Generate plan insights notice if needed
1303                if let Some(trace) = plan_insights_optimizer_trace {
1304                    let target_cluster = catalog.get_cluster(target_cluster_id);
1305                    let features = OptimizerFeatures::from(catalog.system_config())
1306                        .override_from(&target_cluster.config.features())
1307                        // A cluster-scoped LaunchDarkly rule beats a manual
1308                        // `FEATURES` pin.
1309                        .override_from(
1310                            &catalog
1311                                .state()
1312                                .cluster_scoped_optimizer_overrides(target_cluster_id),
1313                        );
1314                    let insights = trace
1315                        .into_plan_insights(
1316                            &features,
1317                            &catalog.for_session(session),
1318                            Some(finishing.clone()),
1319                            Some(target_cluster),
1320                            df_meta.clone(),
1321                            insights_ctx,
1322                        )
1323                        .await?;
1324                    session.add_notice(AdapterNotice::PlanInsights(insights));
1325                }
1326
1327                let watch_set = logging.id().map(|logging_id| {
1328                    WatchSetCreation::new(
1329                        logging_id,
1330                        catalog.state(),
1331                        &input_id_bundle,
1332                        determination.timestamp_context.timestamp_or_default(),
1333                    )
1334                });
1335
1336                let max_result_size = catalog.system_config().max_result_size();
1337
1338                // Clone determination if we need it for emit_timestamp_notice, since it may be
1339                // moved into Command::ExecuteSlowPathPeek.
1340                let determination_for_notice = if session.vars().emit_timestamp_notice() {
1341                    Some(determination.clone())
1342                } else {
1343                    None
1344                };
1345
1346                let response = match peek_plan {
1347                    PeekPlan::FastPath(fast_path_plan) => {
1348                        if let Some(logging_id) = logging.id() {
1349                            // TODO(peek-seq): Actually, we should log it also for
1350                            // FastPathPlan::Constant. The only reason we are not doing so at the
1351                            // moment is to match the old peek sequencing, so that statement logging
1352                            // tests pass with the frontend peek sequencing turned both on and off.
1353                            //
1354                            // When the old sequencing is removed, we should make a couple of
1355                            // changes in how we log timestamps:
1356                            // - Move this up to just after timestamp determination, so that it
1357                            //   appears in the log as soon as possible.
1358                            // - Do it also for Constant peeks.
1359                            // - Currently, slow-path peeks' timestamp logging is done by
1360                            //   `implement_peek_plan`. We could remove it from there, and just do
1361                            //   it here.
1362                            if !matches!(fast_path_plan, FastPathPlan::Constant(..)) {
1363                                self.log_set_timestamp(
1364                                    logging_id,
1365                                    determination.timestamp_context.timestamp_or_default(),
1366                                );
1367                            }
1368                        }
1369
1370                        let row_set_finishing_seconds =
1371                            session.metrics().row_set_finishing_seconds().clone();
1372
1373                        let peek_stash_read_batch_size_bytes =
1374                            mz_compute_types::dyncfgs::PEEK_RESPONSE_STASH_READ_BATCH_SIZE_BYTES
1375                                .get(catalog.system_config().dyncfgs());
1376                        let peek_stash_read_memory_budget_bytes =
1377                            mz_compute_types::dyncfgs::PEEK_RESPONSE_STASH_READ_MEMORY_BUDGET_BYTES
1378                                .get(catalog.system_config().dyncfgs());
1379
1380                        self.implement_fast_path_peek_plan(
1381                            fast_path_plan,
1382                            determination.timestamp_context.timestamp_or_default(),
1383                            finishing,
1384                            target_cluster_id,
1385                            target_replica,
1386                            typ,
1387                            max_result_size,
1388                            max_query_result_size,
1389                            row_set_finishing_seconds,
1390                            read_holds,
1391                            peek_stash_read_batch_size_bytes,
1392                            peek_stash_read_memory_budget_bytes,
1393                            session.conn_id().clone(),
1394                            source_ids,
1395                            watch_set,
1396                            logging,
1397                        )
1398                        .await?
1399                    }
1400                    PeekPlan::SlowPath(dataflow_plan) => {
1401                        if let Some(logging_id) = logging.id() {
1402                            self.log_set_transient_index_id(logging_id, dataflow_plan.id);
1403                        }
1404
1405                        let response = self
1406                            .call_coordinator(|tx| Command::ExecuteSlowPathPeek {
1407                                dataflow_plan: Box::new(dataflow_plan),
1408                                determination,
1409                                finishing,
1410                                compute_instance: target_cluster_id,
1411                                target_replica,
1412                                intermediate_result_type: typ,
1413                                source_ids,
1414                                conn_id: session.conn_id().clone(),
1415                                max_result_size,
1416                                max_query_result_size,
1417                                watch_set,
1418                                tx,
1419                            })
1420                            .await??;
1421                        // On success the peek is registered in `pending_peeks`,
1422                        // which now owns end-of-execution logging. On error the
1423                        // coordinator logs nothing (see
1424                        // `implement_slow_path_peek`), so the guard stays armed
1425                        // and the caller logs the error.
1426                        logging.defuse();
1427                        response
1428                    }
1429                };
1430
1431                // Add timestamp notice if emit_timestamp_notice is enabled
1432                if let Some(determination) = determination_for_notice {
1433                    self.add_timestamp_notice(
1434                        session,
1435                        target_cluster_id,
1436                        input_id_bundle.clone(),
1437                        determination,
1438                    )
1439                    .await?;
1440                }
1441
1442                Ok(Some(match copy_to {
1443                    None => response,
1444                    // COPY TO STDOUT
1445                    Some(format) => ExecuteResponse::CopyTo {
1446                        format,
1447                        resp: Box::new(response),
1448                    },
1449                }))
1450            }
1451            Execution::Subscribe {
1452                subscribe_plan,
1453                df_desc,
1454                df_meta,
1455                optimization_finished_at: _optimization_finished_at,
1456            } => {
1457                if df_desc.as_of.as_ref().expect("as of set") == &df_desc.until {
1458                    session.add_notice(AdapterNotice::EqualSubscribeBounds {
1459                        bound: *df_desc.until.as_option().expect("as of set"),
1460                    });
1461                }
1462                coord::sequencer::emit_optimizer_notices(
1463                    &*catalog,
1464                    session,
1465                    &df_meta.optimizer_notices,
1466                );
1467
1468                // Test-only synchronization point: parks a subscribe between
1469                // sequencing and dispatch, so a test can land a concurrent DROP
1470                // of a dependency in this window. Used by
1471                // workflow_test_drop_index_during_subscribe_sequencing.
1472                fail::fail_point!("subscribe_before_dispatch");
1473
1474                let response = self
1475                    .call_coordinator(|tx| Command::ExecuteSubscribe {
1476                        df_desc,
1477                        dependency_ids: subscribe_plan.from.depends_on(),
1478                        cluster_id: target_cluster_id,
1479                        replica_id: target_replica,
1480                        conn_id: session.conn_id().clone(),
1481                        session_uuid: session.uuid(),
1482                        read_holds,
1483                        plan: subscribe_plan,
1484                        statement_logging_id: logging.id(),
1485                        tx,
1486                    })
1487                    .await??;
1488                // On success the `Subscribing` response carries the
1489                // coordinator-side logging guard and the protocol layer logs
1490                // the end when the subscribe terminates. On error the
1491                // coordinator logs nothing (see the `ExecuteSubscribe`
1492                // handler), so the guard stays armed and the caller logs the
1493                // error.
1494                logging.defuse();
1495                Ok(Some(response))
1496            }
1497            Execution::CopyToS3 {
1498                global_lir_plan,
1499                source_ids,
1500            } => {
1501                let (df_desc, df_meta) = global_lir_plan.unapply();
1502
1503                coord::sequencer::emit_optimizer_notices(
1504                    &*catalog,
1505                    session,
1506                    &df_meta.optimizer_notices,
1507                );
1508
1509                // Extract S3 sink connection info for preflight check
1510                let sink_id = df_desc.sink_id();
1511                let sinks = &df_desc.sink_exports;
1512                if sinks.len() != 1 {
1513                    return Err(AdapterError::Internal(
1514                        "expected exactly one copy to s3 sink".into(),
1515                    ));
1516                }
1517                let (_, sink_desc) = sinks
1518                    .first_key_value()
1519                    .expect("known to be exactly one copy to s3 sink");
1520                let s3_sink_connection = match &sink_desc.connection {
1521                    mz_compute_types::sinks::ComputeSinkConnection::CopyToS3Oneshot(conn) => {
1522                        conn.clone()
1523                    }
1524                    _ => {
1525                        return Err(AdapterError::Internal(
1526                            "expected copy to s3 oneshot sink".into(),
1527                        ));
1528                    }
1529                };
1530
1531                // Perform S3 preflight check in background task (via coordinator).
1532                // This runs slow S3 operations without blocking the coordinator's main task.
1533                self.call_coordinator(|tx| Command::CopyToPreflight {
1534                    s3_sink_connection,
1535                    sink_id,
1536                    tx,
1537                })
1538                .await??;
1539
1540                // Preflight succeeded, now execute the actual COPY TO dataflow
1541                let watch_set = logging.id().map(|logging_id| {
1542                    WatchSetCreation::new(
1543                        logging_id,
1544                        catalog.state(),
1545                        &input_id_bundle,
1546                        determination.timestamp_context.timestamp_or_default(),
1547                    )
1548                });
1549
1550                // We keep ownership of end-of-execution logging:
1551                // `implement_copy_to` logs nothing, and the final response,
1552                // success or error, comes back through this command and is
1553                // logged by the caller.
1554                let response = self
1555                    .call_coordinator(|tx| Command::ExecuteCopyTo {
1556                        df_desc: Box::new(df_desc),
1557                        compute_instance: target_cluster_id,
1558                        target_replica,
1559                        source_ids,
1560                        conn_id: session.conn_id().clone(),
1561                        watch_set,
1562                        tx,
1563                    })
1564                    .await??;
1565
1566                Ok(Some(response))
1567            }
1568        }
1569    }
1570
1571    /// (Similar to Coordinator::determine_timestamp)
1572    /// Determines the timestamp for a query, acquires read holds that ensure the
1573    /// query remains executable at that time, and returns those.
1574    /// The caller is responsible for eventually dropping those read holds.
1575    ///
1576    /// Note: self is taken &mut because of the lazy fetching in `get_compute_instance_client`.
1577    pub(crate) async fn frontend_determine_timestamp(
1578        &mut self,
1579        session: &Session,
1580        id_bundle: &CollectionIdBundle,
1581        when: &QueryWhen,
1582        compute_instance: ComputeInstanceId,
1583        timeline_context: &TimelineContext,
1584        oracle_read_ts: Option<Timestamp>,
1585        real_time_recency_ts: Option<Timestamp>,
1586    ) -> Result<(TimestampDetermination, ReadHolds), AdapterError> {
1587        // this is copy-pasted from Coordinator
1588
1589        let isolation_level = session.vars().transaction_isolation();
1590
1591        let (read_holds, upper) = self
1592            .acquire_read_holds_and_least_valid_write(id_bundle)
1593            .await
1594            .map_err(|err| {
1595                AdapterError::concurrent_dependency_drop_from_collection_lookup_error(
1596                    err,
1597                    compute_instance,
1598                )
1599            })?;
1600        let (det, read_holds) = <Coordinator as TimestampProvider>::determine_timestamp_for_inner(
1601            session,
1602            id_bundle,
1603            when,
1604            timeline_context,
1605            oracle_read_ts,
1606            real_time_recency_ts,
1607            isolation_level,
1608            read_holds,
1609            upper.clone(),
1610        )?;
1611
1612        session
1613            .metrics()
1614            .by_cluster()
1615            .determine_timestamp(
1616                compute_instance,
1617                det.respond_immediately(),
1618                *isolation_level,
1619            )
1620            .inc();
1621        if !det.respond_immediately()
1622            && isolation_level.is_bounded_staleness()
1623            && real_time_recency_ts.is_none()
1624        {
1625            // Note down the difference between BoundedStaleness and Serializable into a metric.
1626            if let Some(bs_ts) = det.timestamp_context.timestamp() {
1627                let (serializable_det, _tmp_read_holds) =
1628                    <Coordinator as TimestampProvider>::determine_timestamp_for_inner(
1629                        session,
1630                        id_bundle,
1631                        when,
1632                        timeline_context,
1633                        oracle_read_ts,
1634                        real_time_recency_ts,
1635                        &IsolationLevel::Serializable,
1636                        read_holds.clone(),
1637                        upper,
1638                    )?;
1639                if let Some(serializable) = serializable_det.timestamp_context.timestamp() {
1640                    session
1641                        .metrics()
1642                        .by_cluster()
1643                        .timestamp_difference_for_bounded_staleness_ms(compute_instance)
1644                        .observe(f64::cast_lossy(u64::from(
1645                            serializable.saturating_sub(*bs_ts),
1646                        )));
1647                }
1648            }
1649        }
1650
1651        Ok((det, read_holds))
1652    }
1653
1654    fn assert_read_holds_correct(
1655        read_holds: &ReadHolds,
1656        execution: &Execution,
1657        determination: &TimestampDetermination,
1658        target_cluster_id: ClusterId,
1659        in_immediate_multi_stmt_txn: bool,
1660    ) {
1661        // Extract source_imports, index_imports, as_of, and execution_name based on Execution variant
1662        let (source_imports, index_imports, as_of, execution_name): (
1663            Vec<GlobalId>,
1664            Vec<GlobalId>,
1665            Timestamp,
1666            &str,
1667        ) = match execution {
1668            Execution::Peek {
1669                global_lir_plan, ..
1670            } => match global_lir_plan.peek_plan() {
1671                PeekPlan::FastPath(fast_path_plan) => {
1672                    let (sources, indexes) = match fast_path_plan {
1673                        FastPathPlan::Constant(..) => (vec![], vec![]),
1674                        FastPathPlan::PeekExisting(_coll_id, idx_id, ..) => (vec![], vec![*idx_id]),
1675                        FastPathPlan::PeekPersist(global_id, ..) => (vec![*global_id], vec![]),
1676                    };
1677                    (
1678                        sources,
1679                        indexes,
1680                        determination.timestamp_context.timestamp_or_default(),
1681                        "FastPath",
1682                    )
1683                }
1684                PeekPlan::SlowPath(dataflow_plan) => {
1685                    let as_of = dataflow_plan
1686                        .desc
1687                        .as_of
1688                        .clone()
1689                        .expect("dataflow has an as_of")
1690                        .into_element();
1691                    (
1692                        dataflow_plan.desc.source_imports.keys().cloned().collect(),
1693                        dataflow_plan.desc.index_imports.keys().cloned().collect(),
1694                        as_of,
1695                        "SlowPath",
1696                    )
1697                }
1698            },
1699            Execution::CopyToS3 {
1700                global_lir_plan, ..
1701            } => {
1702                let df_desc = global_lir_plan.df_desc();
1703                let as_of = df_desc
1704                    .as_of
1705                    .clone()
1706                    .expect("dataflow has an as_of")
1707                    .into_element();
1708                (
1709                    df_desc.source_imports.keys().cloned().collect(),
1710                    df_desc.index_imports.keys().cloned().collect(),
1711                    as_of,
1712                    "CopyToS3",
1713                )
1714            }
1715            Execution::ExplainPlan { .. } | Execution::ExplainPushdown { .. } => {
1716                // No read holds assertions needed for EXPLAIN variants
1717                return;
1718            }
1719            Execution::Subscribe { df_desc, .. } => {
1720                let as_of = df_desc
1721                    .as_of
1722                    .clone()
1723                    .expect("dataflow has an as_of")
1724                    .into_element();
1725                (
1726                    df_desc.source_imports.keys().cloned().collect(),
1727                    df_desc.index_imports.keys().cloned().collect(),
1728                    as_of,
1729                    "Subscribe",
1730                )
1731            }
1732        };
1733
1734        // Assert that we have some read holds for all the imports of the dataflow.
1735        for id in source_imports.iter() {
1736            soft_assert_or_log!(
1737                read_holds.storage_holds.contains_key(id),
1738                "[{}] missing read hold for the source import {}; (in_immediate_multi_stmt_txn: {})",
1739                execution_name,
1740                id,
1741                in_immediate_multi_stmt_txn,
1742            );
1743        }
1744        for id in index_imports.iter() {
1745            soft_assert_or_log!(
1746                read_holds
1747                    .compute_ids()
1748                    .map(|(_instance, coll)| coll)
1749                    .contains(id),
1750                "[{}] missing read hold for the index import {}; (in_immediate_multi_stmt_txn: {})",
1751                execution_name,
1752                id,
1753                in_immediate_multi_stmt_txn,
1754            );
1755        }
1756
1757        // Also check the holds against the as_of.
1758        for (id, h) in read_holds.storage_holds.iter() {
1759            soft_assert_or_log!(
1760                h.since().less_equal(&as_of),
1761                "[{}] storage read hold at {:?} for collection {} is not enough for as_of {:?}, determination: {:?}; (in_immediate_multi_stmt_txn: {})",
1762                execution_name,
1763                h.since(),
1764                id,
1765                as_of,
1766                determination,
1767                in_immediate_multi_stmt_txn,
1768            );
1769        }
1770        for ((instance, id), h) in read_holds.compute_holds.iter() {
1771            soft_assert_eq_or_log!(
1772                *instance,
1773                target_cluster_id,
1774                "[{}] the read hold on {} is on the wrong cluster; (in_immediate_multi_stmt_txn: {})",
1775                execution_name,
1776                id,
1777                in_immediate_multi_stmt_txn,
1778            );
1779            soft_assert_or_log!(
1780                h.since().less_equal(&as_of),
1781                "[{}] compute read hold at {:?} for collection {} is not enough for as_of {:?}, determination: {:?}; (in_immediate_multi_stmt_txn: {})",
1782                execution_name,
1783                h.since(),
1784                id,
1785                as_of,
1786                determination,
1787                in_immediate_multi_stmt_txn,
1788            );
1789        }
1790    }
1791}
1792
1793/// Enum for branching among various execution steps after optimization
1794enum Execution {
1795    Peek {
1796        global_lir_plan: optimize::peek::GlobalLirPlan,
1797        optimization_finished_at: EpochMillis,
1798        plan_insights_optimizer_trace: Option<OptimizerTrace>,
1799        finishing: RowSetFinishing,
1800        copy_to: Option<plan::CopyFormat>,
1801        insights_ctx: Option<Box<PlanInsightsContext>>,
1802    },
1803    Subscribe {
1804        subscribe_plan: SubscribePlan,
1805        df_desc: mz_compute_types::dataflows::LirDataflowDescription,
1806        df_meta: DataflowMetainfo,
1807        optimization_finished_at: EpochMillis,
1808    },
1809    CopyToS3 {
1810        global_lir_plan: optimize::copy_to::GlobalLirPlan,
1811        source_ids: BTreeSet<GlobalId>,
1812    },
1813    ExplainPlan {
1814        df_meta: DataflowMetainfo,
1815        explain_ctx: ExplainPlanContext,
1816        optimizer: optimize::peek::Optimizer,
1817        insights_ctx: Option<Box<PlanInsightsContext>>,
1818    },
1819    ExplainPushdown {
1820        imports: BTreeMap<GlobalId, mz_expr::MapFilterProject>,
1821        determination: TimestampDetermination,
1822    },
1823}