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