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 ©_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 ¶ms,
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, ¶ms, &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}