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::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::tracing::OpenTelemetryContext;
32use mz_persist_client::stats::SnapshotPartStats;
33use mz_repr::explain::{ExprHumanizerExt, TransientItem};
34use mz_repr::{CatalogItemId, Datum, Diff, GlobalId, IntoRowIterator, Row, RowArena, Timestamp};
35use mz_sql::catalog::{CatalogError, SessionCatalog};
36use mz_sql::names::ResolvedIds;
37use mz_sql::plan::{
38    self, AbortTransactionPlan, CommitTransactionPlan, CreateRolePlan, CreateSourcePlanBundle,
39    FetchPlan, HirScalarExpr, MutationKind, Params, Plan, PlanKind, RaisePlan, SideEffectingFunc,
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 tokio::sync::oneshot;
56use tracing::{Instrument, Level, Span, event, warn};
57
58use crate::ExecuteContext;
59use crate::catalog::{Catalog, CatalogState};
60use crate::command::{Command, ExecuteResponse, Response};
61use crate::coord::appends::{DeferredOp, DeferredPlan};
62use crate::coord::validity::PlanValidity;
63use crate::coord::{
64    Coordinator, DeferredPlanStatement, ExplainPlanContext, Message, PlanStatement, TargetCluster,
65    catalog_serving,
66};
67use crate::error::AdapterError;
68use crate::explain::insights::PlanInsightsContext;
69use crate::notice::AdapterNotice;
70use crate::optimize::dataflows::{EvalTime, ExprPrep, ExprPrepOneShot};
71use crate::optimize::peek;
72use crate::session::{
73    EndTransactionAction, Session, StateRevision, TransactionOps, TransactionStatus, WriteOp,
74};
75use crate::util::ClientTransmitter;
76
77// DO NOT make this visible in any way, i.e. do not add any version of
78// `pub` to this mod. The inner `sequence_X` methods are hidden in this
79// private module to prevent anyone from calling them directly. All
80// sequencing should be done through the `sequence_plan` method.
81// This allows us to add catch-all logic that should be applied to all
82// plans in `sequence_plan` and guarantee that no caller can circumvent
83// that logic.
84//
85// The exceptions are:
86//
87// - Creating a role during connection startup. In this scenario, the session has not been properly
88// initialized and we need to skip directly to creating role. We have a specific method,
89// `sequence_create_role_for_startup` for this purpose.
90// - Methods that continue the execution of some plan that was being run asynchronously, such as
91// `sequence_peek_stage` and `sequence_create_connection_stage_finish`.
92// - The frontend peek sequencing temporarily reaches into this module for things that are needed
93//   by both the old and new peek sequencing. TODO(peek-seq): We plan to eliminate this with a
94//   big refactoring after the old peek sequencing is removed.
95
96mod inner;
97pub(crate) use inner::cancel_carried_reconfiguration;
98
99impl Coordinator {
100    /// BOXED FUTURE: As of Nov 2023 the returned Future from this function was 34KB. This would
101    /// get stored on the stack which is bad for runtime performance, and blow up our stack usage.
102    /// Because of that we purposefully move this Future onto the heap (i.e. Box it).
103    pub(crate) fn sequence_plan(
104        &mut self,
105        mut ctx: ExecuteContext,
106        plan: Plan,
107        resolved_ids: ResolvedIds,
108        sql_impl_resolved_ids: ResolvedIds,
109    ) -> LocalBoxFuture<'_, ()> {
110        async move {
111            let responses = ExecuteResponse::generated_from(&PlanKind::from(&plan));
112            ctx.tx_mut().set_allowed(responses);
113
114            if self.controller.read_only() && !plan.allowed_in_read_only() {
115                ctx.retire(Err(AdapterError::ReadOnly));
116                return;
117            }
118
119            // Check if we're still waiting for any of the builtin table appends from when we
120            // started the Session to complete.
121            if let Some((dependencies, wait_future)) =
122                super::appends::waiting_on_startup_appends(self.catalog(), ctx.session_mut(), &plan)
123            {
124                let conn_id = ctx.session().conn_id();
125                tracing::debug!(%conn_id, "deferring plan for startup appends");
126
127                let role_metadata = ctx.session().role_metadata().clone();
128                let validity =
129                    PlanValidity::new(&self.catalog, dependencies, None, None, role_metadata);
130                let deferred_plan = DeferredPlan {
131                    ctx,
132                    plan,
133                    validity,
134                    requires_locks: BTreeSet::default(),
135                    resolved_ids,
136                    sql_impl_resolved_ids,
137                };
138                // Defer op accepts an optional write lock, but there aren't any writes occurring
139                // here, since the map to `None`.
140                let acquire_future = wait_future.map(|()| None);
141
142                self.defer_op(acquire_future, DeferredOp::Plan(deferred_plan));
143
144                // Return early because our op is deferred on waiting for the builtin writes to
145                // complete.
146                return;
147            };
148
149            // Scope the borrow of the Catalog because we need to mutate the Coordinator state below.
150            let target_cluster = match ctx.session().transaction().cluster() {
151                // Use the current transaction's cluster.
152                Some(cluster_id) => TargetCluster::Transaction(cluster_id),
153                // If there isn't a current cluster set for a transaction, then try to auto route.
154                None => {
155                    let session_catalog = self.catalog.for_session(ctx.session());
156                    catalog_serving::auto_run_on_catalog_server(
157                        &session_catalog,
158                        ctx.session(),
159                        &plan,
160                    )
161                }
162            };
163            let (target_cluster_id, target_cluster_name) = match self
164                .catalog()
165                .resolve_target_cluster(target_cluster, ctx.session())
166            {
167                Ok(cluster) => (Some(cluster.id), Some(cluster.name.clone())),
168                Err(_) => (None, None),
169            };
170
171            if let (Some(cluster_id), Some(cluster_name), Some(statement_id)) = (
172                target_cluster_id,
173                target_cluster_name.clone(),
174                ctx.extra().contents(),
175            ) {
176                self.set_statement_execution_cluster(statement_id, cluster_id, cluster_name);
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            // Look up the authenticated role of the connection targeted by
192            // pg_cancel_backend, which check_plan needs for its RBAC check.
193            // Linear search through active connections is fine because this
194            // happens at most once per statement.
195            let target_conn_role = match &plan {
196                Plan::SideEffectingFunc(SideEffectingFunc::PgCancelBackend {
197                    connection_id: Some(connection_id),
198                }) => self
199                    .active_conns()
200                    .into_iter()
201                    .find(|(conn_id, _)| conn_id.unhandled() == *connection_id)
202                    .map(|(_, conn_meta)| *conn_meta.authenticated_role_id()),
203                _ => None,
204            };
205
206            if let Err(e) = rbac::check_plan(
207                &session_catalog,
208                target_conn_role,
209                ctx.session(),
210                &plan,
211                target_cluster_id,
212                &resolved_ids,
213                &sql_impl_resolved_ids,
214            ) {
215                return ctx.retire(Err(e.into()));
216            }
217
218            match plan {
219                Plan::CreateSource(plan) => {
220                    let (item_id, global_id) = return_if_err!(self.allocate_user_id().await, ctx);
221                    let result = self
222                        .sequence_create_source(
223                            &mut ctx,
224                            vec![CreateSourcePlanBundle {
225                                item_id,
226                                global_id,
227                                plan,
228                                resolved_ids,
229                                available_source_references: None,
230                            }],
231                        )
232                        .await;
233                    ctx.retire(result);
234                }
235                Plan::CreateSources(plans) => {
236                    assert!(
237                        resolved_ids.is_empty(),
238                        "each plan has separate resolved_ids"
239                    );
240                    let result = self.sequence_create_source(&mut ctx, plans).await;
241                    ctx.retire(result);
242                }
243                Plan::CreateConnection(plan) => {
244                    self.sequence_create_connection(ctx, plan, resolved_ids)
245                        .await;
246                }
247                Plan::CreateDatabase(plan) => {
248                    let result = self.sequence_create_database(ctx.session_mut(), plan).await;
249                    ctx.retire(result);
250                }
251                Plan::CreateSchema(plan) => {
252                    let result = self.sequence_create_schema(ctx.session_mut(), plan).await;
253                    ctx.retire(result);
254                }
255                Plan::CreateRole(plan) => {
256                    let result = self
257                        .sequence_create_role(Some(ctx.session().conn_id()), plan)
258                        .await;
259                    if let Some(notice) = self.should_emit_rbac_notice(ctx.session()) {
260                        ctx.session().add_notice(notice);
261                    }
262                    ctx.retire(result);
263                }
264                Plan::CreateCluster(plan) => {
265                    let result = self.sequence_create_cluster(ctx.session(), plan).await;
266                    ctx.retire(result);
267                }
268                Plan::CreateClusterReplica(plan) => {
269                    let result = self
270                        .sequence_create_cluster_replica(ctx.session(), plan)
271                        .await;
272                    ctx.retire(result);
273                }
274                Plan::CreateTable(plan) => {
275                    let result = self
276                        .sequence_create_table(&mut ctx, plan, resolved_ids)
277                        .await;
278                    ctx.retire(result);
279                }
280                Plan::CreateSecret(plan) => {
281                    self.sequence_create_secret(ctx, plan).await;
282                }
283                Plan::CreateSink(plan) => {
284                    self.sequence_create_sink(ctx, plan, resolved_ids).await;
285                }
286                Plan::CreateView(plan) => {
287                    self.sequence_create_view(ctx, plan, resolved_ids).await;
288                }
289                Plan::CreateMaterializedView(plan) => {
290                    self.sequence_create_materialized_view(ctx, plan, resolved_ids)
291                        .await;
292                }
293                Plan::CreateIndex(plan) => {
294                    self.sequence_create_index(ctx, plan, resolved_ids).await;
295                }
296                Plan::CreateType(plan) => {
297                    let result = self
298                        .sequence_create_type(ctx.session(), plan, resolved_ids)
299                        .await;
300                    ctx.retire(result);
301                }
302                Plan::CreateNetworkPolicy(plan) => {
303                    let res = self
304                        .sequence_create_network_policy(ctx.session(), plan)
305                        .await;
306                    ctx.retire(res);
307                }
308                Plan::Comment(plan) => {
309                    let result = self.sequence_comment_on(ctx.session(), plan).await;
310                    ctx.retire(result);
311                }
312                Plan::CopyTo(plan) => {
313                    self.sequence_copy_to(ctx, plan, target_cluster).await;
314                }
315                Plan::DropObjects(plan) => {
316                    let result = self.sequence_drop_objects(&mut ctx, plan).await;
317                    ctx.retire(result);
318                }
319                Plan::DropOwned(plan) => {
320                    let result = self.sequence_drop_owned(ctx.session_mut(), plan).await;
321                    ctx.retire(result);
322                }
323                Plan::EmptyQuery => {
324                    ctx.retire(Ok(ExecuteResponse::EmptyQuery));
325                }
326                Plan::ShowAllVariables => {
327                    let result = self.sequence_show_all_variables(ctx.session());
328                    ctx.retire(result);
329                }
330                Plan::ShowVariable(plan) => {
331                    let result = self.sequence_show_variable(ctx.session(), plan);
332                    ctx.retire(result);
333                }
334                Plan::InspectShard(plan) => {
335                    // TODO: Ideally, this await would happen off the main thread.
336                    let result = self.sequence_inspect_shard(ctx.session(), plan).await;
337                    ctx.retire(result);
338                }
339                Plan::SetVariable(plan) => {
340                    let result = self.sequence_set_variable(ctx.session_mut(), plan);
341                    ctx.retire(result);
342                }
343                Plan::ResetVariable(plan) => {
344                    let result = self.sequence_reset_variable(ctx.session_mut(), plan);
345                    ctx.retire(result);
346                }
347                Plan::SetTransaction(plan) => {
348                    let result = self.sequence_set_transaction(ctx.session_mut(), plan);
349                    ctx.retire(result);
350                }
351                Plan::StartTransaction(plan) => {
352                    if matches!(
353                        ctx.session().transaction(),
354                        TransactionStatus::InTransaction(_)
355                    ) {
356                        ctx.session()
357                            .add_notice(AdapterNotice::ExistingTransactionInProgress);
358                    }
359                    let result = ctx.session_mut().start_transaction(
360                        self.now_datetime(),
361                        plan.access,
362                        plan.isolation_level,
363                    );
364                    ctx.retire(result.map(|_| ExecuteResponse::StartedTransaction))
365                }
366                Plan::CommitTransaction(CommitTransactionPlan {
367                    ref transaction_type,
368                })
369                | Plan::AbortTransaction(AbortTransactionPlan {
370                    ref transaction_type,
371                }) => {
372                    // Serialize DDL transactions. Statements that use this mode must return false
373                    // in `must_serialize_ddl()`.
374                    if ctx.session().transaction().is_ddl() {
375                        if let Ok(guard) = self.serialized_ddl.try_lock_owned() {
376                            let prev = self
377                                .active_conns
378                                .get_mut(ctx.session().conn_id())
379                                .expect("connection must exist")
380                                .deferred_lock
381                                .replace(guard);
382                            assert!(
383                                prev.is_none(),
384                                "connections should have at most one lock guard"
385                            );
386                        } else {
387                            self.serialized_ddl.push_back(DeferredPlanStatement {
388                                ctx,
389                                ps: PlanStatement::Plan {
390                                    plan,
391                                    resolved_ids,
392                                    sql_impl_resolved_ids,
393                                },
394                            });
395                            return;
396                        }
397                    }
398
399                    let action = match &plan {
400                        Plan::CommitTransaction(_) => EndTransactionAction::Commit,
401                        Plan::AbortTransaction(_) => EndTransactionAction::Rollback,
402                        _ => unreachable!(),
403                    };
404                    if ctx.session().transaction().is_implicit() && !transaction_type.is_implicit()
405                    {
406                        // In Postgres, if a user sends a COMMIT or ROLLBACK in an
407                        // implicit transaction, a warning is sent warning them.
408                        // (The transaction is still closed and a new implicit
409                        // transaction started, though.)
410                        ctx.session().add_notice(
411                            AdapterNotice::ExplicitTransactionControlInImplicitTransaction,
412                        );
413                    }
414                    self.sequence_end_transaction(ctx, action).await;
415                }
416                Plan::Select(plan) => {
417                    let max = Some(ctx.session().vars().max_query_result_size());
418                    self.sequence_peek(ctx, plan, target_cluster, max).await;
419                }
420                Plan::Subscribe(plan) => {
421                    self.sequence_subscribe(ctx, plan, target_cluster).await;
422                }
423                Plan::SideEffectingFunc(plan) => {
424                    self.sequence_side_effecting_func(ctx, plan).await;
425                }
426                Plan::ShowCreate(plan) => {
427                    ctx.retire(Ok(Self::send_immediate_rows(plan.row)));
428                }
429                Plan::ShowColumns(show_columns_plan) => {
430                    let max = Some(ctx.session().vars().max_query_result_size());
431                    self.sequence_peek(ctx, show_columns_plan.select_plan, target_cluster, max)
432                        .await;
433                }
434                Plan::CopyFrom(plan) => {
435                    self.sequence_copy_from(ctx, plan, target_cluster).await;
436                }
437                Plan::ExplainPlan(plan) => {
438                    self.sequence_explain_plan(ctx, plan, target_cluster).await;
439                }
440                Plan::ExplainPushdown(plan) => {
441                    self.sequence_explain_pushdown(ctx, plan, target_cluster)
442                        .await;
443                }
444                Plan::ExplainSinkSchema(plan) => {
445                    let result = self.sequence_explain_schema(plan);
446                    ctx.retire(result);
447                }
448                Plan::ExplainTimestamp(plan) => {
449                    self.sequence_explain_timestamp(ctx, plan, target_cluster)
450                        .await;
451                }
452                Plan::Insert(plan) => {
453                    self.sequence_insert(ctx, plan).await;
454                }
455                Plan::ReadThenWrite(plan) => {
456                    self.sequence_read_then_write(ctx, plan).await;
457                }
458                Plan::AlterNoop(plan) => {
459                    ctx.retire(Ok(ExecuteResponse::AlteredObject(plan.object_type)));
460                }
461                Plan::AlterCluster(plan) => {
462                    self.sequence_alter_cluster_staged(ctx, plan).await;
463                }
464                Plan::AlterClusterRename(plan) => {
465                    let result = self.sequence_alter_cluster_rename(&mut ctx, plan).await;
466                    ctx.retire(result);
467                }
468                Plan::AlterClusterSwap(plan) => {
469                    let result = self.sequence_alter_cluster_swap(&mut ctx, plan).await;
470                    ctx.retire(result);
471                }
472                Plan::AlterClusterReplicaRename(plan) => {
473                    let result = self
474                        .sequence_alter_cluster_replica_rename(ctx.session(), plan)
475                        .await;
476                    ctx.retire(result);
477                }
478                Plan::AlterConnection(plan) => {
479                    self.sequence_alter_connection(ctx, plan).await;
480                }
481                Plan::AlterSetCluster(plan) => {
482                    let result = self.sequence_alter_set_cluster(ctx.session(), plan).await;
483                    ctx.retire(result);
484                }
485                Plan::AlterRetainHistory(plan) => {
486                    let result = self.sequence_alter_retain_history(&mut ctx, plan).await;
487                    ctx.retire(result);
488                }
489                Plan::AlterSourceTimestampInterval(plan) => {
490                    let result = self
491                        .sequence_alter_source_timestamp_interval(&mut ctx, plan)
492                        .await;
493                    ctx.retire(result);
494                }
495                Plan::AlterItemRename(plan) => {
496                    let result = self.sequence_alter_item_rename(&mut ctx, plan).await;
497                    ctx.retire(result);
498                }
499                Plan::AlterSchemaRename(plan) => {
500                    let result = self.sequence_alter_schema_rename(&mut ctx, plan).await;
501                    ctx.retire(result);
502                }
503                Plan::AlterSchemaSwap(plan) => {
504                    let result = self.sequence_alter_schema_swap(&mut ctx, plan).await;
505                    ctx.retire(result);
506                }
507                Plan::AlterRole(plan) => {
508                    let result = self.sequence_alter_role(ctx.session_mut(), plan).await;
509                    ctx.retire(result);
510                }
511                Plan::AlterSecret(plan) => {
512                    self.sequence_alter_secret(ctx, plan).await;
513                }
514                Plan::AlterSink(plan) => {
515                    self.sequence_alter_sink_prepare(ctx, plan).await;
516                }
517                Plan::AlterSource(plan) => {
518                    let result = self.sequence_alter_source(ctx.session_mut(), plan).await;
519                    ctx.retire(result);
520                }
521                Plan::AlterSystemSet(plan) => {
522                    let result = self.sequence_alter_system_set(ctx.session(), plan).await;
523                    ctx.retire(result);
524                }
525                Plan::AlterSystemReset(plan) => {
526                    let result = self.sequence_alter_system_reset(ctx.session(), plan).await;
527                    ctx.retire(result);
528                }
529                Plan::AlterSystemResetAll(plan) => {
530                    let result = self
531                        .sequence_alter_system_reset_all(ctx.session(), plan)
532                        .await;
533                    ctx.retire(result);
534                }
535                Plan::AlterTableAddColumn(plan) => {
536                    let result = self.sequence_alter_table(&mut ctx, plan).await;
537                    ctx.retire(result);
538                }
539                Plan::AlterMaterializedViewApplyReplacement(plan) => {
540                    self.sequence_alter_materialized_view_apply_replacement_prepare(ctx, plan)
541                        .await;
542                }
543                Plan::AlterNetworkPolicy(plan) => {
544                    let res = self
545                        .sequence_alter_network_policy(ctx.session(), plan)
546                        .await;
547                    ctx.retire(res);
548                }
549                Plan::DiscardTemp => {
550                    self.drop_temp_items(ctx.session().conn_id()).await;
551                    ctx.retire(Ok(ExecuteResponse::DiscardedTemp));
552                }
553                Plan::DiscardAll => {
554                    // Clearing the transaction would silently discard writes staged by an
555                    // earlier statement of the same pipeline.
556                    let txn = ctx.session().transaction();
557                    let discardable =
558                        matches!(txn, TransactionStatus::Started(_)) && !txn.contains_ops();
559                    let ret = if discardable {
560                        let (_, retire_notify) = self.clear_transaction(ctx.session_mut()).await;
561                        ctx.delay_response_until(retire_notify);
562                        self.drop_temp_items(ctx.session().conn_id()).await;
563                        // NOTE: `reset()` resets session variables durably (see
564                        // `SessionVars::reset_all`). This must stay durable and
565                        // must not be reordered to depend on a transaction
566                        // commit: the transaction was just cleared above, so
567                        // there is no commit left to promote a staged reset.
568                        ctx.session_mut().reset();
569                        Ok(ExecuteResponse::DiscardedAll)
570                    } else {
571                        Err(AdapterError::OperationProhibitsTransaction(
572                            "DISCARD ALL".into(),
573                        ))
574                    };
575                    ctx.retire(ret);
576                }
577                Plan::Declare(plan) => {
578                    self.declare(ctx, plan.name, plan.stmt, plan.sql, plan.params);
579                }
580                Plan::Fetch(FetchPlan {
581                    name,
582                    count,
583                    timeout,
584                }) => {
585                    let ctx_extra = std::mem::take(ctx.extra_mut());
586                    ctx.retire(Ok(ExecuteResponse::Fetch {
587                        name,
588                        count,
589                        timeout,
590                        ctx_extra,
591                    }));
592                }
593                Plan::Close(plan) => {
594                    if ctx.session_mut().remove_portal(&plan.name) {
595                        ctx.retire(Ok(ExecuteResponse::ClosedCursor));
596                    } else {
597                        ctx.retire(Err(AdapterError::UnknownCursor(plan.name)));
598                    }
599                }
600                Plan::Prepare(plan) => {
601                    if ctx
602                        .session()
603                        .get_prepared_statement_unverified(&plan.name)
604                        .is_some()
605                    {
606                        ctx.retire(Err(AdapterError::PreparedStatementExists(plan.name)));
607                    } else {
608                        let state_revision = StateRevision {
609                            catalog_revision: self.catalog().transient_revision(),
610                            session_state_revision: ctx.session().state_revision(),
611                        };
612                        ctx.session_mut().set_prepared_statement(
613                            plan.name,
614                            Some(plan.stmt),
615                            plan.sql,
616                            plan.desc,
617                            state_revision,
618                            self.now(),
619                        );
620                        ctx.retire(Ok(ExecuteResponse::Prepare));
621                    }
622                }
623                Plan::Execute(plan) => {
624                    match self.sequence_execute(ctx.session_mut(), plan) {
625                        Ok(portal_name) => {
626                            let (tx, _, session, extra, response_barriers) = ctx.into_parts();
627                            let command = Message::Command(
628                                OpenTelemetryContext::obtain(),
629                                Command::Execute {
630                                    portal_name,
631                                    session,
632                                    tx: tx.take(),
633                                    outer_ctx_extra: Some(extra),
634                                },
635                            );
636                            if response_barriers.is_empty() {
637                                self.internal_cmd_tx
638                                    .send(command)
639                                    .expect("sending to self.internal_cmd_tx cannot fail");
640                            } else {
641                                let internal_cmd_tx = self.internal_cmd_tx.clone();
642                                mz_ore::task::spawn(
643                                    || "execute_after_response_barriers",
644                                    async move {
645                                        for barrier in response_barriers {
646                                            barrier.await;
647                                        }
648                                        let _ = internal_cmd_tx.send(command);
649                                    },
650                                );
651                            }
652                        }
653                        Err(err) => ctx.retire(Err(err)),
654                    };
655                }
656                Plan::Deallocate(plan) => match plan.name {
657                    Some(name) => {
658                        if ctx.session_mut().remove_prepared_statement(&name) {
659                            ctx.retire(Ok(ExecuteResponse::Deallocate { all: false }));
660                        } else {
661                            ctx.retire(Err(AdapterError::UnknownPreparedStatement(name)));
662                        }
663                    }
664                    None => {
665                        ctx.session_mut().remove_all_prepared_statements();
666                        ctx.retire(Ok(ExecuteResponse::Deallocate { all: true }));
667                    }
668                },
669                Plan::Raise(RaisePlan { severity }) => {
670                    ctx.session()
671                        .add_notice(AdapterNotice::UserRequested { severity });
672                    ctx.retire(Ok(ExecuteResponse::Raised));
673                }
674                Plan::GrantPrivileges(plan) => {
675                    let result = self
676                        .sequence_grant_privileges(ctx.session_mut(), plan)
677                        .await;
678                    ctx.retire(result);
679                }
680                Plan::RevokePrivileges(plan) => {
681                    let result = self
682                        .sequence_revoke_privileges(ctx.session_mut(), plan)
683                        .await;
684                    ctx.retire(result);
685                }
686                Plan::AlterDefaultPrivileges(plan) => {
687                    let result = self
688                        .sequence_alter_default_privileges(ctx.session_mut(), plan)
689                        .await;
690                    ctx.retire(result);
691                }
692                Plan::GrantRole(plan) => {
693                    let result = self.sequence_grant_role(ctx.session_mut(), plan).await;
694                    ctx.retire(result);
695                }
696                Plan::RevokeRole(plan) => {
697                    let result = self.sequence_revoke_role(ctx.session_mut(), plan).await;
698                    ctx.retire(result);
699                }
700                Plan::AlterOwner(plan) => {
701                    let result = self.sequence_alter_owner(ctx.session_mut(), plan).await;
702                    ctx.retire(result);
703                }
704                Plan::ReassignOwned(plan) => {
705                    let result = self.sequence_reassign_owned(ctx.session_mut(), plan).await;
706                    ctx.retire(result);
707                }
708                Plan::ValidateConnection(plan) => {
709                    let connection = plan
710                        .connection
711                        .into_inline_connection(self.catalog().state());
712                    let current_storage_configuration = self.controller.storage.config().clone();
713                    mz_ore::task::spawn(|| "coord::validate_connection", async move {
714                        let res = match connection
715                            .validate(plan.id, &current_storage_configuration)
716                            .await
717                        {
718                            Ok(()) => Ok(ExecuteResponse::ValidatedConnection),
719                            Err(err) => Err(err.into()),
720                        };
721                        ctx.retire(res);
722                    });
723                }
724            }
725        }
726        .instrument(tracing::debug_span!("coord::sequencer::sequence_plan"))
727        .boxed_local()
728    }
729
730    #[mz_ore::instrument(level = "debug")]
731    pub(crate) async fn sequence_execute_single_statement_transaction(
732        &mut self,
733        ctx: ExecuteContext,
734        stmt: Arc<Statement<Raw>>,
735        params: Params,
736    ) {
737        // Put the session into single statement implicit so anything can execute.
738        let (tx, internal_cmd_tx, mut session, extra, response_barriers) = ctx.into_parts();
739        assert!(matches!(session.transaction(), TransactionStatus::Default));
740        session.start_transaction_single_stmt(self.now_datetime());
741        let conn_id = session.conn_id().unhandled();
742
743        // Execute the saved statement in a temp transmitter so we can run COMMIT.
744        let (sub_tx, sub_rx) = oneshot::channel();
745        let sub_ct = ClientTransmitter::new(sub_tx, self.internal_cmd_tx.clone());
746        let sub_ctx = ExecuteContext::from_parts_with_response_barriers(
747            sub_ct,
748            internal_cmd_tx,
749            session,
750            extra,
751            response_barriers,
752        );
753        self.handle_execute_inner(stmt, params, sub_ctx).await;
754
755        // The response can need off-thread processing. Wait for it elsewhere so the coordinator can
756        // continue processing.
757        let internal_cmd_tx = self.internal_cmd_tx.clone();
758        mz_ore::task::spawn(
759            || format!("execute_single_statement:{conn_id}"),
760            async move {
761                let Ok(Response {
762                    result,
763                    session,
764                    otel_ctx,
765                }) = sub_rx.await
766                else {
767                    // Coordinator went away.
768                    return;
769                };
770                otel_ctx.attach_as_parent();
771                let (sub_tx, sub_rx) = oneshot::channel();
772                let _ = internal_cmd_tx.send(Message::Command(
773                    otel_ctx,
774                    Command::Commit {
775                        action: EndTransactionAction::Commit,
776                        session,
777                        tx: sub_tx,
778                    },
779                ));
780                let Ok(commit_response) = sub_rx.await else {
781                    // Coordinator went away.
782                    return;
783                };
784                assert!(matches!(
785                    commit_response.session.transaction(),
786                    TransactionStatus::Default
787                ));
788                // The fake, generated response was already sent to the user and we don't need to
789                // ever send an `Ok(result)` to the user, because they are expecting a response from
790                // a `COMMIT`. So, always send the `COMMIT`'s result if the original statement
791                // succeeded. If it failed, we can send an error and don't need to wrap it or send a
792                // later COMMIT or ROLLBACK.
793                let result = match (result, commit_response.result) {
794                    (Ok(_), commit) => commit,
795                    (Err(result), _) => Err(result),
796                };
797                // We ignore the resp.result because it's not clear what to do if it failed since we
798                // can only send a single ExecuteResponse to tx.
799                tx.send(result, commit_response.session);
800            }
801            .instrument(Span::current()),
802        );
803    }
804
805    /// Creates a role during connection startup.
806    ///
807    /// This should not be called from anywhere except connection startup.
808    #[mz_ore::instrument(level = "debug")]
809    pub(crate) async fn sequence_create_role_for_startup(
810        &mut self,
811        plan: CreateRolePlan,
812    ) -> Result<ExecuteResponse, AdapterError> {
813        // This does not set conn_id because it's not yet in active_conns. That is because we can't
814        // make a ConnMeta until we have a role id which we don't have until after the catalog txn
815        // is committed. Passing None here means the audit log won't have a user set in the event's
816        // user field. This seems fine because it is indeed the system that is creating this role,
817        // not a user request, and the user name is still recorded in the plan, so we aren't losing
818        // information.
819        self.sequence_create_role(None, plan).await
820    }
821
822    pub(crate) fn allocate_transient_id(&self) -> (CatalogItemId, GlobalId) {
823        self.transient_id_gen.allocate_id()
824    }
825
826    fn should_emit_rbac_notice(&self, session: &Session) -> Option<AdapterNotice> {
827        if !rbac::is_rbac_enabled_for_session(self.catalog.system_config(), session) {
828            Some(AdapterNotice::RbacUserDisabled)
829        } else {
830            None
831        }
832    }
833
834    /// Inserts the rows from `constants` into the table identified by `target_id`.
835    ///
836    /// # Panics
837    ///
838    /// Panics if `target_id` doesn't refer to a table.
839    /// Panics if `constants` is not an `MirRelationExpr::Constant`.
840    pub(crate) fn insert_constant(
841        catalog: &Catalog,
842        session: &mut Session,
843        target_id: CatalogItemId,
844        constants: MirRelationExpr,
845    ) -> Result<ExecuteResponse, AdapterError> {
846        // Insert can be queued, so we need to re-verify the id exists.
847        let desc = match catalog.try_get_entry(&target_id) {
848            Some(table) => {
849                // Inserts always happen at the latest version of a table.
850                table.relation_desc_latest().expect("table has desc")
851            }
852            None => {
853                return Err(AdapterError::Catalog(mz_catalog::memory::error::Error {
854                    kind: mz_catalog::memory::error::ErrorKind::Sql(CatalogError::UnknownItem(
855                        target_id.to_string(),
856                    )),
857                }));
858            }
859        };
860
861        match constants.as_const() {
862            Some((rows, ..)) => {
863                let rows = rows.clone()?;
864                for (row, _) in &rows {
865                    for (i, datum) in row.iter().enumerate() {
866                        desc.constraints_met(i, &datum)?;
867                    }
868                }
869                let diffs_plan = plan::SendDiffsPlan {
870                    id: target_id,
871                    updates: rows,
872                    kind: MutationKind::Insert,
873                    returning: Vec::new(),
874                    max_result_size: catalog.system_config().max_result_size(),
875                };
876                Self::send_diffs(session, diffs_plan)
877            }
878            None => panic!(
879                "tried using sequence_insert_constant on non-constant MirRelationExpr\n{}",
880                constants.pretty(),
881            ),
882        }
883    }
884
885    #[mz_ore::instrument(level = "debug")]
886    pub(crate) fn send_diffs(
887        session: &mut Session,
888        mut plan: plan::SendDiffsPlan,
889    ) -> Result<ExecuteResponse, AdapterError> {
890        let affected_rows = {
891            let mut affected_rows = Diff::from(0);
892            let mut all_positive_diffs = true;
893            // If all diffs are positive, the number of affected rows is just the
894            // sum of all unconsolidated diffs.
895            for (_, diff) in plan.updates.iter() {
896                if diff.is_negative() {
897                    all_positive_diffs = false;
898                    break;
899                }
900
901                affected_rows += diff;
902            }
903
904            if !all_positive_diffs {
905                // Consolidate rows. This is useful e.g. for an UPDATE where the row
906                // doesn't change, and we need to reflect that in the number of
907                // affected rows.
908                differential_dataflow::consolidation::consolidate(&mut plan.updates);
909
910                affected_rows = Diff::ZERO;
911                // With retractions, the number of affected rows is not the number
912                // of rows we see, but the sum of the absolute value of their diffs,
913                // e.g. if one row is retracted and another is added, the total
914                // number of rows affected is 2.
915                for (_, diff) in plan.updates.iter() {
916                    affected_rows += diff.abs();
917                }
918            }
919
920            usize::try_from(affected_rows.into_inner()).expect("positive Diff must fit")
921        };
922        event!(
923            Level::TRACE,
924            affected_rows,
925            id = format!("{:?}", plan.id),
926            kind = format!("{:?}", plan.kind),
927            updates = plan.updates.len(),
928            returning = plan.returning.len(),
929        );
930
931        session.add_transaction_ops(TransactionOps::Writes(vec![WriteOp {
932            id: plan.id,
933            rows: TableData::Rows(plan.updates),
934        }]))?;
935        if !plan.returning.is_empty() {
936            let finishing = RowSetFinishing {
937                order_by: Vec::new(),
938                limit: None,
939                offset: 0,
940                project: (0..plan.returning[0].0.iter().count()).collect(),
941            };
942            let max_returned_query_size = session.vars().max_query_result_size();
943            let duration_histogram = session.metrics().row_set_finishing_seconds();
944
945            return match finishing.finish(
946                RowCollection::new(plan.returning, &finishing.order_by),
947                plan.max_result_size,
948                Some(max_returned_query_size),
949                duration_histogram,
950            ) {
951                Ok((rows, _size_bytes)) => Ok(Self::send_immediate_rows(rows)),
952                Err(e) => Err(AdapterError::ResultSize(e)),
953            };
954        }
955        Ok(match plan.kind {
956            MutationKind::Delete => ExecuteResponse::Deleted(affected_rows),
957            MutationKind::Insert => ExecuteResponse::Inserted(affected_rows),
958            MutationKind::Update => ExecuteResponse::Updated(affected_rows / 2),
959        })
960    }
961}
962
963/// Checks whether we should emit diagnostic
964/// information associated with reading per-replica sources.
965///
966/// If an unrecoverable error is found (today: an untargeted read on a
967/// cluster with a non-1 number of replicas), return that.  Otherwise,
968/// return a list of associated notices (today: we always emit exactly
969/// one notice if there are any per-replica log dependencies and if
970/// `emit_introspection_query_notice` is set, and none otherwise.)
971pub(crate) fn check_log_reads(
972    catalog: &Catalog,
973    cluster: &Cluster,
974    source_ids: &BTreeSet<GlobalId>,
975    target_replica: &mut Option<ReplicaId>,
976    vars: &SessionVars,
977) -> Result<impl IntoIterator<Item = AdapterNotice>, AdapterError>
978where
979{
980    let log_names = source_ids
981        .iter()
982        .map(|gid| catalog.resolve_item_id(gid))
983        .flat_map(|item_id| catalog.introspection_dependencies(item_id))
984        .map(|item_id| catalog.get_entry(&item_id).name().item.clone())
985        .collect::<Vec<_>>();
986
987    if log_names.is_empty() {
988        return Ok(None);
989    }
990
991    // Reading from log sources on replicated clusters is only allowed if a
992    // target replica is selected. Otherwise, we have no way of knowing which
993    // replica we read the introspection data from.
994    let num_replicas = cluster.replicas().count();
995    if target_replica.is_none() {
996        if num_replicas == 1 {
997            *target_replica = cluster.replicas().map(|r| r.replica_id).next();
998        } else {
999            return Err(AdapterError::UntargetedLogRead { log_names });
1000        }
1001    }
1002
1003    // Ensure that logging is initialized for the target replica, lest
1004    // we try to read from a non-existing arrangement.
1005    let replica_id = target_replica.expect("set to `Some` above");
1006    let replica = &cluster.replica(replica_id).expect("Replica must exist");
1007    if !replica.config.compute.logging.enabled() {
1008        return Err(AdapterError::IntrospectionDisabled { log_names });
1009    }
1010
1011    Ok(vars
1012        .emit_introspection_query_notice()
1013        .then_some(AdapterNotice::PerReplicaLogRead { log_names }))
1014}
1015
1016/// Forward notices that we got from the optimizer.
1017pub(crate) fn emit_optimizer_notices(
1018    catalog: &Catalog,
1019    session: &Session,
1020    notices: &[RawOptimizerNotice],
1021) {
1022    // `for_session` below is expensive, so return early if there's nothing to do.
1023    if notices.is_empty() {
1024        return;
1025    }
1026    let humanizer = catalog.for_session(session);
1027    let system_vars = catalog.system_config();
1028    for notice in notices {
1029        let kind = OptimizerNoticeKind::from(notice);
1030        let notice_enabled = match kind {
1031            OptimizerNoticeKind::EqualsNull => system_vars.enable_notices_for_equals_null(),
1032            OptimizerNoticeKind::IndexAlreadyExists => {
1033                system_vars.enable_notices_for_index_already_exists()
1034            }
1035            OptimizerNoticeKind::IndexTooWideForLiteralConstraints => {
1036                system_vars.enable_notices_for_index_too_wide_for_literal_constraints()
1037            }
1038            OptimizerNoticeKind::IndexKeyEmpty => system_vars.enable_notices_for_index_empty_key(),
1039        };
1040        if notice_enabled {
1041            // We don't need to redact the notice parts because
1042            // `emit_optimizer_notices` is only called by the `sequence_~`
1043            // method for the statement that produces that notice.
1044            session.add_notice(AdapterNotice::OptimizerNotice {
1045                notice: notice.message(&humanizer, false).to_string(),
1046                hint: notice.hint(&humanizer, false).to_string(),
1047            });
1048        }
1049        session
1050            .metrics()
1051            .optimization_notices(&[kind.metric_label()])
1052            .inc_by(1);
1053    }
1054}
1055
1056/// Evaluates a COPY TO target URI expression and validates it.
1057///
1058/// This function is shared between the old peek sequencing (sequence_copy_to)
1059/// and the new frontend peek sequencing to avoid code duplication.
1060pub fn eval_copy_to_uri(
1061    to: HirScalarExpr,
1062    session: &Session,
1063    catalog_state: &CatalogState,
1064) -> Result<Uri, AdapterError> {
1065    let style = ExprPrepOneShot {
1066        logical_time: EvalTime::NotAvailable,
1067        session,
1068        catalog_state,
1069    };
1070    let mut to = to.lower_uncorrelated(catalog_state.system_config())?;
1071    style.prep_scalar_expr(&mut to)?;
1072    let temp_storage = RowArena::new();
1073    let evaled = to.eval(&[], &temp_storage)?;
1074    if evaled == Datum::Null {
1075        coord_bail!("COPY TO target value can not be null");
1076    }
1077    let to_url = match Uri::from_str(evaled.unwrap_str()) {
1078        Ok(url) => {
1079            if url.scheme_str() != Some("s3") && url.scheme_str() != Some("gs") {
1080                coord_bail!("only 's3://...' and 'gs://...' urls are supported as COPY TO target");
1081            }
1082            url
1083        }
1084        Err(e) => coord_bail!("could not parse COPY TO target url: {}", e),
1085    };
1086    Ok(to_url)
1087}
1088
1089/// Returns a future that will execute EXPLAIN FILTER PUSHDOWN, i.e., compute the filter pushdown
1090/// statistics for the given collections with the given MFPs.
1091///
1092/// (Shared helper fn between the old and new sequencing. This doesn't take the Coordinator as a
1093/// parameter, but instead just the specifically necessary things are passed in, so that the
1094/// frontend peek sequencing can also call it.)
1095pub(crate) async fn explain_pushdown_future_inner<
1096    I: IntoIterator<Item = (GlobalId, MapFilterProject)>,
1097>(
1098    session: &Session,
1099    catalog: &Catalog,
1100    storage_collections: &Arc<dyn StorageCollections + Send + Sync>,
1101    as_of: Antichain<Timestamp>,
1102    mz_now: ResultSpec<'static>,
1103    imports: I,
1104) -> impl Future<Output = Result<ExecuteResponse, AdapterError>> + use<I> {
1105    let mut explain_timeout = *session.vars().statement_timeout();
1106    // Timeout of 0 is equivalent to "off", meaning we will wait "forever."
1107    if explain_timeout == Duration::ZERO {
1108        explain_timeout = Duration::MAX;
1109    }
1110    let mut futures = FuturesOrdered::new();
1111    for (id, mfp) in imports {
1112        let catalog_entry = catalog.get_entry_by_global_id(&id);
1113        let full_name = catalog
1114            .for_session(session)
1115            .resolve_full_name(&catalog_entry.name);
1116        let name = format!("{}", full_name);
1117        let relation_desc = catalog_entry
1118            .relation_desc()
1119            .expect("source should have a proper desc")
1120            .into_owned();
1121        let stats_future = storage_collections
1122            .snapshot_parts_stats(id, as_of.clone())
1123            .await;
1124
1125        let mz_now = mz_now.clone();
1126        // These futures may block if the source is not yet readable at the as-of;
1127        // stash them in `futures` and only block on them in a separate task.
1128        // TODO(peek-seq): This complication won't be needed once this function will only be called
1129        // from the new peek sequencing, in which case it will be fine to block the current task.
1130        futures.push_back(async move {
1131            let snapshot_stats = match stats_future.await {
1132                Ok(stats) => stats,
1133                Err(e) => return Err(e),
1134            };
1135            let mut total_bytes = 0;
1136            let mut total_parts = 0;
1137            let mut selected_bytes = 0;
1138            let mut selected_parts = 0;
1139            for SnapshotPartStats {
1140                encoded_size_bytes: bytes,
1141                stats,
1142            } in &snapshot_stats.parts
1143            {
1144                let bytes = u64::cast_from(*bytes);
1145                total_bytes += bytes;
1146                total_parts += 1u64;
1147                let selected = match stats {
1148                    None => true,
1149                    Some(stats) => {
1150                        let stats = stats.decode();
1151                        let stats = RelationPartStats::new(
1152                            name.as_str(),
1153                            &snapshot_stats.metrics.pushdown.part_stats,
1154                            &relation_desc,
1155                            &stats,
1156                        );
1157                        stats.may_match_mfp(mz_now.clone(), &mfp)
1158                    }
1159                };
1160
1161                if selected {
1162                    selected_bytes += bytes;
1163                    selected_parts += 1u64;
1164                }
1165            }
1166            Ok(Row::pack_slice(&[
1167                name.as_str().into(),
1168                total_bytes.into(),
1169                selected_bytes.into(),
1170                total_parts.into(),
1171                selected_parts.into(),
1172            ]))
1173        });
1174    }
1175
1176    let fut = async move {
1177        match tokio::time::timeout(
1178            explain_timeout,
1179            futures::TryStreamExt::try_collect::<Vec<_>>(futures),
1180        )
1181        .await
1182        {
1183            Ok(Ok(rows)) => Ok(ExecuteResponse::SendingRowsImmediate {
1184                rows: Box::new(rows.into_row_iter()),
1185            }),
1186            Ok(Err(err)) => Err(err.into()),
1187            Err(_) => Err(AdapterError::StatementTimeout),
1188        }
1189    };
1190    fut
1191}
1192
1193/// Generates EXPLAIN PLAN output.
1194/// (Shared helper fn between the old and new sequencing.)
1195pub(crate) async fn explain_plan_inner(
1196    session: &Session,
1197    catalog: &Catalog,
1198    df_meta: DataflowMetainfo,
1199    explain_ctx: ExplainPlanContext,
1200    optimizer: peek::Optimizer,
1201    insights_ctx: Option<Box<PlanInsightsContext>>,
1202) -> Result<Vec<Row>, AdapterError> {
1203    let ExplainPlanContext {
1204        config,
1205        format,
1206        stage,
1207        desc,
1208        optimizer_trace,
1209        ..
1210    } = explain_ctx;
1211
1212    let desc = desc.expect("RelationDesc for SelectPlan in EXPLAIN mode");
1213
1214    let session_catalog = catalog.for_session(session);
1215    let expr_humanizer = {
1216        let transient_items = btreemap! {
1217            optimizer.select_id() => TransientItem::new(
1218                Some(vec![GlobalId::Explain.to_string()]),
1219                Some(desc.iter_names().map(|c| c.to_string()).collect()),
1220            )
1221        };
1222        ExprHumanizerExt::new(transient_items, &session_catalog)
1223    };
1224
1225    let finishing = if optimizer.finishing().is_trivial(desc.arity()) {
1226        None
1227    } else {
1228        Some(optimizer.finishing().clone())
1229    };
1230
1231    let target_cluster = catalog.get_cluster(optimizer.cluster_id());
1232    let features = optimizer.config().features.clone();
1233
1234    let rows = optimizer_trace
1235        .into_rows(
1236            format,
1237            &config,
1238            &features,
1239            &expr_humanizer,
1240            finishing,
1241            Some(target_cluster),
1242            df_meta,
1243            stage,
1244            plan::ExplaineeStatementKind::Select,
1245            insights_ctx,
1246        )
1247        .await?;
1248
1249    Ok(rows)
1250}
1251
1252/// Creates a statistics oracle for query optimization.
1253///
1254/// This is a free-standing function that can be called from both the old peek sequencing
1255/// and the new frontend peek sequencing.
1256pub(crate) async fn statistics_oracle(
1257    session: &Session,
1258    source_ids: &BTreeSet<GlobalId>,
1259    query_as_of: &Antichain<Timestamp>,
1260    is_oneshot: bool,
1261    system_config: &vars::SystemVars,
1262    storage_collections: &dyn StorageCollections,
1263) -> Result<Box<dyn StatisticsOracle>, AdapterError> {
1264    if !session.vars().enable_session_cardinality_estimates() {
1265        let stats: Box<dyn StatisticsOracle> = Box::new(EmptyStatisticsOracle);
1266        return Ok(stats);
1267    }
1268
1269    let timeout = if is_oneshot {
1270        // TODO(mgree): ideally, we would shorten the timeout even more if we think the query could take the fast path
1271        system_config.optimizer_oneshot_stats_timeout()
1272    } else {
1273        system_config.optimizer_stats_timeout()
1274    };
1275
1276    let cached_stats = mz_ore::future::timeout(
1277        timeout,
1278        CachedStatisticsOracle::new(source_ids, query_as_of, storage_collections),
1279    )
1280    .await;
1281
1282    match cached_stats {
1283        Ok(stats) => Ok(Box::new(stats)),
1284        Err(mz_ore::future::TimeoutError::DeadlineElapsed) => {
1285            warn!(
1286                is_oneshot = is_oneshot,
1287                "optimizer statistics collection timed out after {}ms",
1288                timeout.as_millis()
1289            );
1290
1291            Ok(Box::new(EmptyStatisticsOracle))
1292        }
1293        Err(mz_ore::future::TimeoutError::Inner(e)) => Err(AdapterError::Storage(e)),
1294    }
1295}
1296
1297#[derive(Debug)]
1298struct CachedStatisticsOracle {
1299    cache: BTreeMap<GlobalId, usize>,
1300}
1301
1302impl CachedStatisticsOracle {
1303    pub async fn new(
1304        ids: &BTreeSet<GlobalId>,
1305        as_of: &Antichain<Timestamp>,
1306        storage_collections: &dyn StorageCollections,
1307    ) -> Result<Self, StorageError> {
1308        let mut cache = BTreeMap::new();
1309
1310        for id in ids {
1311            let stats = storage_collections.snapshot_stats(*id, as_of.clone()).await;
1312
1313            match stats {
1314                Ok(stats) => {
1315                    cache.insert(*id, stats.num_updates);
1316                }
1317                Err(StorageError::IdentifierMissing(id)) => {
1318                    ::tracing::debug!("no statistics for {id}")
1319                }
1320                Err(e) => return Err(e),
1321            }
1322        }
1323
1324        Ok(Self { cache })
1325    }
1326}
1327
1328impl StatisticsOracle for CachedStatisticsOracle {
1329    fn cardinality_estimate(&self, id: GlobalId) -> Option<usize> {
1330        self.cache.get(&id).map(|estimate| *estimate)
1331    }
1332
1333    fn as_map(&self) -> BTreeMap<GlobalId, usize> {
1334        self.cache.clone()
1335    }
1336}