Skip to main content

mz_adapter/coord/
sequencer.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
10// Prevents anyone from accidentally exporting a method from the `inner` module.
11#![allow(clippy::pub_use)]
12
13//! Logic for executing a planned SQL query.
14
15use std::sync::Arc;
16use std::time::Duration;
17
18use futures::future::LocalBoxFuture;
19use futures::stream::FuturesOrdered;
20use inner::return_if_err;
21use mz_compute_types::plan::scalar::LirScalarExpr;
22use mz_expr::row::RowCollection;
23use mz_expr::{MfpPlan, MirRelationExpr, ResultSpec, RowSetFinishing};
24use mz_ore::cast::CastFrom;
25use mz_ore::soft_panic_or_log;
26use mz_persist_client::stats::SnapshotPartStats;
27use mz_repr::{CatalogItemId, Diff, GlobalId, IntoRowIterator, Row, Timestamp};
28use mz_sql::catalog::{CatalogError, SessionCatalog};
29use mz_sql::names::ResolvedIds;
30use mz_sql::plan::{
31    self, AbortTransactionPlan, CommitTransactionPlan, CreateRolePlan, CreateSourcePlanBundle,
32    FetchPlan, MutationKind, Params, Plan, PlanKind, RaisePlan,
33};
34use mz_sql::rbac;
35use mz_sql::session::metadata::SessionMetadata;
36use mz_sql_parser::ast::{Raw, Statement};
37use mz_storage_client::client::TableData;
38use mz_storage_client::storage_collections::StorageCollections;
39use mz_storage_types::connections::inline::IntoInlineConnection;
40use mz_storage_types::stats::RelationPartStats;
41use mz_transform::notice::{OptimizerNoticeApi, OptimizerNoticeKind, RawOptimizerNotice};
42use timely::progress::Antichain;
43use tokio::sync::oneshot;
44use tracing::{Instrument, Level, Span, event, warn};
45
46use crate::ExecuteContext;
47use crate::catalog::Catalog;
48use crate::command::{Command, ExecuteResponse, Response};
49use crate::coord::{Coordinator, Message, TargetCluster, catalog_serving};
50use crate::error::AdapterError;
51use crate::notice::AdapterNotice;
52use crate::session::{
53    EndTransactionAction, Session, StateRevision, TransactionOps, TransactionStatus, WriteOp,
54};
55use crate::util::ClientTransmitter;
56
57// DO NOT make this visible in any way, i.e. do not add any version of
58// `pub` to this mod. The inner `sequence_X` methods are hidden in this
59// private module to prevent anyone from calling them directly. Bound-plan
60// sequencing should be done through the `sequence_plan` method.
61// This allows us to add catch-all logic that should be applied to all
62// plans in `sequence_plan` and guarantee that no caller can circumvent
63// that logic.
64//
65// The exceptions are:
66//
67// - Rename/swap commands enter through `sequence_rename_command` and execute in the catalog.
68// - Creating a role during connection startup. In this scenario, the session has not been properly
69// initialized and we need to skip directly to creating role. We have a specific method,
70// `sequence_create_role_for_startup` for this purpose.
71// - Methods that continue the execution of some plan that was being run asynchronously, such as
72// `sequence_staged` and `sequence_create_connection_stage_finish`.
73// - The frontend sequencing paths call `emit_optimizer_notices` and
74//   `explain_pushdown_future_inner`, which the coordinator uses as well.
75
76mod inner;
77
78impl Coordinator {
79    /// Selects and records the execution cluster, checking its restrictions.
80    pub(super) fn prepare_plan_execution(
81        &mut self,
82        ctx: &ExecuteContext,
83        plan: &Plan,
84    ) -> Result<(TargetCluster, Option<mz_controller_types::ClusterId>), AdapterError> {
85        // Scope the borrow of the Catalog because we need to mutate the Coordinator state below.
86        let target_cluster = match ctx.session().transaction().cluster() {
87            // Use the current transaction's cluster.
88            Some(cluster_id) => TargetCluster::Transaction(cluster_id),
89            // If there isn't a current cluster set for a transaction, then try to auto route.
90            None => {
91                let session_catalog = self.catalog.for_session(ctx.session());
92                catalog_serving::auto_run_on_catalog_server(&session_catalog, ctx.session(), plan)
93            }
94        };
95        let (target_cluster_id, target_cluster_name) = match self
96            .catalog()
97            .resolve_target_cluster(target_cluster, ctx.session())
98        {
99            Ok(cluster) => (Some(cluster.id), Some(cluster.name.clone())),
100            Err(_) => (None, None),
101        };
102
103        if let (Some(cluster_id), Some(cluster_name), Some(statement_id)) = (
104            target_cluster_id,
105            target_cluster_name.clone(),
106            ctx.extra().contents(),
107        ) {
108            self.set_statement_execution_cluster(statement_id, cluster_id, cluster_name);
109        }
110
111        let session_catalog = self.catalog.for_session(ctx.session());
112
113        if let Some(cluster_name) = &target_cluster_name {
114            if let Err(e) =
115                catalog_serving::check_cluster_restrictions(cluster_name, &session_catalog, plan)
116            {
117                return Err(e);
118            }
119        }
120
121        Ok((target_cluster, target_cluster_id))
122    }
123
124    /// Dispatches a plan in a size-bounded, heap-allocated future.
125    pub(crate) fn sequence_plan(
126        &mut self,
127        mut ctx: ExecuteContext,
128        plan: Plan,
129        resolved_ids: ResolvedIds,
130        sql_impl_resolved_ids: ResolvedIds,
131    ) -> LocalBoxFuture<'_, ()> {
132        let future = async move {
133            let responses = ExecuteResponse::generated_from(&PlanKind::from(&plan));
134            ctx.tx_mut().set_allowed(responses);
135
136            if self.controller.read_only() && !plan.allowed_in_read_only() {
137                ctx.retire(Err(AdapterError::ReadOnly));
138                return;
139            }
140
141            let (target_cluster, target_cluster_id) =
142                return_if_err!(self.prepare_plan_execution(&ctx, &plan), ctx);
143            let session_catalog = self.catalog.for_session(ctx.session());
144
145            // The target-connection role only matters for `pg_cancel_backend`,
146            // which the frontend sequences.
147            if let Err(e) = rbac::check_plan(
148                &session_catalog,
149                None,
150                ctx.session(),
151                &plan,
152                target_cluster_id,
153                &resolved_ids,
154                &sql_impl_resolved_ids,
155            ) {
156                return ctx.retire(Err(e.into()));
157            }
158
159            match plan {
160                Plan::CreateSource(plan) => {
161                    if !matches!(
162                        plan.source.data_source,
163                        plan::DataSourceDesc::Webhook { .. }
164                    ) {
165                        let msg = "CREATE SOURCE must execute as a catalog command";
166                        soft_panic_or_log!("{msg}");
167                        ctx.retire(Err(AdapterError::Internal(msg.into())));
168                        return;
169                    }
170                    let (item_id, global_id) = return_if_err!(self.allocate_user_id().await, ctx);
171                    let result = self
172                        .sequence_create_source(
173                            &mut ctx,
174                            vec![CreateSourcePlanBundle {
175                                item_id,
176                                global_id,
177                                plan,
178                                resolved_ids,
179                                available_source_references: None,
180                            }],
181                        )
182                        .await;
183                    ctx.retire(result);
184                }
185                Plan::CreateSources(_) => {
186                    let msg = "CREATE SOURCE must execute as a catalog command";
187                    soft_panic_or_log!("{msg}");
188                    ctx.retire(Err(AdapterError::Internal(msg.into())));
189                }
190                Plan::CreateConnection(plan) => {
191                    self.sequence_create_connection(ctx, plan, resolved_ids)
192                        .await;
193                }
194                Plan::CreateDatabase(plan) => {
195                    let result = self.sequence_create_database(ctx.session_mut(), plan).await;
196                    ctx.retire(result);
197                }
198                Plan::CreateSchema(plan) => {
199                    let result = self.sequence_create_schema(ctx.session_mut(), plan).await;
200                    ctx.retire(result);
201                }
202                Plan::CreateRole(plan) => {
203                    let result = self
204                        .sequence_create_role(Some(ctx.session().conn_id()), plan)
205                        .await;
206                    if let Some(notice) = self.should_emit_rbac_notice(ctx.session()) {
207                        ctx.session().add_notice(notice);
208                    }
209                    ctx.retire(result);
210                }
211                Plan::CreateCluster(plan) => {
212                    let result = self.sequence_create_cluster(ctx.session(), plan).await;
213                    ctx.retire(result);
214                }
215                Plan::CreateClusterReplica(plan) => {
216                    let result = self
217                        .sequence_create_cluster_replica(ctx.session(), plan)
218                        .await;
219                    ctx.retire(result);
220                }
221                Plan::CreateTable(plan) => {
222                    if matches!(
223                        plan.table.data_source,
224                        plan::TableDataSource::DataSource {
225                            desc: plan::DataSourceDesc::IngestionExport { .. },
226                            ..
227                        }
228                    ) {
229                        let msg = "CREATE TABLE FROM SOURCE must execute as a catalog command";
230                        soft_panic_or_log!("{msg}");
231                        ctx.retire(Err(AdapterError::Internal(msg.into())));
232                        return;
233                    }
234                    let result = self
235                        .sequence_create_table(&mut ctx, plan, resolved_ids)
236                        .await;
237                    ctx.retire(result);
238                }
239                Plan::CreateSecret(plan) => {
240                    self.sequence_create_secret(ctx, plan).await;
241                }
242                Plan::CreateSink(plan) => {
243                    self.sequence_create_sink(ctx, plan, resolved_ids).await;
244                }
245                Plan::CreateView(plan) => {
246                    self.sequence_create_view(ctx, plan, resolved_ids).await;
247                }
248                Plan::CreateMaterializedView(plan) => {
249                    self.sequence_create_materialized_view(ctx, plan, resolved_ids)
250                        .await;
251                }
252                Plan::CreateIndex(plan) => {
253                    self.sequence_create_index(ctx, plan, resolved_ids).await;
254                }
255                Plan::CreateMetricSink(plan) => {
256                    self.sequence_create_metric_sink(ctx, plan, resolved_ids)
257                        .await;
258                }
259                Plan::CreateType(plan) => {
260                    let result = self
261                        .sequence_create_type(ctx.session(), plan, resolved_ids)
262                        .await;
263                    ctx.retire(result);
264                }
265                Plan::CreateNetworkPolicy(plan) => {
266                    let res = self
267                        .sequence_create_network_policy(ctx.session(), plan)
268                        .await;
269                    ctx.retire(res);
270                }
271                Plan::Comment(plan) => {
272                    let result = self.sequence_comment_on(ctx.session(), plan).await;
273                    ctx.retire(result);
274                }
275                Plan::DropObjects(plan) => {
276                    let result = self.sequence_drop_objects(&mut ctx, plan).await;
277                    ctx.retire(result);
278                }
279                Plan::DropOwned(plan) => {
280                    let result = self.sequence_drop_owned(ctx.session_mut(), plan).await;
281                    ctx.retire(result);
282                }
283                Plan::EmptyQuery => {
284                    ctx.retire(Ok(ExecuteResponse::EmptyQuery));
285                }
286                Plan::ShowAllVariables => {
287                    let result = self.sequence_show_all_variables(ctx.session());
288                    ctx.retire(result);
289                }
290                Plan::ShowVariable(plan) => {
291                    let result = self.sequence_show_variable(ctx.session(), plan);
292                    ctx.retire(result);
293                }
294                Plan::InspectShard(plan) => {
295                    // TODO: Ideally, this await would happen off the main thread.
296                    let result = self.sequence_inspect_shard(ctx.session(), plan).await;
297                    ctx.retire(result);
298                }
299                Plan::SetVariable(plan) => {
300                    let result = self.sequence_set_variable(ctx.session_mut(), plan);
301                    ctx.retire(result);
302                }
303                Plan::ResetVariable(plan) => {
304                    let result = self.sequence_reset_variable(ctx.session_mut(), plan);
305                    ctx.retire(result);
306                }
307                Plan::SetTransaction(plan) => {
308                    let result = self.sequence_set_transaction(ctx.session_mut(), plan);
309                    ctx.retire(result);
310                }
311                Plan::StartTransaction(plan) => {
312                    if matches!(
313                        ctx.session().transaction(),
314                        TransactionStatus::InTransaction(_)
315                    ) {
316                        ctx.session()
317                            .add_notice(AdapterNotice::ExistingTransactionInProgress);
318                    }
319                    let result = ctx.session_mut().start_transaction(
320                        self.now_datetime(),
321                        plan.access,
322                        plan.isolation_level,
323                    );
324                    ctx.retire(result.map(|_| ExecuteResponse::StartedTransaction))
325                }
326                Plan::CommitTransaction(CommitTransactionPlan {
327                    ref transaction_type,
328                })
329                | Plan::AbortTransaction(AbortTransactionPlan {
330                    ref transaction_type,
331                }) => {
332                    let action = match &plan {
333                        Plan::CommitTransaction(_) => EndTransactionAction::Commit,
334                        Plan::AbortTransaction(_) => EndTransactionAction::Rollback,
335                        _ => unreachable!(),
336                    };
337                    if ctx.session().transaction().is_implicit() && !transaction_type.is_implicit()
338                    {
339                        // In Postgres, if a user sends a COMMIT or ROLLBACK in an
340                        // implicit transaction, a warning is sent warning them.
341                        // (The transaction is still closed and a new implicit
342                        // transaction started, though.)
343                        ctx.session().add_notice(
344                            AdapterNotice::ExplicitTransactionControlInImplicitTransaction,
345                        );
346                    }
347                    self.sequence_end_transaction(ctx, action).await;
348                }
349                Plan::ShowCreate(plan) => {
350                    ctx.retire(Ok(Self::send_immediate_rows(plan.row)));
351                }
352                Plan::CopyFrom(plan) => {
353                    self.sequence_copy_from(ctx, plan, target_cluster).await;
354                }
355                Plan::ExplainPlan(plan) => {
356                    self.sequence_explain_plan(ctx, plan).await;
357                }
358                Plan::ExplainPushdown(plan) => {
359                    self.sequence_explain_pushdown(ctx, plan).await;
360                }
361                Plan::ExplainSinkSchema(plan) => {
362                    let result = self.sequence_explain_schema(plan);
363                    ctx.retire(result);
364                }
365                // `SessionClient::execute_attempts` unrolls SQL `EXECUTE`, and `try_peek` and
366                // `try_read_then_write` take over every statement that plans to one of the
367                // others.
368                // TODO(SQL-760): Drop the manual soft panic once internal errors soft-panic
369                // centrally.
370                plan @ (Plan::CopyTo(_)
371                | Plan::Execute(_)
372                | Plan::ExplainTimestamp(_)
373                | Plan::Insert(_)
374                | Plan::ReadThenWrite(_)
375                | Plan::Select(_)
376                | Plan::ShowColumns(_)
377                | Plan::SideEffectingFunc(_)
378                | Plan::Subscribe(_)) => {
379                    let msg = format!(
380                        "{:?} plan reached the coordinator despite frontend routing",
381                        PlanKind::from(&plan)
382                    );
383                    soft_panic_or_log!("{msg}");
384                    ctx.retire(Err(AdapterError::Internal(msg)));
385                }
386                Plan::AlterNoop(plan) => {
387                    ctx.retire(Ok(ExecuteResponse::AlteredObject(plan.object_type)));
388                }
389                Plan::AlterCluster(plan) => {
390                    self.sequence_alter_cluster_staged(ctx, plan).await;
391                }
392                Plan::AlterClusterRename(_)
393                | Plan::AlterClusterSwap(_)
394                | Plan::AlterClusterReplicaRename(_)
395                | Plan::AlterItemRename(_)
396                | Plan::AlterSchemaRename(_)
397                | Plan::AlterSchemaSwap(_) => {
398                    let msg = "rename/swap must execute as a catalog command";
399                    soft_panic_or_log!("{msg}");
400                    ctx.retire(Err(AdapterError::Internal(msg.into())));
401                }
402                Plan::AlterConnection(plan) => {
403                    self.sequence_alter_connection(ctx, plan).await;
404                }
405                Plan::AlterSetCluster(plan) => {
406                    let result = self.sequence_alter_set_cluster(ctx.session(), plan).await;
407                    ctx.retire(result);
408                }
409                Plan::AlterRetainHistory(plan) => {
410                    let result = self.sequence_alter_retain_history(&mut ctx, plan).await;
411                    ctx.retire(result);
412                }
413                Plan::AlterSourceTimestampInterval(plan) => {
414                    let result = self
415                        .sequence_alter_source_timestamp_interval(&mut ctx, plan)
416                        .await;
417                    ctx.retire(result);
418                }
419                Plan::AlterRole(plan) => {
420                    let result = self.sequence_alter_role(ctx.session_mut(), plan).await;
421                    ctx.retire(result);
422                }
423                Plan::AlterSecret(plan) => {
424                    self.sequence_alter_secret(ctx, plan).await;
425                }
426                Plan::AlterSink(plan) => {
427                    self.sequence_alter_sink_prepare(ctx, plan).await;
428                }
429                Plan::AlterSource(_) => {
430                    let msg = "ALTER SOURCE must execute as a catalog command";
431                    soft_panic_or_log!("{msg}");
432                    ctx.retire(Err(AdapterError::Internal(msg.into())));
433                }
434                Plan::AlterSystemSet(plan) => {
435                    let result = self.sequence_alter_system_set(ctx.session(), plan).await;
436                    ctx.retire(result);
437                }
438                Plan::AlterSystemReset(plan) => {
439                    let result = self.sequence_alter_system_reset(ctx.session(), plan).await;
440                    ctx.retire(result);
441                }
442                Plan::AlterSystemResetAll(plan) => {
443                    let result = self
444                        .sequence_alter_system_reset_all(ctx.session(), plan)
445                        .await;
446                    ctx.retire(result);
447                }
448                Plan::AlterTableAddColumn(plan) => {
449                    let result = self.sequence_alter_table(&mut ctx, plan).await;
450                    ctx.retire(result);
451                }
452                Plan::AlterMaterializedViewApplyReplacement(plan) => {
453                    self.sequence_alter_materialized_view_apply_replacement_prepare(ctx, plan)
454                        .await;
455                }
456                Plan::AlterNetworkPolicy(plan) => {
457                    let res = self
458                        .sequence_alter_network_policy(ctx.session(), plan)
459                        .await;
460                    ctx.retire(res);
461                }
462                Plan::DiscardTemp => {
463                    self.drop_temp_items(ctx.session().conn_id()).await;
464                    ctx.retire(Ok(ExecuteResponse::DiscardedTemp));
465                }
466                Plan::DiscardAll => {
467                    // Clearing the transaction would silently discard writes staged by an
468                    // earlier statement of the same pipeline.
469                    let txn = ctx.session().transaction();
470                    let discardable =
471                        matches!(txn, TransactionStatus::Started(_)) && !txn.contains_ops();
472                    let ret = if discardable {
473                        let (_, retire_notify) = self.clear_transaction(ctx.session_mut()).await;
474                        ctx.delay_response_until(retire_notify);
475                        self.drop_temp_items(ctx.session().conn_id()).await;
476                        // NOTE: `reset()` resets session variables durably (see
477                        // `SessionVars::reset_all`). This must stay durable and
478                        // must not be reordered to depend on a transaction
479                        // commit: the transaction was just cleared above, so
480                        // there is no commit left to promote a staged reset.
481                        let params = ctx.session_mut().reset();
482                        Ok(ExecuteResponse::DiscardedAll { params })
483                    } else {
484                        Err(AdapterError::OperationProhibitsTransaction(
485                            "DISCARD ALL".into(),
486                        ))
487                    };
488                    ctx.retire(ret);
489                }
490                Plan::Declare(plan) => {
491                    self.declare(ctx, plan.name, plan.stmt, plan.sql, plan.params);
492                }
493                Plan::Fetch(FetchPlan {
494                    name,
495                    count,
496                    timeout,
497                }) => {
498                    let ctx_extra = std::mem::take(ctx.extra_mut());
499                    ctx.retire(Ok(ExecuteResponse::Fetch {
500                        name,
501                        count,
502                        timeout,
503                        ctx_extra,
504                    }));
505                }
506                Plan::Close(plan) => {
507                    if ctx.session_mut().remove_portal(&plan.name) {
508                        ctx.retire(Ok(ExecuteResponse::ClosedCursor));
509                    } else {
510                        ctx.retire(Err(AdapterError::UnknownCursor(plan.name)));
511                    }
512                }
513                Plan::Prepare(plan) => {
514                    if ctx
515                        .session()
516                        .get_prepared_statement_unverified(&plan.name)
517                        .is_some()
518                    {
519                        ctx.retire(Err(AdapterError::PreparedStatementExists(plan.name)));
520                    } else {
521                        let state_revision = StateRevision {
522                            catalog_revision: self.catalog().transient_revision(),
523                            session_state_revision: ctx.session().state_revision(),
524                        };
525                        ctx.session_mut().retain_description(
526                            plan.stmt.clone(),
527                            plan.desc.param_types.iter().cloned().map(Some).collect(),
528                            plan.desc.clone(),
529                        );
530                        ctx.session_mut().set_prepared_statement(
531                            plan.name,
532                            Some(plan.stmt),
533                            plan.sql,
534                            plan.desc,
535                            state_revision,
536                            self.now(),
537                        );
538                        ctx.retire(Ok(ExecuteResponse::Prepare));
539                    }
540                }
541                Plan::Deallocate(plan) => match plan.name {
542                    Some(name) => {
543                        if ctx.session_mut().remove_prepared_statement(&name) {
544                            ctx.retire(Ok(ExecuteResponse::Deallocate { all: false }));
545                        } else {
546                            ctx.retire(Err(AdapterError::UnknownPreparedStatement(name)));
547                        }
548                    }
549                    None => {
550                        ctx.session_mut().remove_all_prepared_statements();
551                        ctx.retire(Ok(ExecuteResponse::Deallocate { all: true }));
552                    }
553                },
554                Plan::Raise(RaisePlan { severity }) => {
555                    ctx.session()
556                        .add_notice(AdapterNotice::UserRequested { severity });
557                    ctx.retire(Ok(ExecuteResponse::Raised));
558                }
559                Plan::GrantPrivileges(plan) => {
560                    let result = self
561                        .sequence_grant_privileges(ctx.session_mut(), plan)
562                        .await;
563                    ctx.retire(result);
564                }
565                Plan::RevokePrivileges(plan) => {
566                    let result = self
567                        .sequence_revoke_privileges(ctx.session_mut(), plan)
568                        .await;
569                    ctx.retire(result);
570                }
571                Plan::AlterDefaultPrivileges(plan) => {
572                    let result = self
573                        .sequence_alter_default_privileges(ctx.session_mut(), plan)
574                        .await;
575                    ctx.retire(result);
576                }
577                Plan::GrantRole(plan) => {
578                    let result = self.sequence_grant_role(ctx.session_mut(), plan).await;
579                    ctx.retire(result);
580                }
581                Plan::RevokeRole(plan) => {
582                    let result = self.sequence_revoke_role(ctx.session_mut(), plan).await;
583                    ctx.retire(result);
584                }
585                Plan::AlterOwner(plan) => {
586                    let result = self.sequence_alter_owner(ctx.session_mut(), plan).await;
587                    ctx.retire(result);
588                }
589                Plan::ReassignOwned(plan) => {
590                    let result = self.sequence_reassign_owned(ctx.session_mut(), plan).await;
591                    ctx.retire(result);
592                }
593                Plan::ValidateConnection(plan) => {
594                    let connection = plan
595                        .connection
596                        .into_inline_connection(self.catalog().state());
597                    let current_storage_configuration = self.controller.storage.config().clone();
598                    mz_ore::task::spawn(|| "coord::validate_connection", async move {
599                        let res = match connection
600                            .validate(plan.id, &current_storage_configuration)
601                            .await
602                        {
603                            Ok(()) => Ok(ExecuteResponse::ValidatedConnection),
604                            Err(err) => Err(err.into()),
605                        };
606                        ctx.retire(res);
607                    });
608                }
609            }
610        }
611        .instrument(tracing::debug_span!("coord::sequencer::sequence_plan"));
612        crate::coord::box_dispatcher(future)
613    }
614
615    #[mz_ore::instrument(level = "debug")]
616    pub(crate) async fn sequence_execute_single_statement_transaction(
617        &mut self,
618        ctx: ExecuteContext,
619        stmt: Arc<Statement<Raw>>,
620        params: Params,
621    ) {
622        // Put the session into single statement implicit so anything can execute.
623        let (tx, internal_cmd_tx, mut session, extra, response_barriers) = ctx.into_parts();
624        assert!(matches!(session.transaction(), TransactionStatus::Default));
625        session.start_transaction_single_stmt(self.now_datetime());
626        let conn_id = session.conn_id().unhandled();
627
628        // Execute the saved statement in a temp transmitter so we can run COMMIT.
629        let (sub_tx, sub_rx) = oneshot::channel();
630        let sub_ct = ClientTransmitter::new(sub_tx, self.internal_cmd_tx.clone());
631        let sub_ctx = ExecuteContext::from_parts_with_response_barriers(
632            sub_ct,
633            internal_cmd_tx,
634            session,
635            extra,
636            response_barriers,
637        );
638        self.handle_execute_inner(stmt, params, sub_ctx).await;
639
640        // The response can need off-thread processing. Wait for it elsewhere so the coordinator can
641        // continue processing.
642        let internal_cmd_tx = self.internal_cmd_tx.clone();
643        mz_ore::task::spawn(
644            || format!("execute_single_statement:{conn_id}"),
645            async move {
646                let Ok(Response {
647                    result,
648                    session,
649                    otel_ctx,
650                }) = sub_rx.await
651                else {
652                    // Coordinator went away.
653                    return;
654                };
655                otel_ctx.attach_as_parent();
656                let (sub_tx, sub_rx) = oneshot::channel();
657                let _ = internal_cmd_tx.send(Message::Command(
658                    otel_ctx,
659                    Command::Commit {
660                        action: EndTransactionAction::Commit,
661                        session,
662                        tx: sub_tx,
663                    },
664                ));
665                let Ok(commit_response) = sub_rx.await else {
666                    // Coordinator went away.
667                    return;
668                };
669                assert!(matches!(
670                    commit_response.session.transaction(),
671                    TransactionStatus::Default
672                ));
673                // The fake, generated response was already sent to the user and we don't need to
674                // ever send an `Ok(result)` to the user, because they are expecting a response from
675                // a `COMMIT`. So, always send the `COMMIT`'s result if the original statement
676                // succeeded. If it failed, we can send an error and don't need to wrap it or send a
677                // later COMMIT or ROLLBACK.
678                let result = match (result, commit_response.result) {
679                    (Ok(_), commit) => commit,
680                    (Err(result), _) => Err(result),
681                };
682                // We ignore the resp.result because it's not clear what to do if it failed since we
683                // can only send a single ExecuteResponse to tx.
684                tx.send(result, commit_response.session);
685            }
686            .instrument(Span::current()),
687        );
688    }
689
690    /// Creates a role during connection startup.
691    ///
692    /// This should not be called from anywhere except connection startup.
693    #[mz_ore::instrument(level = "debug")]
694    pub(crate) async fn sequence_create_role_for_startup(
695        &mut self,
696        plan: CreateRolePlan,
697    ) -> Result<ExecuteResponse, AdapterError> {
698        // This does not set conn_id because it's not yet in active_conns. That is because we can't
699        // make a ConnMeta until we have a role id which we don't have until after the catalog txn
700        // is committed. Passing None here means the audit log won't have a user set in the event's
701        // user field. This seems fine because it is indeed the system that is creating this role,
702        // not a user request, and the user name is still recorded in the plan, so we aren't losing
703        // information.
704        self.sequence_create_role(None, plan).await
705    }
706
707    pub(crate) fn allocate_transient_id(&self) -> (CatalogItemId, GlobalId) {
708        self.transient_id_gen.allocate_id()
709    }
710
711    fn should_emit_rbac_notice(&self, session: &Session) -> Option<AdapterNotice> {
712        if !rbac::is_rbac_enabled_for_session(self.catalog.system_config(), session) {
713            Some(AdapterNotice::RbacUserDisabled)
714        } else {
715            None
716        }
717    }
718
719    /// Inserts the rows from `constants` into the table identified by `target_id`.
720    ///
721    /// # Panics
722    ///
723    /// Panics if `target_id` doesn't refer to a table.
724    /// Panics if `constants` is not an `MirRelationExpr::Constant`.
725    pub(crate) fn insert_constant(
726        catalog: &Catalog,
727        session: &mut Session,
728        target_id: CatalogItemId,
729        constants: MirRelationExpr,
730    ) -> Result<ExecuteResponse, AdapterError> {
731        // Insert can be queued, so we need to re-verify the id exists.
732        let desc = match catalog.try_get_entry(&target_id) {
733            Some(table) => {
734                // Inserts always happen at the latest version of a table.
735                table.relation_desc_latest().expect("table has desc")
736            }
737            None => {
738                return Err(AdapterError::Catalog(mz_catalog::memory::error::Error {
739                    kind: mz_catalog::memory::error::ErrorKind::Sql(CatalogError::UnknownItem(
740                        target_id.to_string(),
741                    )),
742                }));
743            }
744        };
745
746        match constants.as_const() {
747            Some((rows, ..)) => {
748                let rows = rows.clone()?;
749                for (row, _) in &rows {
750                    for (i, datum) in row.iter().enumerate() {
751                        desc.constraints_met(i, &datum)?;
752                    }
753                }
754                let diffs_plan = plan::SendDiffsPlan {
755                    id: target_id,
756                    updates: rows,
757                    kind: MutationKind::Insert,
758                    returning: Vec::new(),
759                    max_result_size: catalog.system_config().max_result_size(),
760                };
761                let result = Self::send_diffs(session, diffs_plan);
762                // Let schema-fencing tests land an ALTER after packing rows
763                // against this descriptor but before the transaction commits.
764                fail::fail_point!("insert_after_pack_before_commit");
765                result
766            }
767            None => panic!(
768                "tried using sequence_insert_constant on non-constant MirRelationExpr\n{}",
769                constants.pretty(),
770            ),
771        }
772    }
773
774    #[mz_ore::instrument(level = "debug")]
775    pub(crate) fn send_diffs(
776        session: &mut Session,
777        mut plan: plan::SendDiffsPlan,
778    ) -> Result<ExecuteResponse, AdapterError> {
779        let affected_rows = {
780            let mut affected_rows = Diff::from(0);
781            let mut all_positive_diffs = true;
782            // If all diffs are positive, the number of affected rows is just the
783            // sum of all unconsolidated diffs.
784            for (_, diff) in plan.updates.iter() {
785                if diff.is_negative() {
786                    all_positive_diffs = false;
787                    break;
788                }
789
790                affected_rows += diff;
791            }
792
793            if !all_positive_diffs {
794                // Consolidate rows. This is useful e.g. for an UPDATE where the row
795                // doesn't change, and we need to reflect that in the number of
796                // affected rows.
797                //
798                // NOTE: This differs from PostgreSQL, where `UPDATE t SET x = x`
799                // reports the number of rows matching the WHERE clause even when
800                // no value changes. Because Materialize works in differential
801                // dataflow, the +1 and -1 diffs for an unchanged row cancel out
802                // during consolidation, so it reports 0 affected rows. This is
803                // longstanding behavior and both read-then-write paths agree on
804                // it.
805                differential_dataflow::consolidation::consolidate(&mut plan.updates);
806
807                affected_rows = Diff::ZERO;
808                // With retractions, the number of affected rows is not the number
809                // of rows we see, but the sum of the absolute value of their diffs,
810                // e.g. if one row is retracted and another is added, the total
811                // number of rows affected is 2.
812                for (_, diff) in plan.updates.iter() {
813                    affected_rows += diff.abs();
814                }
815            }
816
817            usize::try_from(affected_rows.into_inner()).expect("positive Diff must fit")
818        };
819        event!(
820            Level::TRACE,
821            affected_rows,
822            id = format!("{:?}", plan.id),
823            kind = format!("{:?}", plan.kind),
824            updates = plan.updates.len(),
825            returning = plan.returning.len(),
826        );
827
828        session.add_transaction_ops(TransactionOps::Writes(vec![WriteOp {
829            id: plan.id,
830            rows: TableData::Rows(plan.updates),
831        }]))?;
832        if !plan.returning.is_empty() {
833            let finishing = RowSetFinishing {
834                order_by: Vec::new(),
835                limit: None,
836                offset: 0,
837                project: (0..plan.returning[0].0.iter().count()).collect(),
838            };
839            let max_returned_query_size = session.vars().max_query_result_size();
840            let duration_histogram = session.metrics().row_set_finishing_seconds();
841
842            return match finishing.finish(
843                RowCollection::new(plan.returning, &finishing.order_by),
844                plan.max_result_size,
845                Some(max_returned_query_size),
846                duration_histogram,
847            ) {
848                Ok((rows, _size_bytes)) => Ok(Self::send_immediate_rows(rows)),
849                Err(e) => Err(AdapterError::ResultSize(e)),
850            };
851        }
852        Ok(match plan.kind {
853            MutationKind::Delete => ExecuteResponse::Deleted(affected_rows),
854            MutationKind::Insert => ExecuteResponse::Inserted(affected_rows),
855            MutationKind::Update => ExecuteResponse::Updated(affected_rows / 2),
856        })
857    }
858}
859
860/// Forward notices that we got from the optimizer.
861pub(crate) fn emit_optimizer_notices(
862    catalog: &Catalog,
863    session: &Session,
864    notices: &[RawOptimizerNotice],
865) {
866    // `for_session` below is expensive, so return early if there's nothing to do.
867    if notices.is_empty() {
868        return;
869    }
870    let humanizer = catalog.for_session(session);
871    let system_vars = catalog.system_config();
872    for notice in notices {
873        let kind = OptimizerNoticeKind::from(notice);
874        let notice_enabled = match kind {
875            OptimizerNoticeKind::EqualsNull => system_vars.enable_notices_for_equals_null(),
876            OptimizerNoticeKind::IndexAlreadyExists => {
877                system_vars.enable_notices_for_index_already_exists()
878            }
879            OptimizerNoticeKind::IndexTooWideForLiteralConstraints => {
880                system_vars.enable_notices_for_index_too_wide_for_literal_constraints()
881            }
882            OptimizerNoticeKind::IndexKeyEmpty => system_vars.enable_notices_for_index_empty_key(),
883        };
884        if notice_enabled {
885            // We don't need to redact the notice parts because
886            // `emit_optimizer_notices` is only called by the `sequence_~`
887            // method for the statement that produces that notice.
888            session.add_notice(AdapterNotice::OptimizerNotice {
889                notice: notice.message(&humanizer, false).to_string(),
890                hint: notice.hint(&humanizer, false).to_string(),
891            });
892        }
893        session
894            .metrics()
895            .optimization_notices(&[kind.metric_label()])
896            .inc_by(1);
897    }
898}
899
900/// Returns a future that will execute EXPLAIN FILTER PUSHDOWN, i.e., compute the filter pushdown
901/// statistics for the given collections with the given MFPs.
902///
903/// The coordinator calls this for materialized views and the frontend peek sequencing for
904/// SELECTs, so it takes only what it needs instead of the `Coordinator`.
905pub(crate) async fn explain_pushdown_future_inner<
906    I: IntoIterator<Item = (GlobalId, MfpPlan<LirScalarExpr>)>,
907>(
908    session: &Session,
909    catalog: &Catalog,
910    storage_collections: &Arc<dyn StorageCollections + Send + Sync>,
911    as_of: Antichain<Timestamp>,
912    mz_now: ResultSpec<'static>,
913    imports: I,
914) -> impl Future<Output = Result<ExecuteResponse, AdapterError>> + use<I> {
915    let mut explain_timeout = *session.vars().statement_timeout();
916    // Timeout of 0 is equivalent to "off", meaning we will wait "forever."
917    if explain_timeout == Duration::ZERO {
918        explain_timeout = Duration::MAX;
919    }
920    let mut futures = FuturesOrdered::new();
921    for (id, plan) in imports {
922        let catalog_entry = catalog.get_entry_by_global_id(&id);
923        let full_name = catalog
924            .for_session(session)
925            .resolve_full_name(&catalog_entry.name);
926        let name = format!("{}", full_name);
927        let relation_desc = catalog_entry
928            .relation_desc()
929            .expect("source should have a proper desc")
930            .into_owned();
931        let stats_future = storage_collections
932            .snapshot_parts_stats(id, as_of.clone())
933            .await;
934
935        let mz_now = mz_now.clone();
936        // These futures may block if the source is not yet readable at the as-of;
937        // stash them in `futures` and only block on them in a separate task.
938        // TODO(peek-seq): This complication won't be needed once only the frontend calls this
939        // function, in which case it will be fine to block the current task.
940        futures.push_back(async move {
941            let snapshot_stats = match stats_future.await {
942                Ok(stats) => stats,
943                Err(e) => return Err(e),
944            };
945            let mut total_bytes = 0;
946            let mut total_parts = 0;
947            let mut selected_bytes = 0;
948            let mut selected_parts = 0;
949            for SnapshotPartStats {
950                encoded_size_bytes: bytes,
951                stats,
952            } in &snapshot_stats.parts
953            {
954                let bytes = u64::cast_from(*bytes);
955                total_bytes += bytes;
956                total_parts += 1u64;
957                let selected = match stats.as_ref().and_then(|x| x.try_decode().ok()) {
958                    // Also the arm for stats that do not decode, which a
959                    // newer writer's stats kind can produce. Both report the
960                    // part as selected, matching what a read of it would do.
961                    None => true,
962                    Some(stats) => {
963                        let stats = RelationPartStats::new(
964                            name.as_str(),
965                            &snapshot_stats.metrics.pushdown.part_stats,
966                            &relation_desc,
967                            &stats,
968                        );
969                        stats.may_match_mfp_plan(mz_now.clone(), &plan)
970                    }
971                };
972
973                if selected {
974                    selected_bytes += bytes;
975                    selected_parts += 1u64;
976                }
977            }
978            Ok(Row::pack_slice(&[
979                name.as_str().into(),
980                total_bytes.into(),
981                selected_bytes.into(),
982                total_parts.into(),
983                selected_parts.into(),
984            ]))
985        });
986    }
987
988    let fut = async move {
989        match tokio::time::timeout(
990            explain_timeout,
991            futures::TryStreamExt::try_collect::<Vec<_>>(futures),
992        )
993        .await
994        {
995            Ok(Ok(rows)) => Ok(ExecuteResponse::SendingRowsImmediate {
996                rows: Box::new(rows.into_row_iter()),
997            }),
998            Ok(Err(err)) => Err(err.into()),
999            Err(_) => Err(AdapterError::StatementTimeout),
1000        }
1001    };
1002    fut
1003}