1#![allow(clippy::pub_use)]
12
13use std::collections::{BTreeMap, BTreeSet};
16use std::str::FromStr;
17use std::sync::Arc;
18use std::time::Duration;
19
20use futures::FutureExt;
21use futures::future::LocalBoxFuture;
22use futures::stream::FuturesOrdered;
23use http::Uri;
24use inner::return_if_err;
25use maplit::btreemap;
26use mz_catalog::memory::objects::Cluster;
27use mz_controller_types::ReplicaId;
28use mz_expr::row::RowCollection;
29use mz_expr::{Eval, MapFilterProject, MirRelationExpr, ResultSpec, RowSetFinishing};
30use mz_ore::cast::CastFrom;
31use mz_ore::soft_panic_or_log;
32use mz_ore::tracing::OpenTelemetryContext;
33use mz_persist_client::stats::SnapshotPartStats;
34use mz_repr::explain::{ExprHumanizerExt, TransientItem};
35use mz_repr::{CatalogItemId, Datum, Diff, GlobalId, IntoRowIterator, Row, RowArena, Timestamp};
36use mz_sql::catalog::{CatalogError, SessionCatalog};
37use mz_sql::names::ResolvedIds;
38use mz_sql::plan::{
39 self, AbortTransactionPlan, CommitTransactionPlan, CreateRolePlan, CreateSourcePlanBundle,
40 FetchPlan, HirScalarExpr, MutationKind, Params, Plan, PlanKind, RaisePlan,
41};
42use mz_sql::rbac;
43use mz_sql::session::metadata::SessionMetadata;
44use mz_sql::session::vars;
45use mz_sql::session::vars::SessionVars;
46use mz_sql_parser::ast::{Raw, Statement};
47use mz_storage_client::client::TableData;
48use mz_storage_client::storage_collections::StorageCollections;
49use mz_storage_types::connections::inline::IntoInlineConnection;
50use mz_storage_types::controller::StorageError;
51use mz_storage_types::stats::RelationPartStats;
52use mz_transform::dataflow::DataflowMetainfo;
53use mz_transform::notice::{OptimizerNoticeApi, OptimizerNoticeKind, RawOptimizerNotice};
54use mz_transform::{EmptyStatisticsOracle, StatisticsOracle};
55use timely::progress::Antichain;
56use tokio::sync::oneshot;
57use tracing::{Instrument, Level, Span, event, warn};
58
59use crate::ExecuteContext;
60use crate::catalog::{Catalog, CatalogState};
61use crate::command::{Command, ExecuteResponse, Response};
62use crate::coord::appends::DeferredPlan;
63use crate::coord::validity::PlanValidity;
64use crate::coord::{
65 Coordinator, DeferredPlanStatement, ExplainPlanContext, Message, PlanStatement, TargetCluster,
66 catalog_serving,
67};
68use crate::error::AdapterError;
69use crate::explain::insights::PlanInsightsContext;
70use crate::notice::AdapterNotice;
71use crate::optimize::dataflows::{EvalTime, ExprPrep, ExprPrepOneShot};
72use crate::optimize::peek;
73use crate::session::{
74 EndTransactionAction, Session, StateRevision, TransactionOps, TransactionStatus, WriteOp,
75};
76use crate::util::ClientTransmitter;
77
78mod inner;
97
98impl Coordinator {
99 pub(crate) fn sequence_plan(
103 &mut self,
104 mut ctx: ExecuteContext,
105 plan: Plan,
106 resolved_ids: ResolvedIds,
107 sql_impl_resolved_ids: ResolvedIds,
108 ) -> LocalBoxFuture<'_, ()> {
109 async move {
110 let responses = ExecuteResponse::generated_from(&PlanKind::from(&plan));
111 ctx.tx_mut().set_allowed(responses);
112
113 if self.controller.read_only() && !plan.allowed_in_read_only() {
114 ctx.retire(Err(AdapterError::ReadOnly));
115 return;
116 }
117
118 if let Some((dependencies, wait_future)) =
121 super::appends::waiting_on_startup_appends(self.catalog(), ctx.session_mut(), &plan)
122 {
123 let conn_id = ctx.session().conn_id();
124 tracing::debug!(%conn_id, "deferring plan for startup appends");
125
126 let role_metadata = ctx.session().role_metadata().clone();
127 let validity =
128 PlanValidity::new(&self.catalog, dependencies, None, None, role_metadata);
129 let deferred_plan = DeferredPlan {
130 ctx,
131 plan,
132 validity,
133 resolved_ids,
134 sql_impl_resolved_ids,
135 };
136 self.defer_plan(wait_future, deferred_plan);
137
138 return;
141 };
142
143 let target_cluster = match ctx.session().transaction().cluster() {
145 Some(cluster_id) => TargetCluster::Transaction(cluster_id),
147 None => {
149 let session_catalog = self.catalog.for_session(ctx.session());
150 catalog_serving::auto_run_on_catalog_server(
151 &session_catalog,
152 ctx.session(),
153 &plan,
154 )
155 }
156 };
157 let (target_cluster_id, target_cluster_name) = match self
158 .catalog()
159 .resolve_target_cluster(target_cluster, ctx.session())
160 {
161 Ok(cluster) => (Some(cluster.id), Some(cluster.name.clone())),
162 Err(_) => (None, None),
163 };
164
165 if let (Some(cluster_id), Some(cluster_name), Some(statement_id)) = (
166 target_cluster_id,
167 target_cluster_name.clone(),
168 ctx.extra().contents(),
169 ) {
170 self.set_statement_execution_cluster(statement_id, cluster_id, cluster_name);
171 }
172
173 let session_catalog = self.catalog.for_session(ctx.session());
174
175 if let Some(cluster_name) = &target_cluster_name {
176 if let Err(e) = catalog_serving::check_cluster_restrictions(
177 cluster_name,
178 &session_catalog,
179 &plan,
180 ) {
181 return ctx.retire(Err(e));
182 }
183 }
184
185 if let Err(e) = rbac::check_plan(
188 &session_catalog,
189 None,
190 ctx.session(),
191 &plan,
192 target_cluster_id,
193 &resolved_ids,
194 &sql_impl_resolved_ids,
195 ) {
196 return ctx.retire(Err(e.into()));
197 }
198
199 match plan {
200 Plan::CreateSource(plan) => {
201 let (item_id, global_id) = return_if_err!(self.allocate_user_id().await, ctx);
202 let result = self
203 .sequence_create_source(
204 &mut ctx,
205 vec![CreateSourcePlanBundle {
206 item_id,
207 global_id,
208 plan,
209 resolved_ids,
210 available_source_references: None,
211 }],
212 )
213 .await;
214 ctx.retire(result);
215 }
216 Plan::CreateSources(plans) => {
217 assert!(
218 resolved_ids.is_empty(),
219 "each plan has separate resolved_ids"
220 );
221 let result = self.sequence_create_source(&mut ctx, plans).await;
222 ctx.retire(result);
223 }
224 Plan::CreateConnection(plan) => {
225 self.sequence_create_connection(ctx, plan, resolved_ids)
226 .await;
227 }
228 Plan::CreateDatabase(plan) => {
229 let result = self.sequence_create_database(ctx.session_mut(), plan).await;
230 ctx.retire(result);
231 }
232 Plan::CreateSchema(plan) => {
233 let result = self.sequence_create_schema(ctx.session_mut(), plan).await;
234 ctx.retire(result);
235 }
236 Plan::CreateRole(plan) => {
237 let result = self
238 .sequence_create_role(Some(ctx.session().conn_id()), plan)
239 .await;
240 if let Some(notice) = self.should_emit_rbac_notice(ctx.session()) {
241 ctx.session().add_notice(notice);
242 }
243 ctx.retire(result);
244 }
245 Plan::CreateCluster(plan) => {
246 let result = self.sequence_create_cluster(ctx.session(), plan).await;
247 ctx.retire(result);
248 }
249 Plan::CreateClusterReplica(plan) => {
250 let result = self
251 .sequence_create_cluster_replica(ctx.session(), plan)
252 .await;
253 ctx.retire(result);
254 }
255 Plan::CreateTable(plan) => {
256 let result = self
257 .sequence_create_table(&mut ctx, plan, resolved_ids)
258 .await;
259 ctx.retire(result);
260 }
261 Plan::CreateSecret(plan) => {
262 self.sequence_create_secret(ctx, plan).await;
263 }
264 Plan::CreateSink(plan) => {
265 self.sequence_create_sink(ctx, plan, resolved_ids).await;
266 }
267 Plan::CreateView(plan) => {
268 self.sequence_create_view(ctx, plan, resolved_ids).await;
269 }
270 Plan::CreateMaterializedView(plan) => {
271 self.sequence_create_materialized_view(ctx, plan, resolved_ids)
272 .await;
273 }
274 Plan::CreateIndex(plan) => {
275 self.sequence_create_index(ctx, plan, resolved_ids).await;
276 }
277 Plan::CreateMetricSink(plan) => {
278 self.sequence_create_metric_sink(ctx, plan, resolved_ids)
279 .await;
280 }
281 Plan::CreateType(plan) => {
282 let result = self
283 .sequence_create_type(ctx.session(), plan, resolved_ids)
284 .await;
285 ctx.retire(result);
286 }
287 Plan::CreateNetworkPolicy(plan) => {
288 let res = self
289 .sequence_create_network_policy(ctx.session(), plan)
290 .await;
291 ctx.retire(res);
292 }
293 Plan::Comment(plan) => {
294 let result = self.sequence_comment_on(ctx.session(), plan).await;
295 ctx.retire(result);
296 }
297 Plan::DropObjects(plan) => {
298 let result = self.sequence_drop_objects(&mut ctx, plan).await;
299 ctx.retire(result);
300 }
301 Plan::DropOwned(plan) => {
302 let result = self.sequence_drop_owned(ctx.session_mut(), plan).await;
303 ctx.retire(result);
304 }
305 Plan::EmptyQuery => {
306 ctx.retire(Ok(ExecuteResponse::EmptyQuery));
307 }
308 Plan::ShowAllVariables => {
309 let result = self.sequence_show_all_variables(ctx.session());
310 ctx.retire(result);
311 }
312 Plan::ShowVariable(plan) => {
313 let result = self.sequence_show_variable(ctx.session(), plan);
314 ctx.retire(result);
315 }
316 Plan::InspectShard(plan) => {
317 let result = self.sequence_inspect_shard(ctx.session(), plan).await;
319 ctx.retire(result);
320 }
321 Plan::SetVariable(plan) => {
322 let result = self.sequence_set_variable(ctx.session_mut(), plan);
323 ctx.retire(result);
324 }
325 Plan::ResetVariable(plan) => {
326 let result = self.sequence_reset_variable(ctx.session_mut(), plan);
327 ctx.retire(result);
328 }
329 Plan::SetTransaction(plan) => {
330 let result = self.sequence_set_transaction(ctx.session_mut(), plan);
331 ctx.retire(result);
332 }
333 Plan::StartTransaction(plan) => {
334 if matches!(
335 ctx.session().transaction(),
336 TransactionStatus::InTransaction(_)
337 ) {
338 ctx.session()
339 .add_notice(AdapterNotice::ExistingTransactionInProgress);
340 }
341 let result = ctx.session_mut().start_transaction(
342 self.now_datetime(),
343 plan.access,
344 plan.isolation_level,
345 );
346 ctx.retire(result.map(|_| ExecuteResponse::StartedTransaction))
347 }
348 Plan::CommitTransaction(CommitTransactionPlan {
349 ref transaction_type,
350 })
351 | Plan::AbortTransaction(AbortTransactionPlan {
352 ref transaction_type,
353 }) => {
354 if ctx.session().transaction().is_ddl() {
357 if let Ok(guard) = self.serialized_ddl.try_lock_owned() {
358 let prev = self
359 .active_conns
360 .get_mut(ctx.session().conn_id())
361 .expect("connection must exist")
362 .deferred_lock
363 .replace(guard);
364 assert!(
365 prev.is_none(),
366 "connections should have at most one lock guard"
367 );
368 } else {
369 self.serialized_ddl.push_back(DeferredPlanStatement {
370 ctx,
371 ps: PlanStatement::Plan {
372 plan,
373 resolved_ids,
374 sql_impl_resolved_ids,
375 },
376 });
377 return;
378 }
379 }
380
381 let action = match &plan {
382 Plan::CommitTransaction(_) => EndTransactionAction::Commit,
383 Plan::AbortTransaction(_) => EndTransactionAction::Rollback,
384 _ => unreachable!(),
385 };
386 if ctx.session().transaction().is_implicit() && !transaction_type.is_implicit()
387 {
388 ctx.session().add_notice(
393 AdapterNotice::ExplicitTransactionControlInImplicitTransaction,
394 );
395 }
396 self.sequence_end_transaction(ctx, action).await;
397 }
398 Plan::Subscribe(plan) => {
399 self.sequence_subscribe(ctx, plan, target_cluster).await;
400 }
401 Plan::ShowCreate(plan) => {
402 ctx.retire(Ok(Self::send_immediate_rows(plan.row)));
403 }
404 Plan::CopyFrom(plan) => {
405 self.sequence_copy_from(ctx, plan, target_cluster).await;
406 }
407 Plan::ExplainPlan(plan) => {
408 self.sequence_explain_plan(ctx, plan, target_cluster).await;
409 }
410 Plan::ExplainPushdown(plan) => {
411 self.sequence_explain_pushdown(ctx, plan).await;
412 }
413 Plan::ExplainSinkSchema(plan) => {
414 let result = self.sequence_explain_schema(plan);
415 ctx.retire(result);
416 }
417 Plan::ExplainTimestamp(plan) => {
418 self.sequence_explain_timestamp(ctx, plan, target_cluster)
419 .await;
420 }
421 plan @ (Plan::CopyTo(_)
426 | Plan::Insert(_)
427 | Plan::ReadThenWrite(_)
428 | Plan::Select(_)
429 | Plan::ShowColumns(_)
430 | Plan::SideEffectingFunc(_)) => {
431 let msg = format!(
432 "{:?} plan reached the coordinator despite frontend routing",
433 PlanKind::from(&plan)
434 );
435 soft_panic_or_log!("{msg}");
436 ctx.retire(Err(AdapterError::Internal(msg)));
437 }
438 Plan::AlterNoop(plan) => {
439 ctx.retire(Ok(ExecuteResponse::AlteredObject(plan.object_type)));
440 }
441 Plan::AlterCluster(plan) => {
442 self.sequence_alter_cluster_staged(ctx, plan).await;
443 }
444 Plan::AlterClusterRename(plan) => {
445 let result = self.sequence_alter_cluster_rename(&mut ctx, plan).await;
446 ctx.retire(result);
447 }
448 Plan::AlterClusterSwap(plan) => {
449 let result = self.sequence_alter_cluster_swap(&mut ctx, plan).await;
450 ctx.retire(result);
451 }
452 Plan::AlterClusterReplicaRename(plan) => {
453 let result = self
454 .sequence_alter_cluster_replica_rename(ctx.session(), plan)
455 .await;
456 ctx.retire(result);
457 }
458 Plan::AlterConnection(plan) => {
459 self.sequence_alter_connection(ctx, plan).await;
460 }
461 Plan::AlterSetCluster(plan) => {
462 let result = self.sequence_alter_set_cluster(ctx.session(), plan).await;
463 ctx.retire(result);
464 }
465 Plan::AlterRetainHistory(plan) => {
466 let result = self.sequence_alter_retain_history(&mut ctx, plan).await;
467 ctx.retire(result);
468 }
469 Plan::AlterSourceTimestampInterval(plan) => {
470 let result = self
471 .sequence_alter_source_timestamp_interval(&mut ctx, plan)
472 .await;
473 ctx.retire(result);
474 }
475 Plan::AlterItemRename(plan) => {
476 let result = self.sequence_alter_item_rename(&mut ctx, plan).await;
477 ctx.retire(result);
478 }
479 Plan::AlterSchemaRename(plan) => {
480 let result = self.sequence_alter_schema_rename(&mut ctx, plan).await;
481 ctx.retire(result);
482 }
483 Plan::AlterSchemaSwap(plan) => {
484 let result = self.sequence_alter_schema_swap(&mut ctx, plan).await;
485 ctx.retire(result);
486 }
487 Plan::AlterRole(plan) => {
488 let result = self.sequence_alter_role(ctx.session_mut(), plan).await;
489 ctx.retire(result);
490 }
491 Plan::AlterSecret(plan) => {
492 self.sequence_alter_secret(ctx, plan).await;
493 }
494 Plan::AlterSink(plan) => {
495 self.sequence_alter_sink_prepare(ctx, plan).await;
496 }
497 Plan::AlterSource(plan) => {
498 let result = self.sequence_alter_source(ctx.session_mut(), plan).await;
499 ctx.retire(result);
500 }
501 Plan::AlterSystemSet(plan) => {
502 let result = self.sequence_alter_system_set(ctx.session(), plan).await;
503 ctx.retire(result);
504 }
505 Plan::AlterSystemReset(plan) => {
506 let result = self.sequence_alter_system_reset(ctx.session(), plan).await;
507 ctx.retire(result);
508 }
509 Plan::AlterSystemResetAll(plan) => {
510 let result = self
511 .sequence_alter_system_reset_all(ctx.session(), plan)
512 .await;
513 ctx.retire(result);
514 }
515 Plan::AlterTableAddColumn(plan) => {
516 let result = self.sequence_alter_table(&mut ctx, plan).await;
517 ctx.retire(result);
518 }
519 Plan::AlterMaterializedViewApplyReplacement(plan) => {
520 self.sequence_alter_materialized_view_apply_replacement_prepare(ctx, plan)
521 .await;
522 }
523 Plan::AlterNetworkPolicy(plan) => {
524 let res = self
525 .sequence_alter_network_policy(ctx.session(), plan)
526 .await;
527 ctx.retire(res);
528 }
529 Plan::DiscardTemp => {
530 self.drop_temp_items(ctx.session().conn_id()).await;
531 ctx.retire(Ok(ExecuteResponse::DiscardedTemp));
532 }
533 Plan::DiscardAll => {
534 let txn = ctx.session().transaction();
537 let discardable =
538 matches!(txn, TransactionStatus::Started(_)) && !txn.contains_ops();
539 let ret = if discardable {
540 let (_, retire_notify) = self.clear_transaction(ctx.session_mut()).await;
541 ctx.delay_response_until(retire_notify);
542 self.drop_temp_items(ctx.session().conn_id()).await;
543 let params = ctx.session_mut().reset();
549 Ok(ExecuteResponse::DiscardedAll { params })
550 } else {
551 Err(AdapterError::OperationProhibitsTransaction(
552 "DISCARD ALL".into(),
553 ))
554 };
555 ctx.retire(ret);
556 }
557 Plan::Declare(plan) => {
558 self.declare(ctx, plan.name, plan.stmt, plan.sql, plan.params);
559 }
560 Plan::Fetch(FetchPlan {
561 name,
562 count,
563 timeout,
564 }) => {
565 let ctx_extra = std::mem::take(ctx.extra_mut());
566 ctx.retire(Ok(ExecuteResponse::Fetch {
567 name,
568 count,
569 timeout,
570 ctx_extra,
571 }));
572 }
573 Plan::Close(plan) => {
574 if ctx.session_mut().remove_portal(&plan.name) {
575 ctx.retire(Ok(ExecuteResponse::ClosedCursor));
576 } else {
577 ctx.retire(Err(AdapterError::UnknownCursor(plan.name)));
578 }
579 }
580 Plan::Prepare(plan) => {
581 if ctx
582 .session()
583 .get_prepared_statement_unverified(&plan.name)
584 .is_some()
585 {
586 ctx.retire(Err(AdapterError::PreparedStatementExists(plan.name)));
587 } else {
588 let state_revision = StateRevision {
589 catalog_revision: self.catalog().transient_revision(),
590 session_state_revision: ctx.session().state_revision(),
591 };
592 ctx.session_mut().set_prepared_statement(
593 plan.name,
594 Some(plan.stmt),
595 plan.sql,
596 plan.desc,
597 state_revision,
598 self.now(),
599 );
600 ctx.retire(Ok(ExecuteResponse::Prepare));
601 }
602 }
603 Plan::Execute(plan) => {
604 match self.sequence_execute(ctx.session_mut(), plan) {
605 Ok(portal_name) => {
606 let (tx, _, session, extra, response_barriers) = ctx.into_parts();
607 let command = Message::Command(
614 OpenTelemetryContext::obtain(),
615 Command::Execute {
616 portal_name,
617 session,
618 tx: tx.take(),
619 outer_ctx_extra: Some(extra.defuse()),
620 },
621 );
622 if response_barriers.is_empty() {
623 self.internal_cmd_tx
624 .send(command)
625 .expect("sending to self.internal_cmd_tx cannot fail");
626 } else {
627 let internal_cmd_tx = self.internal_cmd_tx.clone();
628 mz_ore::task::spawn(
629 || "execute_after_response_barriers",
630 async move {
631 for barrier in response_barriers {
632 barrier.await;
633 }
634 let _ = internal_cmd_tx.send(command);
635 },
636 );
637 }
638 }
639 Err(err) => ctx.retire(Err(err)),
640 };
641 }
642 Plan::Deallocate(plan) => match plan.name {
643 Some(name) => {
644 if ctx.session_mut().remove_prepared_statement(&name) {
645 ctx.retire(Ok(ExecuteResponse::Deallocate { all: false }));
646 } else {
647 ctx.retire(Err(AdapterError::UnknownPreparedStatement(name)));
648 }
649 }
650 None => {
651 ctx.session_mut().remove_all_prepared_statements();
652 ctx.retire(Ok(ExecuteResponse::Deallocate { all: true }));
653 }
654 },
655 Plan::Raise(RaisePlan { severity }) => {
656 ctx.session()
657 .add_notice(AdapterNotice::UserRequested { severity });
658 ctx.retire(Ok(ExecuteResponse::Raised));
659 }
660 Plan::GrantPrivileges(plan) => {
661 let result = self
662 .sequence_grant_privileges(ctx.session_mut(), plan)
663 .await;
664 ctx.retire(result);
665 }
666 Plan::RevokePrivileges(plan) => {
667 let result = self
668 .sequence_revoke_privileges(ctx.session_mut(), plan)
669 .await;
670 ctx.retire(result);
671 }
672 Plan::AlterDefaultPrivileges(plan) => {
673 let result = self
674 .sequence_alter_default_privileges(ctx.session_mut(), plan)
675 .await;
676 ctx.retire(result);
677 }
678 Plan::GrantRole(plan) => {
679 let result = self.sequence_grant_role(ctx.session_mut(), plan).await;
680 ctx.retire(result);
681 }
682 Plan::RevokeRole(plan) => {
683 let result = self.sequence_revoke_role(ctx.session_mut(), plan).await;
684 ctx.retire(result);
685 }
686 Plan::AlterOwner(plan) => {
687 let result = self.sequence_alter_owner(ctx.session_mut(), plan).await;
688 ctx.retire(result);
689 }
690 Plan::ReassignOwned(plan) => {
691 let result = self.sequence_reassign_owned(ctx.session_mut(), plan).await;
692 ctx.retire(result);
693 }
694 Plan::ValidateConnection(plan) => {
695 let connection = plan
696 .connection
697 .into_inline_connection(self.catalog().state());
698 let current_storage_configuration = self.controller.storage.config().clone();
699 mz_ore::task::spawn(|| "coord::validate_connection", async move {
700 let res = match connection
701 .validate(plan.id, ¤t_storage_configuration)
702 .await
703 {
704 Ok(()) => Ok(ExecuteResponse::ValidatedConnection),
705 Err(err) => Err(err.into()),
706 };
707 ctx.retire(res);
708 });
709 }
710 }
711 }
712 .instrument(tracing::debug_span!("coord::sequencer::sequence_plan"))
713 .boxed_local()
714 }
715
716 #[mz_ore::instrument(level = "debug")]
717 pub(crate) async fn sequence_execute_single_statement_transaction(
718 &mut self,
719 ctx: ExecuteContext,
720 stmt: Arc<Statement<Raw>>,
721 params: Params,
722 ) {
723 let (tx, internal_cmd_tx, mut session, extra, response_barriers) = ctx.into_parts();
725 assert!(matches!(session.transaction(), TransactionStatus::Default));
726 session.start_transaction_single_stmt(self.now_datetime());
727 let conn_id = session.conn_id().unhandled();
728
729 let (sub_tx, sub_rx) = oneshot::channel();
731 let sub_ct = ClientTransmitter::new(sub_tx, self.internal_cmd_tx.clone());
732 let sub_ctx = ExecuteContext::from_parts_with_response_barriers(
733 sub_ct,
734 internal_cmd_tx,
735 session,
736 extra,
737 response_barriers,
738 );
739 self.handle_execute_inner(stmt, params, sub_ctx).await;
740
741 let internal_cmd_tx = self.internal_cmd_tx.clone();
744 mz_ore::task::spawn(
745 || format!("execute_single_statement:{conn_id}"),
746 async move {
747 let Ok(Response {
748 result,
749 session,
750 otel_ctx,
751 }) = sub_rx.await
752 else {
753 return;
755 };
756 otel_ctx.attach_as_parent();
757 let (sub_tx, sub_rx) = oneshot::channel();
758 let _ = internal_cmd_tx.send(Message::Command(
759 otel_ctx,
760 Command::Commit {
761 action: EndTransactionAction::Commit,
762 session,
763 tx: sub_tx,
764 },
765 ));
766 let Ok(commit_response) = sub_rx.await else {
767 return;
769 };
770 assert!(matches!(
771 commit_response.session.transaction(),
772 TransactionStatus::Default
773 ));
774 let result = match (result, commit_response.result) {
780 (Ok(_), commit) => commit,
781 (Err(result), _) => Err(result),
782 };
783 tx.send(result, commit_response.session);
786 }
787 .instrument(Span::current()),
788 );
789 }
790
791 #[mz_ore::instrument(level = "debug")]
795 pub(crate) async fn sequence_create_role_for_startup(
796 &mut self,
797 plan: CreateRolePlan,
798 ) -> Result<ExecuteResponse, AdapterError> {
799 self.sequence_create_role(None, plan).await
806 }
807
808 pub(crate) fn allocate_transient_id(&self) -> (CatalogItemId, GlobalId) {
809 self.transient_id_gen.allocate_id()
810 }
811
812 fn should_emit_rbac_notice(&self, session: &Session) -> Option<AdapterNotice> {
813 if !rbac::is_rbac_enabled_for_session(self.catalog.system_config(), session) {
814 Some(AdapterNotice::RbacUserDisabled)
815 } else {
816 None
817 }
818 }
819
820 pub(crate) fn insert_constant(
827 catalog: &Catalog,
828 session: &mut Session,
829 target_id: CatalogItemId,
830 constants: MirRelationExpr,
831 ) -> Result<ExecuteResponse, AdapterError> {
832 let desc = match catalog.try_get_entry(&target_id) {
834 Some(table) => {
835 table.relation_desc_latest().expect("table has desc")
837 }
838 None => {
839 return Err(AdapterError::Catalog(mz_catalog::memory::error::Error {
840 kind: mz_catalog::memory::error::ErrorKind::Sql(CatalogError::UnknownItem(
841 target_id.to_string(),
842 )),
843 }));
844 }
845 };
846
847 match constants.as_const() {
848 Some((rows, ..)) => {
849 let rows = rows.clone()?;
850 for (row, _) in &rows {
851 for (i, datum) in row.iter().enumerate() {
852 desc.constraints_met(i, &datum)?;
853 }
854 }
855 let diffs_plan = plan::SendDiffsPlan {
856 id: target_id,
857 updates: rows,
858 kind: MutationKind::Insert,
859 returning: Vec::new(),
860 max_result_size: catalog.system_config().max_result_size(),
861 };
862 let result = Self::send_diffs(session, diffs_plan);
863 fail::fail_point!("insert_after_pack_before_commit");
866 result
867 }
868 None => panic!(
869 "tried using sequence_insert_constant on non-constant MirRelationExpr\n{}",
870 constants.pretty(),
871 ),
872 }
873 }
874
875 #[mz_ore::instrument(level = "debug")]
876 pub(crate) fn send_diffs(
877 session: &mut Session,
878 mut plan: plan::SendDiffsPlan,
879 ) -> Result<ExecuteResponse, AdapterError> {
880 let affected_rows = {
881 let mut affected_rows = Diff::from(0);
882 let mut all_positive_diffs = true;
883 for (_, diff) in plan.updates.iter() {
886 if diff.is_negative() {
887 all_positive_diffs = false;
888 break;
889 }
890
891 affected_rows += diff;
892 }
893
894 if !all_positive_diffs {
895 differential_dataflow::consolidation::consolidate(&mut plan.updates);
907
908 affected_rows = Diff::ZERO;
909 for (_, diff) in plan.updates.iter() {
914 affected_rows += diff.abs();
915 }
916 }
917
918 usize::try_from(affected_rows.into_inner()).expect("positive Diff must fit")
919 };
920 event!(
921 Level::TRACE,
922 affected_rows,
923 id = format!("{:?}", plan.id),
924 kind = format!("{:?}", plan.kind),
925 updates = plan.updates.len(),
926 returning = plan.returning.len(),
927 );
928
929 session.add_transaction_ops(TransactionOps::Writes(vec![WriteOp {
930 id: plan.id,
931 rows: TableData::Rows(plan.updates),
932 }]))?;
933 if !plan.returning.is_empty() {
934 let finishing = RowSetFinishing {
935 order_by: Vec::new(),
936 limit: None,
937 offset: 0,
938 project: (0..plan.returning[0].0.iter().count()).collect(),
939 };
940 let max_returned_query_size = session.vars().max_query_result_size();
941 let duration_histogram = session.metrics().row_set_finishing_seconds();
942
943 return match finishing.finish(
944 RowCollection::new(plan.returning, &finishing.order_by),
945 plan.max_result_size,
946 Some(max_returned_query_size),
947 duration_histogram,
948 ) {
949 Ok((rows, _size_bytes)) => Ok(Self::send_immediate_rows(rows)),
950 Err(e) => Err(AdapterError::ResultSize(e)),
951 };
952 }
953 Ok(match plan.kind {
954 MutationKind::Delete => ExecuteResponse::Deleted(affected_rows),
955 MutationKind::Insert => ExecuteResponse::Inserted(affected_rows),
956 MutationKind::Update => ExecuteResponse::Updated(affected_rows / 2),
957 })
958 }
959}
960
961pub(crate) fn check_log_reads(
970 catalog: &Catalog,
971 cluster: &Cluster,
972 source_ids: &BTreeSet<GlobalId>,
973 target_replica: &mut Option<ReplicaId>,
974 vars: &SessionVars,
975) -> Result<impl IntoIterator<Item = AdapterNotice>, AdapterError>
976where
977{
978 let log_names = source_ids
979 .iter()
980 .map(|gid| catalog.resolve_item_id(gid))
981 .flat_map(|item_id| catalog.introspection_dependencies(item_id))
982 .map(|item_id| catalog.get_entry(&item_id).name().item.clone())
983 .collect::<Vec<_>>();
984
985 if log_names.is_empty() {
986 return Ok(None);
987 }
988
989 let num_replicas = cluster.replicas().count();
993 if target_replica.is_none() {
994 if num_replicas == 1 {
995 *target_replica = cluster.replicas().map(|r| r.replica_id).next();
996 } else {
997 return Err(AdapterError::UntargetedLogRead { log_names });
998 }
999 }
1000
1001 let replica_id = target_replica.expect("set to `Some` above");
1004 let replica = &cluster.replica(replica_id).expect("Replica must exist");
1005 if !replica.config.compute.logging.enabled() {
1006 return Err(AdapterError::IntrospectionDisabled { log_names });
1007 }
1008
1009 Ok(vars
1010 .emit_introspection_query_notice()
1011 .then_some(AdapterNotice::PerReplicaLogRead { log_names }))
1012}
1013
1014pub(crate) fn emit_optimizer_notices(
1016 catalog: &Catalog,
1017 session: &Session,
1018 notices: &[RawOptimizerNotice],
1019) {
1020 if notices.is_empty() {
1022 return;
1023 }
1024 let humanizer = catalog.for_session(session);
1025 let system_vars = catalog.system_config();
1026 for notice in notices {
1027 let kind = OptimizerNoticeKind::from(notice);
1028 let notice_enabled = match kind {
1029 OptimizerNoticeKind::EqualsNull => system_vars.enable_notices_for_equals_null(),
1030 OptimizerNoticeKind::IndexAlreadyExists => {
1031 system_vars.enable_notices_for_index_already_exists()
1032 }
1033 OptimizerNoticeKind::IndexTooWideForLiteralConstraints => {
1034 system_vars.enable_notices_for_index_too_wide_for_literal_constraints()
1035 }
1036 OptimizerNoticeKind::IndexKeyEmpty => system_vars.enable_notices_for_index_empty_key(),
1037 };
1038 if notice_enabled {
1039 session.add_notice(AdapterNotice::OptimizerNotice {
1043 notice: notice.message(&humanizer, false).to_string(),
1044 hint: notice.hint(&humanizer, false).to_string(),
1045 });
1046 }
1047 session
1048 .metrics()
1049 .optimization_notices(&[kind.metric_label()])
1050 .inc_by(1);
1051 }
1052}
1053
1054pub fn eval_copy_to_uri(
1056 to: HirScalarExpr,
1057 session: &Session,
1058 catalog_state: &CatalogState,
1059) -> Result<Uri, AdapterError> {
1060 let style = ExprPrepOneShot {
1061 logical_time: EvalTime::NotAvailable,
1062 session,
1063 catalog_state,
1064 };
1065 let mut to = to.lower_uncorrelated(catalog_state.system_config())?;
1066 style.prep_scalar_expr(&mut to)?;
1067 let temp_storage = RowArena::new();
1068 let evaled = to.eval(&[], &temp_storage)?;
1069 if evaled == Datum::Null {
1070 coord_bail!("COPY TO target value can not be null");
1071 }
1072 let to_url = match Uri::from_str(evaled.unwrap_str()) {
1073 Ok(url) => {
1074 if url.scheme_str() != Some("s3") && url.scheme_str() != Some("gs") {
1075 coord_bail!("only 's3://...' and 'gs://...' urls are supported as COPY TO target");
1076 }
1077 url
1078 }
1079 Err(e) => coord_bail!("could not parse COPY TO target url: {}", e),
1080 };
1081 Ok(to_url)
1082}
1083
1084pub(crate) async fn explain_pushdown_future_inner<
1091 I: IntoIterator<Item = (GlobalId, MapFilterProject)>,
1092>(
1093 session: &Session,
1094 catalog: &Catalog,
1095 storage_collections: &Arc<dyn StorageCollections + Send + Sync>,
1096 as_of: Antichain<Timestamp>,
1097 mz_now: ResultSpec<'static>,
1098 imports: I,
1099) -> impl Future<Output = Result<ExecuteResponse, AdapterError>> + use<I> {
1100 let mut explain_timeout = *session.vars().statement_timeout();
1101 if explain_timeout == Duration::ZERO {
1103 explain_timeout = Duration::MAX;
1104 }
1105 let mut futures = FuturesOrdered::new();
1106 for (id, mfp) in imports {
1107 let catalog_entry = catalog.get_entry_by_global_id(&id);
1108 let full_name = catalog
1109 .for_session(session)
1110 .resolve_full_name(&catalog_entry.name);
1111 let name = format!("{}", full_name);
1112 let relation_desc = catalog_entry
1113 .relation_desc()
1114 .expect("source should have a proper desc")
1115 .into_owned();
1116 let stats_future = storage_collections
1117 .snapshot_parts_stats(id, as_of.clone())
1118 .await;
1119
1120 let mz_now = mz_now.clone();
1121 futures.push_back(async move {
1126 let snapshot_stats = match stats_future.await {
1127 Ok(stats) => stats,
1128 Err(e) => return Err(e),
1129 };
1130 let mut total_bytes = 0;
1131 let mut total_parts = 0;
1132 let mut selected_bytes = 0;
1133 let mut selected_parts = 0;
1134 for SnapshotPartStats {
1135 encoded_size_bytes: bytes,
1136 stats,
1137 } in &snapshot_stats.parts
1138 {
1139 let bytes = u64::cast_from(*bytes);
1140 total_bytes += bytes;
1141 total_parts += 1u64;
1142 let selected = match stats.as_ref().and_then(|x| x.try_decode().ok()) {
1143 None => true,
1147 Some(stats) => {
1148 let stats = RelationPartStats::new(
1149 name.as_str(),
1150 &snapshot_stats.metrics.pushdown.part_stats,
1151 &relation_desc,
1152 &stats,
1153 );
1154 stats.may_match_mfp(mz_now.clone(), &mfp)
1155 }
1156 };
1157
1158 if selected {
1159 selected_bytes += bytes;
1160 selected_parts += 1u64;
1161 }
1162 }
1163 Ok(Row::pack_slice(&[
1164 name.as_str().into(),
1165 total_bytes.into(),
1166 selected_bytes.into(),
1167 total_parts.into(),
1168 selected_parts.into(),
1169 ]))
1170 });
1171 }
1172
1173 let fut = async move {
1174 match tokio::time::timeout(
1175 explain_timeout,
1176 futures::TryStreamExt::try_collect::<Vec<_>>(futures),
1177 )
1178 .await
1179 {
1180 Ok(Ok(rows)) => Ok(ExecuteResponse::SendingRowsImmediate {
1181 rows: Box::new(rows.into_row_iter()),
1182 }),
1183 Ok(Err(err)) => Err(err.into()),
1184 Err(_) => Err(AdapterError::StatementTimeout),
1185 }
1186 };
1187 fut
1188}
1189
1190pub(crate) async fn explain_plan_inner(
1193 session: &Session,
1194 catalog: &Catalog,
1195 df_meta: DataflowMetainfo,
1196 explain_ctx: ExplainPlanContext,
1197 optimizer: peek::Optimizer,
1198 insights_ctx: Option<Box<PlanInsightsContext>>,
1199) -> Result<Vec<Row>, AdapterError> {
1200 let ExplainPlanContext {
1201 config,
1202 format,
1203 stage,
1204 desc,
1205 optimizer_trace,
1206 ..
1207 } = explain_ctx;
1208
1209 let desc = desc.expect("RelationDesc for SelectPlan in EXPLAIN mode");
1210
1211 let session_catalog = catalog.for_session(session);
1212 let expr_humanizer = {
1213 let transient_items = btreemap! {
1214 optimizer.select_id() => TransientItem::new(
1215 Some(vec![GlobalId::Explain.to_string()]),
1216 Some(desc.iter_names().map(|c| c.to_string()).collect()),
1217 )
1218 };
1219 ExprHumanizerExt::new(transient_items, &session_catalog)
1220 };
1221
1222 let finishing = if optimizer.finishing().is_trivial(desc.arity()) {
1223 None
1224 } else {
1225 Some(optimizer.finishing().clone())
1226 };
1227
1228 let target_cluster = catalog.get_cluster(optimizer.cluster_id());
1229 let features = optimizer.config().features.clone();
1230
1231 let rows = optimizer_trace
1232 .into_rows(
1233 format,
1234 &config,
1235 &features,
1236 &expr_humanizer,
1237 finishing,
1238 Some(target_cluster),
1239 df_meta,
1240 stage,
1241 plan::ExplaineeStatementKind::Select,
1242 insights_ctx,
1243 )
1244 .await?;
1245
1246 Ok(rows)
1247}
1248
1249pub(crate) async fn statistics_oracle(
1254 session: &Session,
1255 source_ids: &BTreeSet<GlobalId>,
1256 query_as_of: &Antichain<Timestamp>,
1257 is_oneshot: bool,
1258 system_config: &vars::SystemVars,
1259 storage_collections: &dyn StorageCollections,
1260) -> Result<Box<dyn StatisticsOracle>, AdapterError> {
1261 if !session.vars().enable_session_cardinality_estimates() {
1262 let stats: Box<dyn StatisticsOracle> = Box::new(EmptyStatisticsOracle);
1263 return Ok(stats);
1264 }
1265
1266 let timeout = if is_oneshot {
1267 system_config.optimizer_oneshot_stats_timeout()
1269 } else {
1270 system_config.optimizer_stats_timeout()
1271 };
1272
1273 let cached_stats = mz_ore::future::timeout(
1274 timeout,
1275 CachedStatisticsOracle::new(source_ids, query_as_of, storage_collections),
1276 )
1277 .await;
1278
1279 match cached_stats {
1280 Ok(stats) => Ok(Box::new(stats)),
1281 Err(mz_ore::future::TimeoutError::DeadlineElapsed) => {
1282 warn!(
1283 is_oneshot = is_oneshot,
1284 "optimizer statistics collection timed out after {}ms",
1285 timeout.as_millis()
1286 );
1287
1288 Ok(Box::new(EmptyStatisticsOracle))
1289 }
1290 Err(mz_ore::future::TimeoutError::Inner(e)) => Err(AdapterError::Storage(e)),
1291 }
1292}
1293
1294#[derive(Debug)]
1295struct CachedStatisticsOracle {
1296 cache: BTreeMap<GlobalId, usize>,
1297}
1298
1299impl CachedStatisticsOracle {
1300 pub async fn new(
1301 ids: &BTreeSet<GlobalId>,
1302 as_of: &Antichain<Timestamp>,
1303 storage_collections: &dyn StorageCollections,
1304 ) -> Result<Self, StorageError> {
1305 let mut cache = BTreeMap::new();
1306
1307 for id in ids {
1308 let stats = storage_collections.snapshot_stats(*id, as_of.clone()).await;
1309
1310 match stats {
1311 Ok(stats) => {
1312 cache.insert(*id, stats.num_updates);
1313 }
1314 Err(StorageError::IdentifierMissing(id)) => {
1315 ::tracing::debug!("no statistics for {id}")
1316 }
1317 Err(e) => return Err(e),
1318 }
1319 }
1320
1321 Ok(Self { cache })
1322 }
1323}
1324
1325impl StatisticsOracle for CachedStatisticsOracle {
1326 fn cardinality_estimate(&self, id: GlobalId) -> Option<usize> {
1327 self.cache.get(&id).map(|estimate| *estimate)
1328 }
1329
1330 fn as_map(&self) -> BTreeMap<GlobalId, usize> {
1331 self.cache.clone()
1332 }
1333}