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