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