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