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 is the first (non-AS OF) query in a multi-statement transaction, store
736                // the read holds in the coordinator, so subsequent queries can validate against
737                // them.
738                if in_immediate_multi_stmt_txn {
739                    self.call_coordinator(|tx| Command::StoreTransactionReadHolds {
740                        conn_id: session.conn_id().clone(),
741                        read_holds: read_holds.clone(),
742                        tx,
743                    })
744                    .await;
745                }
746
747                (determination, read_holds)
748            }
749        };
750
751        {
752            // Assert that we have a read hold for all the collections in our `input_id_bundle`.
753            for id in input_id_bundle.iter() {
754                let s = read_holds.storage_holds.contains_key(&id);
755                let c = read_holds
756                    .compute_ids()
757                    .map(|(_instance, coll)| coll)
758                    .contains(&id);
759                soft_assert_or_log!(
760                    s || c,
761                    "missing read hold for collection {} in `input_id_bundle`; (in_immediate_multi_stmt_txn: {})",
762                    id,
763                    in_immediate_multi_stmt_txn,
764                );
765            }
766
767            // Assert that each part of the `input_id_bundle` corresponds to the right part of
768            // `read_holds`.
769            for id in input_id_bundle.storage_ids.iter() {
770                soft_assert_or_log!(
771                    read_holds.storage_holds.contains_key(id),
772                    "missing storage read hold for collection {} in `input_id_bundle`; (in_immediate_multi_stmt_txn: {})",
773                    id,
774                    in_immediate_multi_stmt_txn,
775                );
776            }
777            for id in input_id_bundle
778                .compute_ids
779                .iter()
780                .flat_map(|(_instance, colls)| colls)
781            {
782                soft_assert_or_log!(
783                    read_holds
784                        .compute_ids()
785                        .map(|(_instance, coll)| coll)
786                        .contains(id),
787                    "missing compute read hold for collection {} in `input_id_bundle`; (in_immediate_multi_stmt_txn: {})",
788                    id,
789                    in_immediate_multi_stmt_txn,
790                );
791            }
792        }
793
794        // (TODO(peek-seq): The below TODO is copied from the old peek sequencing. We should resolve
795        // this when we decide what to with `AS OF` in transactions.)
796        // TODO: Checking for only `InTransaction` and not `Implied` (also `Started`?) seems
797        // arbitrary and we don't recall why we did it (possibly an error!). Change this to always
798        // set the transaction ops. Decide and document what our policy should be on AS OF queries.
799        // Maybe they shouldn't be allowed in transactions at all because it's hard to explain
800        // what's going on there. This should probably get a small design document.
801
802        // We only track the peeks in the session if the query doesn't use AS
803        // OF or we're inside an explicit transaction. The latter case is
804        // necessary to support PG's `BEGIN` semantics, whose behavior can
805        // depend on whether or not reads have occurred in the txn.
806        let requires_linearization = (&explain_ctx).into();
807        let mut transaction_determination = determination.clone();
808        match query_plan {
809            QueryPlan::Subscribe { .. } => {
810                if when.is_transactional() {
811                    session.add_transaction_ops(TransactionOps::Subscribe)?;
812                }
813            }
814            QueryPlan::Select(..) | QueryPlan::CopyTo(..) => {
815                if when.is_transactional() {
816                    session.add_transaction_ops(TransactionOps::Peeks {
817                        determination: transaction_determination,
818                        cluster_id: target_cluster_id,
819                        requires_linearization,
820                    })?;
821                } else if matches!(session.transaction(), &TransactionStatus::InTransaction(_)) {
822                    // If the query uses AS OF, then ignore the timestamp.
823                    transaction_determination.timestamp_context = TimestampContext::NoTimestamp;
824                    session.add_transaction_ops(TransactionOps::Peeks {
825                        determination: transaction_determination,
826                        cluster_id: target_cluster_id,
827                        requires_linearization,
828                    })?;
829                }
830            }
831        }
832
833        // # From peek_optimize
834
835        let stats = statistics_oracle(
836            session,
837            &source_ids,
838            &determination.timestamp_context.antichain(),
839            true,
840            catalog.system_config(),
841            &*self.storage_collections,
842        )
843        .await
844        .unwrap_or_else(|_| Box::new(EmptyStatisticsOracle));
845
846        // Generate data structures that can be moved to another task where we will perform possibly
847        // expensive optimizations.
848        let timestamp_context = determination.timestamp_context.clone();
849        let session_meta = session.meta();
850        let now = catalog.config().now.clone();
851        let target_cluster_name = target_cluster_name.clone();
852        let needs_plan_insights = explain_ctx.needs_plan_insights();
853        let determination_for_pushdown = if matches!(explain_ctx, ExplainContext::Pushdown) {
854            // This is a hairy data structure, so avoid this clone if we are not in
855            // EXPLAIN FILTER PUSHDOWN.
856            Some(determination.clone())
857        } else {
858            None
859        };
860
861        let span = Span::current();
862
863        // Prepare data for plan insights if needed
864        let catalog_for_insights = if needs_plan_insights {
865            Some(Arc::clone(&catalog))
866        } else {
867            None
868        };
869        let mut compute_instances = BTreeMap::new();
870        if needs_plan_insights {
871            for user_cluster in catalog.user_clusters() {
872                let snapshot = ComputeInstanceSnapshot::new_without_collections(user_cluster.id);
873                compute_instances.insert(user_cluster.name.clone(), snapshot);
874            }
875        }
876
877        let source_ids_for_closure = source_ids.clone();
878
879        let optimization_future: JoinHandle<Result<_, AdapterError>> = match query_plan {
880            QueryPlan::CopyTo(select_plan, mut copy_to_ctx) => {
881                let raw_expr = select_plan.source.clone();
882
883                // COPY TO path: calculate output_batch_count and create copy_to optimizer
884                let worker_counts = cluster.replicas().map(|r| {
885                    let loc = &r.config.location;
886                    loc.workers().unwrap_or_else(|| loc.num_processes())
887                });
888                let max_worker_count = match worker_counts.max() {
889                    Some(count) => u64::cast_from(count),
890                    None => {
891                        return Err(AdapterError::NoClusterReplicasAvailable {
892                            name: cluster.name.clone(),
893                            is_managed: cluster.is_managed(),
894                        });
895                    }
896                };
897                copy_to_ctx.output_batch_count = Some(max_worker_count);
898
899                let mut optimizer = optimize::copy_to::Optimizer::new(
900                    Arc::clone(&catalog),
901                    compute_instance_snapshot,
902                    view_id,
903                    copy_to_ctx,
904                    optimizer_config,
905                    self.optimizer_metrics.clone(),
906                );
907
908                mz_ore::task::spawn_blocking(
909                    || "optimize copy-to",
910                    move || {
911                        span.in_scope(|| {
912                            let _dispatch_guard = explain_ctx.dispatch_guard();
913
914                            // COPY TO path: HIR ⇒ local MIR ⇒ resolve ⇒ global LIR.
915                            let global_lir_plan = optimize::optimize_oneshot(
916                                &mut optimizer,
917                                raw_expr.clone(),
918                                |local_mir_plan| {
919                                    local_mir_plan.resolve(
920                                        timestamp_context.clone(),
921                                        &session_meta,
922                                        stats,
923                                    )
924                                },
925                            )?;
926                            Ok(Execution::CopyToS3 {
927                                global_lir_plan,
928                                source_ids: source_ids_for_closure,
929                            })
930                        })
931                    },
932                )
933            }
934            QueryPlan::Select(select_plan) => {
935                let select_plan = select_plan.clone();
936                let raw_expr = select_plan.source.clone();
937
938                // SELECT/EXPLAIN path: create peek optimizer
939                let mut optimizer = optimize::peek::Optimizer::new(
940                    Arc::clone(&catalog),
941                    compute_instance_snapshot,
942                    select_plan.finishing.clone(),
943                    view_id,
944                    index_id,
945                    optimizer_config,
946                    self.optimizer_metrics.clone(),
947                );
948
949                mz_ore::task::spawn_blocking(
950                    || "optimize peek",
951                    move || {
952                        span.in_scope(|| {
953                            let _dispatch_guard = explain_ctx.dispatch_guard();
954
955                            // SELECT/EXPLAIN path: HIR ⇒ local MIR ⇒ resolve ⇒
956                            // global LIR. We capture the result (rather than
957                            // propagating with `?`) so that a failure can still be
958                            // routed to `EXPLAIN BROKEN` below.
959                            let global_lir_plan_result = optimize::optimize_oneshot(
960                                &mut optimizer,
961                                raw_expr.clone(),
962                                |local_mir_plan| {
963                                    local_mir_plan.resolve(
964                                        timestamp_context.clone(),
965                                        &session_meta,
966                                        stats,
967                                    )
968                                },
969                            )
970                            .map_err(AdapterError::from);
971                            let optimization_finished_at = now();
972
973                            let create_insights_ctx =
974                                |optimizer: &optimize::peek::Optimizer,
975                                 is_notice: bool|
976                                 -> Option<Box<PlanInsightsContext>> {
977                                    if !needs_plan_insights {
978                                        return None;
979                                    }
980
981                                    let catalog = catalog_for_insights.as_ref()?;
982
983                                    let enable_re_optimize = if needs_plan_insights {
984                                        // Disable any plan insights that use the optimizer if we only want the
985                                        // notice and plan optimization took longer than the threshold. This is
986                                        // to prevent a situation where optimizing takes a while and there are
987                                        // lots of clusters, which would delay peek execution by the product of
988                                        // those.
989                                        //
990                                        // (This heuristic doesn't work well, see #9492.)
991                                        let dyncfgs = catalog.system_config().dyncfgs();
992                                        let opt_limit = mz_adapter_types::dyncfgs
993                                        ::PLAN_INSIGHTS_NOTICE_FAST_PATH_CLUSTERS_OPTIMIZE_DURATION
994                                            .get(dyncfgs);
995                                        !(is_notice && optimizer.duration() > opt_limit)
996                                    } else {
997                                        false
998                                    };
999
1000                                    Some(Box::new(PlanInsightsContext {
1001                                        stmt: select_plan
1002                                            .select
1003                                            .as_deref()
1004                                            .map(Clone::clone)
1005                                            .map(Statement::Select),
1006                                        raw_expr: raw_expr.clone(),
1007                                        catalog: Arc::clone(catalog),
1008                                        compute_instances,
1009                                        target_instance: target_cluster_name,
1010                                        metrics: optimizer.metrics().clone(),
1011                                        finishing: optimizer.finishing().clone(),
1012                                        optimizer_config: optimizer.config().clone(),
1013                                        session: session_meta,
1014                                        timestamp_context,
1015                                        view_id: optimizer.select_id(),
1016                                        index_id: optimizer.index_id(),
1017                                        enable_re_optimize,
1018                                    }))
1019                                };
1020
1021                            let global_lir_plan = match global_lir_plan_result {
1022                                Ok(plan) => plan,
1023                                Err(err) => {
1024                                    let result = if let ExplainContext::Plan(explain_ctx) =
1025                                        explain_ctx
1026                                        && explain_ctx.broken
1027                                    {
1028                                        // EXPLAIN BROKEN: log error and continue with defaults
1029                                        tracing::error!(
1030                                            "error while handling EXPLAIN statement: {}",
1031                                            err
1032                                        );
1033                                        Ok(Execution::ExplainPlan {
1034                                            df_meta: Default::default(),
1035                                            explain_ctx,
1036                                            optimizer,
1037                                            insights_ctx: None,
1038                                        })
1039                                    } else {
1040                                        Err(err)
1041                                    };
1042                                    return result;
1043                                }
1044                            };
1045
1046                            match explain_ctx {
1047                                ExplainContext::Plan(explain_ctx) => {
1048                                    let (_, df_meta, _) = global_lir_plan.unapply();
1049                                    let insights_ctx = create_insights_ctx(&optimizer, false);
1050                                    Ok(Execution::ExplainPlan {
1051                                        df_meta,
1052                                        explain_ctx,
1053                                        optimizer,
1054                                        insights_ctx,
1055                                    })
1056                                }
1057                                ExplainContext::None => Ok(Execution::Peek {
1058                                    global_lir_plan,
1059                                    optimization_finished_at,
1060                                    plan_insights_optimizer_trace: None,
1061                                    finishing: select_plan.finishing,
1062                                    copy_to: select_plan.copy_to,
1063                                    insights_ctx: None,
1064                                }),
1065                                ExplainContext::PlanInsightsNotice(optimizer_trace) => {
1066                                    let insights_ctx = create_insights_ctx(&optimizer, true);
1067                                    Ok(Execution::Peek {
1068                                        global_lir_plan,
1069                                        optimization_finished_at,
1070                                        plan_insights_optimizer_trace: Some(optimizer_trace),
1071                                        finishing: select_plan.finishing,
1072                                        copy_to: select_plan.copy_to,
1073                                        insights_ctx,
1074                                    })
1075                                }
1076                                ExplainContext::Pushdown => {
1077                                    let (plan, _, _) = global_lir_plan.unapply();
1078                                    let imports = match plan {
1079                                        PeekPlan::SlowPath(plan) => plan
1080                                            .desc
1081                                            .source_imports
1082                                            .into_iter()
1083                                            .filter_map(|(id, import)| {
1084                                                import.desc.arguments.operators.map(|mfp| (id, mfp))
1085                                            })
1086                                            .collect(),
1087                                        PeekPlan::FastPath(_) => {
1088                                            std::collections::BTreeMap::default()
1089                                        }
1090                                    };
1091                                    Ok(Execution::ExplainPushdown {
1092                                        imports,
1093                                        determination: determination_for_pushdown
1094                                            .expect("it's present for the ExplainPushdown case"),
1095                                    })
1096                                }
1097                            }
1098                        })
1099                    },
1100                )
1101            }
1102            QueryPlan::Subscribe(plan) => {
1103                let plan = plan.clone();
1104                let catalog: Arc<Catalog> = Arc::clone(&catalog);
1105                let debug_name = format!("subscribe-{}", index_id);
1106                let mut optimizer = optimize::subscribe::Optimizer::new(
1107                    catalog,
1108                    compute_instance_snapshot.clone(),
1109                    view_id,
1110                    index_id,
1111                    plan.with_snapshot,
1112                    plan.up_to,
1113                    debug_name,
1114                    optimizer_config,
1115                    self.optimizer_metrics.clone(),
1116                );
1117                mz_ore::task::spawn_blocking(
1118                    || "optimize subscribe",
1119                    move || {
1120                        span.in_scope(|| {
1121                            let _dispatch_guard = explain_ctx.dispatch_guard();
1122
1123                            let global_mir_plan = optimizer.catch_unwind_optimize(plan.clone())?;
1124                            let as_of = timestamp_context.timestamp_or_default();
1125
1126                            if let Some(up_to) = optimizer.up_to() {
1127                                if as_of > up_to {
1128                                    return Err(AdapterError::AbsurdSubscribeBounds {
1129                                        as_of,
1130                                        up_to,
1131                                    });
1132                                }
1133                            }
1134                            let local_mir_plan =
1135                                global_mir_plan.resolve(Antichain::from_elem(as_of));
1136
1137                            let global_lir_plan =
1138                                optimizer.catch_unwind_optimize(local_mir_plan)?;
1139                            let optimization_finished_at = now();
1140
1141                            let (df_desc, df_meta) = global_lir_plan.unapply();
1142                            Ok(Execution::Subscribe {
1143                                subscribe_plan: plan,
1144                                df_desc,
1145                                df_meta,
1146                                optimization_finished_at,
1147                            })
1148                        })
1149                    },
1150                )
1151            }
1152        };
1153
1154        let mut optimization_timeout = *session.vars().statement_timeout();
1155        // Timeout of 0 is equivalent to "off", meaning we will wait "forever."
1156        if optimization_timeout == Duration::ZERO {
1157            optimization_timeout = Duration::MAX;
1158        }
1159        let optimization_result =
1160            // Note: spawn_blocking tasks cannot be cancelled, so on timeout we stop waiting but the
1161            // optimization task continues running in the background until completion. See
1162            // https://github.com/MaterializeInc/database-issues/issues/8644 for properly cancelling
1163            // optimizer runs.
1164            match tokio::time::timeout(optimization_timeout, optimization_future).await {
1165                Ok(Ok(result)) => result,
1166                Ok(Err(AdapterError::Optimizer(err))) => {
1167                    return Err(AdapterError::Internal(format!(
1168                        "internal error in optimizer: {}",
1169                        err
1170                    )));
1171                }
1172                Ok(Err(err)) => {
1173                    return Err(err);
1174                }
1175                Err(_elapsed) => {
1176                    warn!("optimize peek timed out after {:?}", optimization_timeout);
1177                    return Err(AdapterError::StatementTimeout);
1178                }
1179            };
1180
1181        // Log optimization finished
1182        if let Some(logging_id) = logging_guard.id() {
1183            self.log_lifecycle_event(logging_id, StatementLifecycleEvent::OptimizationFinished);
1184        }
1185
1186        // Assert that read holds are correct for the execution plan
1187        Self::assert_read_holds_correct(
1188            &read_holds,
1189            &optimization_result,
1190            &determination,
1191            target_cluster_id,
1192            in_immediate_multi_stmt_txn,
1193        );
1194
1195        // Handle the optimization result: either generate EXPLAIN output or continue with execution
1196        match optimization_result {
1197            Execution::ExplainPlan {
1198                df_meta,
1199                explain_ctx,
1200                optimizer,
1201                insights_ctx,
1202            } => {
1203                let rows = coord::sequencer::explain_plan_inner(
1204                    session,
1205                    &catalog,
1206                    df_meta,
1207                    explain_ctx,
1208                    optimizer,
1209                    insights_ctx,
1210                )
1211                .await?;
1212
1213                Ok(Some(ExecuteResponse::SendingRowsImmediate {
1214                    rows: Box::new(rows.into_row_iter()),
1215                }))
1216            }
1217            Execution::ExplainPushdown {
1218                imports,
1219                determination,
1220            } => {
1221                // # From peek_explain_pushdown
1222
1223                let as_of = determination.timestamp_context.antichain();
1224                let mz_now = determination
1225                    .timestamp_context
1226                    .timestamp()
1227                    .map(|t| ResultSpec::value(Datum::MzTimestamp(*t)))
1228                    .unwrap_or_else(ResultSpec::value_all);
1229
1230                Ok(Some(
1231                    coord::sequencer::explain_pushdown_future_inner(
1232                        session,
1233                        &*catalog,
1234                        &self.storage_collections,
1235                        as_of,
1236                        mz_now,
1237                        imports,
1238                    )
1239                    .await
1240                    .await?,
1241                ))
1242            }
1243            Execution::Peek {
1244                global_lir_plan,
1245                optimization_finished_at: _optimization_finished_at,
1246                plan_insights_optimizer_trace,
1247                finishing,
1248                copy_to,
1249                insights_ctx,
1250            } => {
1251                // Continue with normal execution
1252                // # From peek_finish
1253
1254                // The typ here was generated from the HIR SQL type and simply stored in LIR.
1255                let (peek_plan, df_meta, typ) = global_lir_plan.unapply();
1256
1257                coord::sequencer::emit_optimizer_notices(
1258                    &*catalog,
1259                    session,
1260                    &df_meta.optimizer_notices,
1261                );
1262
1263                // Generate plan insights notice if needed
1264                if let Some(trace) = plan_insights_optimizer_trace {
1265                    let target_cluster = catalog.get_cluster(target_cluster_id);
1266                    let features = OptimizerFeatures::from(catalog.system_config())
1267                        .override_from(&target_cluster.config.features())
1268                        // A cluster-scoped LaunchDarkly rule beats a manual
1269                        // `FEATURES` pin.
1270                        .override_from(
1271                            &catalog
1272                                .state()
1273                                .cluster_scoped_optimizer_overrides(target_cluster_id),
1274                        );
1275                    let insights = trace
1276                        .into_plan_insights(
1277                            &features,
1278                            &catalog.for_session(session),
1279                            Some(finishing.clone()),
1280                            Some(target_cluster),
1281                            df_meta.clone(),
1282                            insights_ctx,
1283                        )
1284                        .await?;
1285                    session.add_notice(AdapterNotice::PlanInsights(insights));
1286                }
1287
1288                // # Now back to peek_finish
1289
1290                let watch_set = logging_guard.id().map(|logging_id| {
1291                    WatchSetCreation::new(
1292                        logging_id,
1293                        catalog.state(),
1294                        &input_id_bundle,
1295                        determination.timestamp_context.timestamp_or_default(),
1296                    )
1297                });
1298
1299                let max_result_size = catalog.system_config().max_result_size();
1300
1301                // Clone determination if we need it for emit_timestamp_notice, since it may be
1302                // moved into Command::ExecuteSlowPathPeek.
1303                let determination_for_notice = if session.vars().emit_timestamp_notice() {
1304                    Some(determination.clone())
1305                } else {
1306                    None
1307                };
1308
1309                let response = match peek_plan {
1310                    PeekPlan::FastPath(fast_path_plan) => {
1311                        if let Some(logging_id) = logging_guard.id() {
1312                            // TODO(peek-seq): Actually, we should log it also for
1313                            // FastPathPlan::Constant. The only reason we are not doing so at the
1314                            // moment is to match the old peek sequencing, so that statement logging
1315                            // tests pass with the frontend peek sequencing turned both on and off.
1316                            //
1317                            // When the old sequencing is removed, we should make a couple of
1318                            // changes in how we log timestamps:
1319                            // - Move this up to just after timestamp determination, so that it
1320                            //   appears in the log as soon as possible.
1321                            // - Do it also for Constant peeks.
1322                            // - Currently, slow-path peeks' timestamp logging is done by
1323                            //   `implement_peek_plan`. We could remove it from there, and just do
1324                            //   it here.
1325                            if !matches!(fast_path_plan, FastPathPlan::Constant(..)) {
1326                                self.log_set_timestamp(
1327                                    logging_id,
1328                                    determination.timestamp_context.timestamp_or_default(),
1329                                );
1330                            }
1331                        }
1332
1333                        let row_set_finishing_seconds =
1334                            session.metrics().row_set_finishing_seconds().clone();
1335
1336                        let peek_stash_read_batch_size_bytes =
1337                            mz_compute_types::dyncfgs::PEEK_RESPONSE_STASH_READ_BATCH_SIZE_BYTES
1338                                .get(catalog.system_config().dyncfgs());
1339                        let peek_stash_read_memory_budget_bytes =
1340                            mz_compute_types::dyncfgs::PEEK_RESPONSE_STASH_READ_MEMORY_BUDGET_BYTES
1341                                .get(catalog.system_config().dyncfgs());
1342
1343                        self.implement_fast_path_peek_plan(
1344                            fast_path_plan,
1345                            determination.timestamp_context.timestamp_or_default(),
1346                            finishing,
1347                            target_cluster_id,
1348                            target_replica,
1349                            typ,
1350                            max_result_size,
1351                            max_query_result_size,
1352                            row_set_finishing_seconds,
1353                            read_holds,
1354                            peek_stash_read_batch_size_bytes,
1355                            peek_stash_read_memory_budget_bytes,
1356                            session.conn_id().clone(),
1357                            source_ids,
1358                            watch_set,
1359                            logging_guard,
1360                        )
1361                        .await?
1362                    }
1363                    PeekPlan::SlowPath(dataflow_plan) => {
1364                        if let Some(logging_id) = logging_guard.id() {
1365                            self.log_set_transient_index_id(logging_id, dataflow_plan.id);
1366                        }
1367
1368                        let response = self
1369                            .call_coordinator(|tx| Command::ExecuteSlowPathPeek {
1370                                dataflow_plan: Box::new(dataflow_plan),
1371                                determination,
1372                                finishing,
1373                                compute_instance: target_cluster_id,
1374                                target_replica,
1375                                intermediate_result_type: typ,
1376                                source_ids,
1377                                conn_id: session.conn_id().clone(),
1378                                max_result_size,
1379                                max_query_result_size,
1380                                watch_set,
1381                                tx,
1382                            })
1383                            .await?;
1384                        // On success the peek is registered in `pending_peeks`,
1385                        // which now owns end-of-execution logging. On error the
1386                        // coordinator logs nothing (see
1387                        // `implement_slow_path_peek`), so the guard stays armed
1388                        // and the caller logs the error.
1389                        logging_guard.defuse();
1390                        response
1391                    }
1392                };
1393
1394                // Add timestamp notice if emit_timestamp_notice is enabled
1395                if let Some(determination) = determination_for_notice {
1396                    let explanation = self
1397                        .call_coordinator(|tx| Command::ExplainTimestamp {
1398                            conn_id: session.conn_id().clone(),
1399                            session_wall_time: session.pcx().wall_time,
1400                            cluster_id: target_cluster_id,
1401                            id_bundle: input_id_bundle.clone(),
1402                            determination,
1403                            tx,
1404                        })
1405                        .await;
1406                    session.add_notice(AdapterNotice::QueryTimestamp { explanation });
1407                }
1408
1409                Ok(Some(match copy_to {
1410                    None => response,
1411                    // COPY TO STDOUT
1412                    Some(format) => ExecuteResponse::CopyTo {
1413                        format,
1414                        resp: Box::new(response),
1415                    },
1416                }))
1417            }
1418            Execution::Subscribe {
1419                subscribe_plan,
1420                df_desc,
1421                df_meta,
1422                optimization_finished_at: _optimization_finished_at,
1423            } => {
1424                if df_desc.as_of.as_ref().expect("as of set") == &df_desc.until {
1425                    session.add_notice(AdapterNotice::EqualSubscribeBounds {
1426                        bound: *df_desc.until.as_option().expect("as of set"),
1427                    });
1428                }
1429                coord::sequencer::emit_optimizer_notices(
1430                    &*catalog,
1431                    session,
1432                    &df_meta.optimizer_notices,
1433                );
1434
1435                // Test-only synchronization point: parks a subscribe between
1436                // sequencing and dispatch, so a test can land a concurrent DROP
1437                // of a dependency in this window. Used by
1438                // workflow_test_drop_index_during_subscribe_sequencing.
1439                fail::fail_point!("subscribe_before_dispatch");
1440
1441                let response = self
1442                    .call_coordinator(|tx| Command::ExecuteSubscribe {
1443                        df_desc,
1444                        dependency_ids: subscribe_plan.from.depends_on(),
1445                        cluster_id: target_cluster_id,
1446                        replica_id: target_replica,
1447                        conn_id: session.conn_id().clone(),
1448                        session_uuid: session.uuid(),
1449                        read_holds,
1450                        plan: subscribe_plan,
1451                        statement_logging_id: logging_guard.id(),
1452                        tx,
1453                    })
1454                    .await?;
1455                // On success the `Subscribing` response carries the
1456                // coordinator-side logging guard and the protocol layer logs
1457                // the end when the subscribe terminates. On error the
1458                // coordinator logs nothing (see the `ExecuteSubscribe`
1459                // handler), so the guard stays armed and the caller logs the
1460                // error.
1461                logging_guard.defuse();
1462                Ok(Some(response))
1463            }
1464            Execution::CopyToS3 {
1465                global_lir_plan,
1466                source_ids,
1467            } => {
1468                let (df_desc, df_meta) = global_lir_plan.unapply();
1469
1470                coord::sequencer::emit_optimizer_notices(
1471                    &*catalog,
1472                    session,
1473                    &df_meta.optimizer_notices,
1474                );
1475
1476                // Extract S3 sink connection info for preflight check
1477                let sink_id = df_desc.sink_id();
1478                let sinks = &df_desc.sink_exports;
1479                if sinks.len() != 1 {
1480                    return Err(AdapterError::Internal(
1481                        "expected exactly one copy to s3 sink".into(),
1482                    ));
1483                }
1484                let (_, sink_desc) = sinks
1485                    .first_key_value()
1486                    .expect("known to be exactly one copy to s3 sink");
1487                let s3_sink_connection = match &sink_desc.connection {
1488                    mz_compute_types::sinks::ComputeSinkConnection::CopyToS3Oneshot(conn) => {
1489                        conn.clone()
1490                    }
1491                    _ => {
1492                        return Err(AdapterError::Internal(
1493                            "expected copy to s3 oneshot sink".into(),
1494                        ));
1495                    }
1496                };
1497
1498                // Perform S3 preflight check in background task (via coordinator).
1499                // This runs slow S3 operations without blocking the coordinator's main task.
1500                self.call_coordinator(|tx| Command::CopyToPreflight {
1501                    s3_sink_connection,
1502                    sink_id,
1503                    tx,
1504                })
1505                .await?;
1506
1507                // Preflight succeeded, now execute the actual COPY TO dataflow
1508                let watch_set = logging_guard.id().map(|logging_id| {
1509                    WatchSetCreation::new(
1510                        logging_id,
1511                        catalog.state(),
1512                        &input_id_bundle,
1513                        determination.timestamp_context.timestamp_or_default(),
1514                    )
1515                });
1516
1517                // We keep ownership of end-of-execution logging:
1518                // `implement_copy_to` logs nothing, and the final response,
1519                // success or error, comes back through this command and is
1520                // logged by the caller.
1521                let response = self
1522                    .call_coordinator(|tx| Command::ExecuteCopyTo {
1523                        df_desc: Box::new(df_desc),
1524                        compute_instance: target_cluster_id,
1525                        target_replica,
1526                        source_ids,
1527                        conn_id: session.conn_id().clone(),
1528                        watch_set,
1529                        tx,
1530                    })
1531                    .await?;
1532
1533                Ok(Some(response))
1534            }
1535        }
1536    }
1537
1538    /// (Similar to Coordinator::determine_timestamp)
1539    /// Determines the timestamp for a query, acquires read holds that ensure the
1540    /// query remains executable at that time, and returns those.
1541    /// The caller is responsible for eventually dropping those read holds.
1542    ///
1543    /// Note: self is taken &mut because of the lazy fetching in `get_compute_instance_client`.
1544    pub(crate) async fn frontend_determine_timestamp(
1545        &mut self,
1546        session: &Session,
1547        id_bundle: &CollectionIdBundle,
1548        when: &QueryWhen,
1549        compute_instance: ComputeInstanceId,
1550        timeline_context: &TimelineContext,
1551        oracle_read_ts: Option<Timestamp>,
1552        real_time_recency_ts: Option<Timestamp>,
1553    ) -> Result<(TimestampDetermination, ReadHolds), AdapterError> {
1554        // this is copy-pasted from Coordinator
1555
1556        let isolation_level = session.vars().transaction_isolation();
1557
1558        let (read_holds, upper) = self
1559            .acquire_read_holds_and_least_valid_write(id_bundle)
1560            .await
1561            .map_err(|err| {
1562                AdapterError::concurrent_dependency_drop_from_collection_lookup_error(
1563                    err,
1564                    compute_instance,
1565                )
1566            })?;
1567        let (det, read_holds) = <Coordinator as TimestampProvider>::determine_timestamp_for_inner(
1568            session,
1569            id_bundle,
1570            when,
1571            timeline_context,
1572            oracle_read_ts,
1573            real_time_recency_ts,
1574            isolation_level,
1575            read_holds,
1576            upper.clone(),
1577        )?;
1578
1579        session
1580            .metrics()
1581            .determine_timestamp(&[
1582                match det.respond_immediately() {
1583                    true => "true",
1584                    false => "false",
1585                },
1586                isolation_level.as_variant_str(),
1587                &compute_instance.to_string(),
1588            ])
1589            .inc();
1590        if !det.respond_immediately()
1591            && isolation_level == &IsolationLevel::StrictSerializable
1592            && real_time_recency_ts.is_none()
1593        {
1594            // Note down the difference between StrictSerializable and Serializable into a metric.
1595            if let Some(strict) = det.timestamp_context.timestamp() {
1596                let (serializable_det, _tmp_read_holds) =
1597                    <Coordinator as TimestampProvider>::determine_timestamp_for_inner(
1598                        session,
1599                        id_bundle,
1600                        when,
1601                        timeline_context,
1602                        oracle_read_ts,
1603                        real_time_recency_ts,
1604                        &IsolationLevel::Serializable,
1605                        read_holds.clone(),
1606                        upper.clone(),
1607                    )?;
1608                if let Some(serializable) = serializable_det.timestamp_context.timestamp() {
1609                    session
1610                        .metrics()
1611                        .timestamp_difference_for_strict_serializable_ms(&[compute_instance
1612                            .to_string()
1613                            .as_ref()])
1614                        .observe(f64::cast_lossy(u64::from(
1615                            strict.saturating_sub(*serializable),
1616                        )));
1617                }
1618            }
1619        }
1620        if !det.respond_immediately()
1621            && isolation_level.is_bounded_staleness()
1622            && real_time_recency_ts.is_none()
1623        {
1624            // Note down the difference between BoundedStaleness and Serializable into a metric.
1625            if let Some(bs_ts) = det.timestamp_context.timestamp() {
1626                let (serializable_det, _tmp_read_holds) =
1627                    <Coordinator as TimestampProvider>::determine_timestamp_for_inner(
1628                        session,
1629                        id_bundle,
1630                        when,
1631                        timeline_context,
1632                        oracle_read_ts,
1633                        real_time_recency_ts,
1634                        &IsolationLevel::Serializable,
1635                        read_holds.clone(),
1636                        upper,
1637                    )?;
1638                if let Some(serializable) = serializable_det.timestamp_context.timestamp() {
1639                    session
1640                        .metrics()
1641                        .timestamp_difference_for_bounded_staleness_ms(&[compute_instance
1642                            .to_string()
1643                            .as_ref()])
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: DataflowDescription<mz_compute_types::plan::LirRelationExpr>,
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}