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