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