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