mz_adapter/
command.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
10use std::collections::{BTreeMap, BTreeSet};
11use std::net::IpAddr;
12use std::pin::Pin;
13use std::sync::Arc;
14use std::time::Duration;
15
16use chrono::{DateTime, Utc};
17use derivative::Derivative;
18use enum_kinds::EnumKind;
19use futures::Stream;
20use mz_adapter_types::connection::{ConnectionId, ConnectionIdType};
21use mz_auth::password::Password;
22use mz_cluster_client::ReplicaId;
23use mz_compute_types::ComputeInstanceId;
24use mz_controller_types::ClusterId;
25use mz_expr::RowSetFinishing;
26use mz_ore::collections::CollectionExt;
27use mz_ore::soft_assert_no_log;
28use mz_ore::tracing::OpenTelemetryContext;
29use mz_persist_client::PersistClient;
30use mz_pgcopy::CopyFormatParams;
31use mz_repr::global_id::TransientIdGen;
32use mz_repr::role_id::RoleId;
33use mz_repr::{CatalogItemId, ColumnIndex, Diff, GlobalId, Row, RowIterator, SqlRelationType};
34use mz_sql::ast::{FetchDirection, Raw, Statement};
35use mz_sql::catalog::ObjectType;
36use mz_sql::optimizer_metrics::OptimizerMetrics;
37use mz_sql::plan;
38use mz_sql::plan::{ExecuteTimeout, Plan, PlanKind, SideEffectingFunc};
39use mz_sql::session::user::User;
40use mz_sql::session::vars::{OwnedVarInput, SystemVars};
41use mz_sql_parser::ast::{AlterObjectRenameStatement, AlterOwnerStatement, DropObjectsStatement};
42use mz_storage_types::sources::Timeline;
43use mz_timestamp_oracle::TimestampOracle;
44use tokio::sync::{Semaphore, mpsc, oneshot, watch};
45use uuid::Uuid;
46
47use crate::active_compute_sink::ActiveSubscribeOwner;
48use crate::catalog::Catalog;
49use crate::config::{ScopedParameters, ScopedParametersScope, SystemParameterFrontend};
50use crate::coord::appends::{BuiltinTableAppendNotify, WriteResult};
51use crate::coord::consistency::CoordinatorInconsistencies;
52use crate::coord::peek::{PeekDataflowPlan, PeekResponseUnary};
53use crate::coord::timestamp_selection::TimestampDetermination;
54use crate::coord::{ExecuteContextExtra, ExecuteContextGuard};
55use crate::error::AdapterError;
56use crate::optimize::LirDataflowDescription;
57use crate::session::{EndTransactionAction, RowBatchStream, Session};
58use crate::statement_logging::{
59    FrontendStatementLoggingEvent, StatementEndedExecutionReason, StatementExecutionStrategy,
60    StatementLoggingFrontend,
61};
62use crate::statement_logging::{StatementLoggingId, WatchSetCreation};
63use crate::util::Transmittable;
64use crate::webhook::AppendWebhookResponse;
65use crate::{
66    AdapterNotice, AppendWebhookError, CollectionIdBundle, ReadHolds, TimestampExplanation,
67};
68
69/// A handle for pgwire to stream raw byte chunks to the coordinator's
70/// parallel background batch builder tasks during COPY FROM STDIN.
71#[derive(Debug)]
72pub struct CopyFromStdinWriter {
73    /// Channels for distributing raw byte chunks (split at row boundaries)
74    /// across parallel worker tasks. pgwire round-robins chunks across these.
75    pub batch_txs: Vec<mpsc::Sender<Vec<u8>>>,
76    /// Receives the final result (`Vec<ProtoBatch>` + total row count, or error)
77    /// from the collector task. pgwire uses this to commit the batches.
78    pub completion_rx: oneshot::Receiver<
79        Result<(Vec<mz_persist_client::batch::ProtoBatch>, u64), crate::AdapterError>,
80    >,
81}
82
83#[derive(Debug)]
84pub struct CatalogSnapshot {
85    pub catalog: Arc<Catalog>,
86}
87
88#[derive(Debug)]
89pub enum Command {
90    CatalogSnapshot {
91        tx: oneshot::Sender<CatalogSnapshot>,
92    },
93
94    Startup {
95        tx: oneshot::Sender<Result<StartupResponse, AdapterError>>,
96        user: User,
97        conn_id: ConnectionId,
98        client_ip: Option<IpAddr>,
99        secret_key: u32,
100        uuid: Uuid,
101        application_name: String,
102        notice_tx: mpsc::UnboundedSender<AdapterNotice>,
103    },
104
105    AuthenticatePassword {
106        tx: oneshot::Sender<Result<(), AdapterError>>,
107        role_name: String,
108        password: Option<Password>,
109    },
110
111    AuthenticateGetSASLChallenge {
112        tx: oneshot::Sender<Result<SASLChallengeResponse, AdapterError>>,
113        role_name: String,
114        nonce: String,
115    },
116
117    AuthenticateVerifySASLProof {
118        tx: oneshot::Sender<Result<SASLVerifyProofResponse, AdapterError>>,
119        role_name: String,
120        proof: String,
121        auth_message: String,
122        mock_hash: String,
123    },
124
125    CheckRoleCanLogin {
126        tx: oneshot::Sender<Result<(), AdapterError>>,
127        role_name: String,
128    },
129
130    Execute {
131        portal_name: String,
132        session: Session,
133        tx: oneshot::Sender<Response<ExecuteResponse>>,
134        /// The end-of-execution obligation of the statement this execution
135        /// serves, if any. `Coordinator::handle_execute` arms it again on
136        /// receipt. `None` means no log entry exists yet, so the coordinator
137        /// begins one.
138        outer_ctx_extra: Option<ExecuteContextExtra>,
139    },
140
141    /// Attempts to commit or abort the session's transaction. Guarantees that the Coordinator's
142    /// transaction state has been cleared, even if the commit or abort fails. (A failure can
143    /// happen, for example, if the session's role id has been dropped which will prevent
144    /// sequence_end_transaction from running.)
145    Commit {
146        action: EndTransactionAction,
147        session: Session,
148        tx: oneshot::Sender<Response<ExecuteResponse>>,
149    },
150
151    CancelRequest {
152        conn_id: ConnectionIdType,
153        secret_key: u32,
154    },
155
156    PrivilegedCancelRequest {
157        conn_id: ConnectionId,
158    },
159
160    GetWebhook {
161        database: String,
162        schema: String,
163        name: String,
164        tx: oneshot::Sender<Result<AppendWebhookResponse, AppendWebhookError>>,
165    },
166
167    GetSystemVars {
168        tx: oneshot::Sender<SystemVars>,
169    },
170
171    SetSystemVars {
172        vars: BTreeMap<String, String>,
173        conn_id: ConnectionId,
174        tx: oneshot::Sender<Result<(), AdapterError>>,
175    },
176
177    /// Replace the scoped feature-flag overrides (the complete desired state).
178    /// Computed by the system-parameter sync loop from continuous LaunchDarkly
179    /// evaluation. The coordinator stores this working copy and reconciles it
180    /// into the per-scope resolution boundaries (the compute controller's
181    /// per-replica dyncfg layer for `replica`-scoped parameters). See the
182    /// scoped feature flags design.
183    UpdateScopedSystemParameters {
184        overrides: ScopedParameters,
185        /// Bounds which objects' durable rows the reconcile may prune. See
186        /// [`crate::catalog::Op::UpdateScopedSystemParameters`].
187        prune_scope: ScopedParametersScope,
188        tx: oneshot::Sender<()>,
189    },
190
191    /// Install (or replace) the shared system-parameter frontend on the
192    /// coordinator, so catalog transactions can resolve a new object's scoped
193    /// overrides before its replica is provisioned or its first dataflow is
194    /// planned, rather than waiting for the next sync tick. Sent by the sync loop
195    /// whenever it (re)initializes the frontend. See the scoped feature flags
196    /// design.
197    InstallScopedSystemParameterFrontend {
198        frontend: Arc<SystemParameterFrontend>,
199    },
200
201    InjectAuditEvents {
202        events: Vec<crate::catalog::InjectedAuditEvent>,
203        conn_id: ConnectionId,
204        tx: oneshot::Sender<Result<(), AdapterError>>,
205    },
206
207    Terminate {
208        conn_id: ConnectionId,
209        tx: Option<oneshot::Sender<Result<(), AdapterError>>>,
210    },
211
212    /// Sets up a streaming COPY FROM STDIN operation. The coordinator
213    /// creates parallel background batch builder tasks and returns a
214    /// [`CopyFromStdinWriter`] that pgwire uses to stream raw byte chunks.
215    StartCopyFromStdin {
216        target_id: CatalogItemId,
217        target_name: String,
218        columns: Vec<ColumnIndex>,
219        /// The row description for the target table (used for constraint checks).
220        row_desc: mz_repr::RelationDesc,
221        /// Copy format parameters (text/csv/binary) for decoding raw bytes.
222        params: mz_pgcopy::CopyFormatParams<'static>,
223        session: Session,
224        tx: oneshot::Sender<Response<CopyFromStdinWriter>>,
225    },
226
227    /// Performs any cleanup and logging actions necessary for
228    /// finalizing a statement execution.
229    ///
230    /// Only used for cases that terminate in the protocol layer and
231    /// otherwise have no reason to hand control back to the coordinator.
232    /// In other cases, we piggy-back on another command.
233    RetireExecute {
234        data: ExecuteContextExtra,
235        reason: StatementEndedExecutionReason,
236    },
237
238    CheckConsistency {
239        tx: oneshot::Sender<Result<(), CoordinatorInconsistencies>>,
240    },
241
242    Dump {
243        tx: oneshot::Sender<Result<serde_json::Value, anyhow::Error>>,
244    },
245
246    GetComputeInstanceClient {
247        instance_id: ComputeInstanceId,
248        tx: oneshot::Sender<
249            Result<
250                mz_compute_client::controller::instance_client::InstanceClient,
251                mz_compute_client::controller::error::InstanceMissing,
252            >,
253        >,
254    },
255
256    GetOracle {
257        timeline: Timeline,
258        tx: oneshot::Sender<
259            Result<Arc<dyn TimestampOracle<mz_repr::Timestamp> + Send + Sync>, AdapterError>,
260        >,
261    },
262
263    DetermineRealTimeRecentTimestamp {
264        source_ids: BTreeSet<GlobalId>,
265        real_time_recency_timeout: Duration,
266        tx: oneshot::Sender<Result<Option<mz_repr::Timestamp>, AdapterError>>,
267    },
268
269    GetTransactionReadHoldsBundle {
270        conn_id: ConnectionId,
271        tx: oneshot::Sender<Option<ReadHolds>>,
272    },
273
274    /// _Merges_ the given read holds into the given connection's stored transaction read holds.
275    StoreTransactionReadHolds {
276        conn_id: ConnectionId,
277        read_holds: ReadHolds,
278        tx: oneshot::Sender<()>,
279    },
280
281    ExecuteSlowPathPeek {
282        dataflow_plan: Box<PeekDataflowPlan>,
283        determination: TimestampDetermination,
284        finishing: RowSetFinishing,
285        compute_instance: ComputeInstanceId,
286        target_replica: Option<ReplicaId>,
287        intermediate_result_type: SqlRelationType,
288        source_ids: BTreeSet<GlobalId>,
289        conn_id: ConnectionId,
290        max_result_size: u64,
291        max_query_result_size: Option<u64>,
292        /// If statement logging is enabled, contains all info needed for installing watch sets
293        /// and logging the statement execution.
294        watch_set: Option<WatchSetCreation>,
295        tx: oneshot::Sender<Result<ExecuteResponse, AdapterError>>,
296    },
297
298    ExecuteSubscribe {
299        df_desc: mz_compute_types::dataflows::LirDataflowDescription,
300        dependency_ids: BTreeSet<GlobalId>,
301        cluster_id: ComputeInstanceId,
302        replica_id: Option<ReplicaId>,
303        conn_id: ConnectionId,
304        session_uuid: Uuid,
305        read_holds: ReadHolds,
306        plan: plan::SubscribePlan,
307        statement_logging_id: Option<StatementLoggingId>,
308        tx: oneshot::Sender<Result<ExecuteResponse, AdapterError>>,
309    },
310
311    /// Preflight check for COPY TO S3 operation. This runs the slow S3 operations
312    /// (loading SDK config, checking bucket path, verifying permissions, uploading sentinel)
313    /// in a background task to avoid blocking the coordinator.
314    CopyToPreflight {
315        /// The S3 connection info needed for preflight checks.
316        s3_sink_connection: mz_compute_types::sinks::CopyToS3OneshotSinkConnection,
317        /// The sink ID for logging and S3 key management.
318        sink_id: GlobalId,
319        /// Response channel for the preflight result.
320        tx: oneshot::Sender<Result<(), AdapterError>>,
321    },
322
323    ExecuteCopyTo {
324        df_desc: Box<mz_compute_types::dataflows::LirDataflowDescription>,
325        compute_instance: ComputeInstanceId,
326        target_replica: Option<ReplicaId>,
327        source_ids: BTreeSet<GlobalId>,
328        conn_id: ConnectionId,
329        /// If statement logging is enabled, contains all info needed for installing watch sets
330        /// and logging the statement execution.
331        watch_set: Option<WatchSetCreation>,
332        tx: oneshot::Sender<Result<ExecuteResponse, AdapterError>>,
333    },
334
335    /// Execute a side-effecting function from the frontend peek path.
336    ExecuteSideEffectingFunc {
337        plan: SideEffectingFunc,
338        conn_id: ConnectionId,
339        tx: oneshot::Sender<Result<ExecuteResponse, AdapterError>>,
340    },
341
342    /// Look up an active connection by its raw connection ID, returning its
343    /// `ConnectionId` handle and authenticated role, or `None` if there is no
344    /// such connection.
345    ///
346    /// While the caller holds the returned `ConnectionId` handle, the raw
347    /// connection ID cannot be reused by a new connection. Frontend peek
348    /// sequencing relies on this to ensure that the connection it performs an
349    /// RBAC check against is the same one that a subsequent
350    /// `ExecuteSideEffectingFunc` acts on.
351    LookupConnection {
352        connection_id: u32,
353        tx: oneshot::Sender<Option<(ConnectionId, RoleId)>>,
354    },
355
356    /// Register a pending peek initiated by frontend sequencing. This is needed for:
357    /// - statement logging
358    /// - query cancellation
359    RegisterFrontendPeek {
360        uuid: Uuid,
361        conn_id: ConnectionId,
362        cluster_id: mz_controller_types::ClusterId,
363        depends_on: BTreeSet<GlobalId>,
364        is_fast_path: bool,
365        /// If statement logging is enabled, contains all info needed for installing watch sets
366        /// and logging the statement execution.
367        watch_set: Option<WatchSetCreation>,
368        tx: oneshot::Sender<Result<(), AdapterError>>,
369    },
370
371    /// Unregister and retire a pending peek that was registered but then
372    /// failed to issue, ending its statement-logging execution with the given
373    /// reason.
374    ///
375    /// Registration handed ownership of end-of-execution logging to the
376    /// coordinator, so the frontend must not log the end itself. If a
377    /// concurrent teardown (e.g. a `DROP CLUSTER`) already retired the peek
378    /// and logged its end, this is a no-op.
379    UnregisterFrontendPeek {
380        uuid: Uuid,
381        reason: StatementEndedExecutionReason,
382        tx: oneshot::Sender<()>,
383    },
384
385    /// Generate a timestamp explanation, or fail with `ConcurrentDependencyDrop`
386    /// if a collection in `id_bundle` no longer exists.
387    /// This is used when `emit_timestamp_notice` is enabled.
388    ExplainTimestamp {
389        conn_id: ConnectionId,
390        session_wall_time: DateTime<Utc>,
391        cluster_id: ClusterId,
392        id_bundle: CollectionIdBundle,
393        determination: TimestampDetermination,
394        tx: oneshot::Sender<Result<TimestampExplanation, AdapterError>>,
395    },
396
397    /// Statement logging event from frontend peek sequencing.
398    /// No response channel needed - this is fire-and-forget.
399    FrontendStatementLogging(FrontendStatementLoggingEvent),
400
401    /// Registers a connection-scoped cancellation watch and returns a receiver
402    /// that becomes `true` when cancellation is requested for the connection.
403    ///
404    /// Registration always installs a fresh channel, so the caller cannot
405    /// observe a cancellation aimed at an earlier statement.
406    RegisterConnectionCancelWatch {
407        conn_id: ConnectionId,
408        tx: oneshot::Sender<watch::Receiver<bool>>,
409    },
410
411    /// Creates an internal subscribe, meaning one that writes no
412    /// `mz_subscriptions` row, and returns the response channel. Used by
413    /// frontend-sequenced read-then-write (DELETE/UPDATE/INSERT...SELECT)
414    /// operations via OCC.
415    CreateInternalSubscribe {
416        df_desc: Box<LirDataflowDescription>,
417        cluster_id: ComputeInstanceId,
418        replica_id: Option<ReplicaId>,
419        depends_on: BTreeSet<GlobalId>,
420        as_of: mz_repr::Timestamp,
421        arity: usize,
422        sink_id: GlobalId,
423        owner: ActiveSubscribeOwner,
424        start_time: mz_ore::now::EpochMillis,
425        read_holds: ReadHolds,
426        tx: oneshot::Sender<Result<mpsc::UnboundedReceiver<PeekResponseUnary>, AdapterError>>,
427    },
428
429    /// Submits a write attempt. Carries the accumulated diffs to write.
430    ///
431    /// `write_ts` selects between two modes:
432    /// - `Some(ts)`: the write must land at exactly `ts`, and reports
433    ///   `WriteResult::TimestampPassed` if the table's timestamp is already past
434    ///   it. The caller decides whether to recompute the diffs and try again.
435    /// - `None`: the coordinator picks the timestamp from the oracle during
436    ///   group commit, so the timestamp cannot be passed. Every other outcome,
437    ///   including read-only, a changed target and cancellation, is reported the
438    ///   same way in both modes.
439    AttemptWrite {
440        attempt: WriteAttemptKind,
441        target_id: CatalogItemId,
442        target_global_id: GlobalId,
443        diffs: Vec<(Row, Diff)>,
444        tx: oneshot::Sender<WriteResult>,
445    },
446
447    /// Drops an internal subscribe. Fire-and-forget, the caller does not wait
448    /// for completion.
449    DropInternalSubscribe {
450        sink_id: GlobalId,
451    },
452}
453
454/// Who a read-then-write commits on behalf of, and how its timestamp is chosen.
455///
456/// Group commit picking the timestamp requires a connection to answer through,
457/// so that combination is only reachable from a session.
458#[derive(Debug)]
459pub enum WriteAttemptKind {
460    /// A session's write, cancelled with `conn_id` if the connection goes away
461    /// before it commits. A `write_ts` of `None` lets group commit pick the
462    /// timestamp, which then cannot be reported as passed.
463    Session {
464        conn_id: ConnectionId,
465        write_ts: Option<mz_repr::Timestamp>,
466    },
467    /// Coordinator background work. There is no connection to cancel with, so
468    /// the caller names the timestamp and handles `TimestampPassed` itself.
469    Background { write_ts: mz_repr::Timestamp },
470}
471
472impl Command {
473    pub fn session(&self) -> Option<&Session> {
474        match self {
475            Command::Execute { session, .. }
476            | Command::Commit { session, .. }
477            | Command::StartCopyFromStdin { session, .. } => Some(session),
478            Command::CancelRequest { .. }
479            | Command::Startup { .. }
480            | Command::AuthenticatePassword { .. }
481            | Command::AuthenticateGetSASLChallenge { .. }
482            | Command::AuthenticateVerifySASLProof { .. }
483            | Command::CheckRoleCanLogin { .. }
484            | Command::CatalogSnapshot { .. }
485            | Command::PrivilegedCancelRequest { .. }
486            | Command::GetWebhook { .. }
487            | Command::Terminate { .. }
488            | Command::GetSystemVars { .. }
489            | Command::SetSystemVars { .. }
490            | Command::UpdateScopedSystemParameters { .. }
491            | Command::InstallScopedSystemParameterFrontend { .. }
492            | Command::RetireExecute { .. }
493            | Command::CheckConsistency { .. }
494            | Command::Dump { .. }
495            | Command::GetComputeInstanceClient { .. }
496            | Command::GetOracle { .. }
497            | Command::DetermineRealTimeRecentTimestamp { .. }
498            | Command::GetTransactionReadHoldsBundle { .. }
499            | Command::StoreTransactionReadHolds { .. }
500            | Command::ExecuteSlowPathPeek { .. }
501            | Command::ExecuteSubscribe { .. }
502            | Command::CopyToPreflight { .. }
503            | Command::ExecuteCopyTo { .. }
504            | Command::ExecuteSideEffectingFunc { .. }
505            | Command::LookupConnection { .. }
506            | Command::RegisterFrontendPeek { .. }
507            | Command::UnregisterFrontendPeek { .. }
508            | Command::ExplainTimestamp { .. }
509            | Command::FrontendStatementLogging(..)
510            | Command::InjectAuditEvents { .. }
511            | Command::RegisterConnectionCancelWatch { .. }
512            | Command::CreateInternalSubscribe { .. }
513            | Command::AttemptWrite { .. }
514            | Command::DropInternalSubscribe { .. } => None,
515        }
516    }
517
518    pub fn session_mut(&mut self) -> Option<&mut Session> {
519        match self {
520            Command::Execute { session, .. }
521            | Command::Commit { session, .. }
522            | Command::StartCopyFromStdin { session, .. } => Some(session),
523            Command::CancelRequest { .. }
524            | Command::Startup { .. }
525            | Command::AuthenticatePassword { .. }
526            | Command::AuthenticateGetSASLChallenge { .. }
527            | Command::AuthenticateVerifySASLProof { .. }
528            | Command::CheckRoleCanLogin { .. }
529            | Command::CatalogSnapshot { .. }
530            | Command::PrivilegedCancelRequest { .. }
531            | Command::GetWebhook { .. }
532            | Command::Terminate { .. }
533            | Command::GetSystemVars { .. }
534            | Command::SetSystemVars { .. }
535            | Command::UpdateScopedSystemParameters { .. }
536            | Command::InstallScopedSystemParameterFrontend { .. }
537            | Command::RetireExecute { .. }
538            | Command::CheckConsistency { .. }
539            | Command::Dump { .. }
540            | Command::GetComputeInstanceClient { .. }
541            | Command::GetOracle { .. }
542            | Command::DetermineRealTimeRecentTimestamp { .. }
543            | Command::GetTransactionReadHoldsBundle { .. }
544            | Command::StoreTransactionReadHolds { .. }
545            | Command::ExecuteSlowPathPeek { .. }
546            | Command::ExecuteSubscribe { .. }
547            | Command::CopyToPreflight { .. }
548            | Command::ExecuteCopyTo { .. }
549            | Command::ExecuteSideEffectingFunc { .. }
550            | Command::LookupConnection { .. }
551            | Command::RegisterFrontendPeek { .. }
552            | Command::UnregisterFrontendPeek { .. }
553            | Command::ExplainTimestamp { .. }
554            | Command::FrontendStatementLogging(..)
555            | Command::InjectAuditEvents { .. }
556            | Command::RegisterConnectionCancelWatch { .. }
557            | Command::CreateInternalSubscribe { .. }
558            | Command::AttemptWrite { .. }
559            | Command::DropInternalSubscribe { .. } => None,
560        }
561    }
562}
563
564#[derive(Debug)]
565pub struct Response<T> {
566    pub result: Result<T, AdapterError>,
567    pub session: Session,
568    pub otel_ctx: OpenTelemetryContext,
569}
570
571#[derive(Debug, Clone, Copy)]
572pub struct SuperuserAttribute(pub Option<bool>);
573
574/// The response to [`Client::startup`](crate::Client::startup).
575#[derive(Derivative)]
576#[derivative(Debug)]
577pub struct StartupResponse {
578    /// Cancellation notifications for session-owned completion waits.
579    pub frontend_cancel_rx: tokio::sync::watch::Receiver<()>,
580    /// RoleId for the user.
581    pub role_id: RoleId,
582    /// The role's superuser attribute in the Catalog.
583    /// This attribute is None for Cloud. Cloud is able
584    /// to derive the role's superuser status from
585    /// external_metadata_rx.
586    pub superuser_attribute: SuperuserAttribute,
587    /// A future that completes when all necessary Builtin Table writes have completed.
588    #[derivative(Debug = "ignore")]
589    pub write_notify: BuiltinTableAppendNotify,
590    /// Map of (name, VarInput::Flat) tuples of session default variables that should be set.
591    pub session_defaults: BTreeMap<String, OwnedVarInput>,
592    pub catalog: Arc<Catalog>,
593    pub storage_collections:
594        Arc<dyn mz_storage_client::storage_collections::StorageCollections + Send + Sync>,
595    pub transient_id_gen: Arc<TransientIdGen>,
596    pub optimizer_metrics: OptimizerMetrics,
597    pub persist_client: PersistClient,
598    pub statement_logging_frontend: StatementLoggingFrontend,
599    /// Semaphore for limiting concurrent OCC (optimistic concurrency control)
600    /// write operations.
601    pub occ_write_semaphore: Arc<Semaphore>,
602    /// Requests a group commit, which is how the frontend asks for the write
603    /// timeline to advance without having anything to write.
604    pub group_commit_notifier: crate::coord::appends::GroupCommitNotifier,
605    /// Whether the coordinator is in read-only mode (e.g. during 0dt upgrades).
606    /// The frontend path must reject mutations when this is true.
607    pub read_only: bool,
608}
609
610#[derive(Derivative)]
611#[derivative(Debug)]
612pub struct SASLChallengeResponse {
613    pub iteration_count: usize,
614    /// Base64-encoded salt for the SASL challenge.
615    pub salt: String,
616    pub nonce: String,
617}
618
619#[derive(Derivative)]
620#[derivative(Debug)]
621pub struct SASLVerifyProofResponse {
622    pub verifier: String,
623}
624
625// Facile implementation for `StartupResponse`, which does not use the `allowed`
626// feature of `ClientTransmitter`.
627impl Transmittable for StartupResponse {
628    type Allowed = bool;
629    fn to_allowed(&self) -> Self::Allowed {
630        true
631    }
632}
633
634/// The response to [`SessionClient::dump_catalog`](crate::SessionClient::dump_catalog).
635#[derive(Debug, Clone)]
636pub struct CatalogDump(String);
637
638impl CatalogDump {
639    pub fn new(raw: String) -> Self {
640        CatalogDump(raw)
641    }
642
643    pub fn into_string(self) -> String {
644        self.0
645    }
646}
647
648impl Transmittable for CatalogDump {
649    type Allowed = bool;
650    fn to_allowed(&self) -> Self::Allowed {
651        true
652    }
653}
654
655impl Transmittable for SystemVars {
656    type Allowed = bool;
657    fn to_allowed(&self) -> Self::Allowed {
658        true
659    }
660}
661
662/// The response to [`SessionClient::execute`](crate::SessionClient::execute).
663#[derive(EnumKind, Derivative)]
664#[derivative(Debug)]
665#[enum_kind(ExecuteResponseKind, derive(PartialOrd, Ord))]
666pub enum ExecuteResponse {
667    /// The default privileges were altered.
668    AlteredDefaultPrivileges,
669    /// The requested object was altered.
670    AlteredObject(ObjectType),
671    /// The role was altered.
672    AlteredRole,
673    /// The system configuration was altered.
674    AlteredSystemConfiguration,
675    /// The requested cursor was closed.
676    ClosedCursor,
677    /// The provided comment was created.
678    Comment,
679    /// The specified number of rows were copied into the requested output.
680    Copied(usize),
681    /// The response for a COPY TO STDOUT query.
682    CopyTo {
683        format: mz_sql::plan::CopyFormat,
684        resp: Box<ExecuteResponse>,
685    },
686    CopyFrom {
687        /// Table we're copying into.
688        target_id: CatalogItemId,
689        /// Human-readable full name of the target table.
690        target_name: String,
691        columns: Vec<ColumnIndex>,
692        params: CopyFormatParams<'static>,
693        ctx_extra: ExecuteContextGuard,
694    },
695    /// The requested connection was created.
696    CreatedConnection,
697    /// The requested database was created.
698    CreatedDatabase,
699    /// The requested schema was created.
700    CreatedSchema,
701    /// The requested role was created.
702    CreatedRole,
703    /// The requested cluster was created.
704    CreatedCluster,
705    /// The requested cluster replica was created.
706    CreatedClusterReplica,
707    /// The requested index was created.
708    CreatedIndex,
709    /// The requested metric sink was created.
710    CreatedMetricSink,
711    /// The requested introspection subscribe was created.
712    CreatedIntrospectionSubscribe,
713    /// The requested secret was created.
714    CreatedSecret,
715    /// The requested sink was created.
716    CreatedSink,
717    /// The requested source was created.
718    CreatedSource,
719    /// The requested table was created.
720    CreatedTable,
721    /// The requested view was created.
722    CreatedView,
723    /// The requested views were created.
724    CreatedViews,
725    /// The requested materialized view was created.
726    CreatedMaterializedView,
727    /// The requested type was created.
728    CreatedType,
729    /// The requested network policy was created.
730    CreatedNetworkPolicy,
731    /// The requested prepared statement was removed.
732    Deallocate { all: bool },
733    /// The requested cursor was declared.
734    DeclaredCursor,
735    /// The specified number of rows were deleted from the requested table.
736    Deleted(usize),
737    /// The temporary objects associated with the session have been discarded.
738    DiscardedTemp,
739    /// All state associated with the session has been discarded.
740    DiscardedAll {
741        /// Session parameters that changed because the session was reset.
742        params: BTreeMap<&'static str, String>,
743    },
744    /// The requested object was dropped.
745    DroppedObject(ObjectType),
746    /// The requested objects were dropped.
747    DroppedOwned,
748    /// The provided query was empty.
749    EmptyQuery,
750    /// Fetch results from a cursor.
751    Fetch {
752        /// The name of the cursor from which to fetch results.
753        name: String,
754        /// The number of results to fetch.
755        count: Option<FetchDirection>,
756        /// How long to wait for results to arrive.
757        timeout: ExecuteTimeout,
758        ctx_extra: ExecuteContextGuard,
759    },
760    /// The requested privilege was granted.
761    GrantedPrivilege,
762    /// The requested role was granted.
763    GrantedRole,
764    /// The specified number of rows were inserted into the requested table.
765    Inserted(usize),
766    /// The specified prepared statement was created.
767    Prepare,
768    /// A user-requested warning was raised.
769    Raised,
770    /// The requested objects were reassigned.
771    ReassignOwned,
772    /// The requested privilege was revoked.
773    RevokedPrivilege,
774    /// The requested role was revoked.
775    RevokedRole,
776    /// Rows will be delivered via the specified stream.
777    SendingRowsStreaming {
778        #[derivative(Debug = "ignore")]
779        rows: Pin<Box<dyn Stream<Item = PeekResponseUnary> + Send + Sync>>,
780        instance_id: ComputeInstanceId,
781        strategy: StatementExecutionStrategy,
782    },
783    /// Rows are known to be available immediately, and thus the execution is
784    /// considered ended in the coordinator.
785    SendingRowsImmediate {
786        #[derivative(Debug = "ignore")]
787        rows: Box<dyn RowIterator + Send + Sync>,
788    },
789    /// The specified variable was set to a new value.
790    SetVariable {
791        name: String,
792        /// Whether the operation was a `RESET` rather than a set.
793        reset: bool,
794    },
795    /// A new transaction was started.
796    StartedTransaction,
797    /// Updates to the requested source or view will be streamed to the
798    /// contained receiver.
799    Subscribing {
800        #[derivative(Debug = "ignore")]
801        rx: RowBatchStream,
802        ctx_extra: ExecuteContextGuard,
803        instance_id: ComputeInstanceId,
804    },
805    /// The active transaction committed.
806    TransactionCommitted {
807        /// Session parameters that changed because the transaction ended.
808        params: BTreeMap<&'static str, String>,
809    },
810    /// The active transaction rolled back.
811    TransactionRolledBack {
812        /// Session parameters that changed because the transaction ended.
813        params: BTreeMap<&'static str, String>,
814    },
815    /// The specified number of rows were updated in the requested table.
816    Updated(usize),
817    /// A connection was validated.
818    ValidatedConnection,
819}
820
821impl TryFrom<&Statement<Raw>> for ExecuteResponse {
822    type Error = ();
823
824    /// Returns Ok if this Statement always produces a single, trivial ExecuteResponse.
825    fn try_from(stmt: &Statement<Raw>) -> Result<Self, Self::Error> {
826        let resp_kinds = Plan::generated_from(&stmt.into())
827            .iter()
828            .map(ExecuteResponse::generated_from)
829            .flatten()
830            .cloned()
831            .collect::<BTreeSet<ExecuteResponseKind>>();
832        let resps = resp_kinds
833            .iter()
834            .map(|r| (*r).try_into())
835            .collect::<Result<Vec<ExecuteResponse>, _>>();
836        // Check if this statement's possible plans yield exactly one possible ExecuteResponse.
837        if let Ok(resps) = resps {
838            if resps.len() == 1 {
839                return Ok(resps.into_element());
840            }
841        }
842        let resp = match stmt {
843            Statement::DropObjects(DropObjectsStatement { object_type, .. }) => {
844                ExecuteResponse::DroppedObject((*object_type).into())
845            }
846            Statement::AlterObjectRename(AlterObjectRenameStatement { object_type, .. })
847            | Statement::AlterOwner(AlterOwnerStatement { object_type, .. }) => {
848                ExecuteResponse::AlteredObject((*object_type).into())
849            }
850            _ => return Err(()),
851        };
852        // Ensure that if the planner ever adds possible plans we complain here.
853        soft_assert_no_log!(
854            resp_kinds.len() == 1
855                && resp_kinds.first().expect("must exist") == &ExecuteResponseKind::from(&resp),
856            "ExecuteResponses out of sync with planner"
857        );
858        Ok(resp)
859    }
860}
861
862impl TryInto<ExecuteResponse> for ExecuteResponseKind {
863    type Error = ();
864
865    /// Attempts to convert into an ExecuteResponse. Returns an error if not possible without
866    /// actually executing a statement.
867    fn try_into(self) -> Result<ExecuteResponse, Self::Error> {
868        match self {
869            ExecuteResponseKind::AlteredDefaultPrivileges => {
870                Ok(ExecuteResponse::AlteredDefaultPrivileges)
871            }
872            ExecuteResponseKind::AlteredObject => Err(()),
873            ExecuteResponseKind::AlteredRole => Ok(ExecuteResponse::AlteredRole),
874            ExecuteResponseKind::AlteredSystemConfiguration => {
875                Ok(ExecuteResponse::AlteredSystemConfiguration)
876            }
877            ExecuteResponseKind::ClosedCursor => Ok(ExecuteResponse::ClosedCursor),
878            ExecuteResponseKind::Comment => Ok(ExecuteResponse::Comment),
879            ExecuteResponseKind::Copied => Err(()),
880            ExecuteResponseKind::CopyTo => Err(()),
881            ExecuteResponseKind::CopyFrom => Err(()),
882            ExecuteResponseKind::CreatedConnection => Ok(ExecuteResponse::CreatedConnection),
883            ExecuteResponseKind::CreatedDatabase => Ok(ExecuteResponse::CreatedDatabase),
884            ExecuteResponseKind::CreatedSchema => Ok(ExecuteResponse::CreatedSchema),
885            ExecuteResponseKind::CreatedRole => Ok(ExecuteResponse::CreatedRole),
886            ExecuteResponseKind::CreatedCluster => Ok(ExecuteResponse::CreatedCluster),
887            ExecuteResponseKind::CreatedClusterReplica => {
888                Ok(ExecuteResponse::CreatedClusterReplica)
889            }
890            ExecuteResponseKind::CreatedIndex => Ok(ExecuteResponse::CreatedIndex),
891            ExecuteResponseKind::CreatedMetricSink => Ok(ExecuteResponse::CreatedMetricSink),
892            ExecuteResponseKind::CreatedSecret => Ok(ExecuteResponse::CreatedSecret),
893            ExecuteResponseKind::CreatedSink => Ok(ExecuteResponse::CreatedSink),
894            ExecuteResponseKind::CreatedSource => Ok(ExecuteResponse::CreatedSource),
895            ExecuteResponseKind::CreatedTable => Ok(ExecuteResponse::CreatedTable),
896            ExecuteResponseKind::CreatedView => Ok(ExecuteResponse::CreatedView),
897            ExecuteResponseKind::CreatedViews => Ok(ExecuteResponse::CreatedViews),
898            ExecuteResponseKind::CreatedMaterializedView => {
899                Ok(ExecuteResponse::CreatedMaterializedView)
900            }
901            ExecuteResponseKind::CreatedNetworkPolicy => Ok(ExecuteResponse::CreatedNetworkPolicy),
902            ExecuteResponseKind::CreatedType => Ok(ExecuteResponse::CreatedType),
903            ExecuteResponseKind::Deallocate => Err(()),
904            ExecuteResponseKind::DeclaredCursor => Ok(ExecuteResponse::DeclaredCursor),
905            ExecuteResponseKind::Deleted => Err(()),
906            ExecuteResponseKind::DiscardedTemp => Ok(ExecuteResponse::DiscardedTemp),
907            ExecuteResponseKind::DiscardedAll => Err(()),
908            ExecuteResponseKind::DroppedObject => Err(()),
909            ExecuteResponseKind::DroppedOwned => Ok(ExecuteResponse::DroppedOwned),
910            ExecuteResponseKind::EmptyQuery => Ok(ExecuteResponse::EmptyQuery),
911            ExecuteResponseKind::Fetch => Err(()),
912            ExecuteResponseKind::GrantedPrivilege => Ok(ExecuteResponse::GrantedPrivilege),
913            ExecuteResponseKind::GrantedRole => Ok(ExecuteResponse::GrantedRole),
914            ExecuteResponseKind::Inserted => Err(()),
915            ExecuteResponseKind::Prepare => Ok(ExecuteResponse::Prepare),
916            ExecuteResponseKind::Raised => Ok(ExecuteResponse::Raised),
917            ExecuteResponseKind::ReassignOwned => Ok(ExecuteResponse::ReassignOwned),
918            ExecuteResponseKind::RevokedPrivilege => Ok(ExecuteResponse::RevokedPrivilege),
919            ExecuteResponseKind::RevokedRole => Ok(ExecuteResponse::RevokedRole),
920            ExecuteResponseKind::SetVariable => Err(()),
921            ExecuteResponseKind::StartedTransaction => Ok(ExecuteResponse::StartedTransaction),
922            ExecuteResponseKind::Subscribing => Err(()),
923            ExecuteResponseKind::TransactionCommitted => Err(()),
924            ExecuteResponseKind::TransactionRolledBack => Err(()),
925            ExecuteResponseKind::Updated => Err(()),
926            ExecuteResponseKind::ValidatedConnection => Ok(ExecuteResponse::ValidatedConnection),
927            ExecuteResponseKind::SendingRowsStreaming => Err(()),
928            ExecuteResponseKind::SendingRowsImmediate => Err(()),
929            ExecuteResponseKind::CreatedIntrospectionSubscribe => {
930                Ok(ExecuteResponse::CreatedIntrospectionSubscribe)
931            }
932        }
933    }
934}
935
936impl ExecuteResponse {
937    pub fn tag(&self) -> Option<String> {
938        use ExecuteResponse::*;
939        match self {
940            AlteredDefaultPrivileges => Some("ALTER DEFAULT PRIVILEGES".into()),
941            AlteredObject(o) => Some(format!("ALTER {}", o)),
942            AlteredRole => Some("ALTER ROLE".into()),
943            AlteredSystemConfiguration => Some("ALTER SYSTEM".into()),
944            ClosedCursor => Some("CLOSE CURSOR".into()),
945            Comment => Some("COMMENT".into()),
946            Copied(n) => Some(format!("COPY {}", n)),
947            CopyTo { .. } => None,
948            CopyFrom { .. } => None,
949            CreatedConnection { .. } => Some("CREATE CONNECTION".into()),
950            CreatedDatabase { .. } => Some("CREATE DATABASE".into()),
951            CreatedSchema { .. } => Some("CREATE SCHEMA".into()),
952            CreatedRole => Some("CREATE ROLE".into()),
953            CreatedCluster { .. } => Some("CREATE CLUSTER".into()),
954            CreatedClusterReplica { .. } => Some("CREATE CLUSTER REPLICA".into()),
955            CreatedIndex { .. } => Some("CREATE INDEX".into()),
956            CreatedMetricSink { .. } => Some("CREATE METRIC SINK".into()),
957            CreatedSecret { .. } => Some("CREATE SECRET".into()),
958            CreatedSink { .. } => Some("CREATE SINK".into()),
959            CreatedSource { .. } => Some("CREATE SOURCE".into()),
960            CreatedTable { .. } => Some("CREATE TABLE".into()),
961            CreatedView { .. } => Some("CREATE VIEW".into()),
962            CreatedViews { .. } => Some("CREATE VIEWS".into()),
963            CreatedMaterializedView { .. } => Some("CREATE MATERIALIZED VIEW".into()),
964            CreatedType => Some("CREATE TYPE".into()),
965            CreatedNetworkPolicy => Some("CREATE NETWORKPOLICY".into()),
966            Deallocate { all } => Some(format!("DEALLOCATE{}", if *all { " ALL" } else { "" })),
967            DeclaredCursor => Some("DECLARE CURSOR".into()),
968            Deleted(n) => Some(format!("DELETE {}", n)),
969            DiscardedTemp => Some("DISCARD TEMP".into()),
970            DiscardedAll { .. } => Some("DISCARD ALL".into()),
971            DroppedObject(o) => Some(format!("DROP {o}")),
972            DroppedOwned => Some("DROP OWNED".into()),
973            EmptyQuery => None,
974            Fetch { .. } => None,
975            GrantedPrivilege => Some("GRANT".into()),
976            GrantedRole => Some("GRANT ROLE".into()),
977            Inserted(n) => {
978                // "On successful completion, an INSERT command returns a
979                // command tag of the form `INSERT <oid> <count>`."
980                //     -- https://www.postgresql.org/docs/11/sql-insert.html
981                //
982                // OIDs are a PostgreSQL-specific historical quirk, but we
983                // can return a 0 OID to indicate that the table does not
984                // have OIDs.
985                Some(format!("INSERT 0 {}", n))
986            }
987            Prepare => Some("PREPARE".into()),
988            Raised => Some("RAISE".into()),
989            ReassignOwned => Some("REASSIGN OWNED".into()),
990            RevokedPrivilege => Some("REVOKE".into()),
991            RevokedRole => Some("REVOKE ROLE".into()),
992            SendingRowsStreaming { .. } | SendingRowsImmediate { .. } => None,
993            SetVariable { reset: true, .. } => Some("RESET".into()),
994            SetVariable { reset: false, .. } => Some("SET".into()),
995            StartedTransaction { .. } => Some("BEGIN".into()),
996            Subscribing { .. } => None,
997            TransactionCommitted { .. } => Some("COMMIT".into()),
998            TransactionRolledBack { .. } => Some("ROLLBACK".into()),
999            Updated(n) => Some(format!("UPDATE {}", n)),
1000            ValidatedConnection => Some("VALIDATE CONNECTION".into()),
1001            CreatedIntrospectionSubscribe => Some("CREATE INTROSPECTION SUBSCRIBE".into()),
1002        }
1003    }
1004
1005    /// Expresses which [`PlanKind`] generate which set of [`ExecuteResponseKind`].
1006    /// `ExecuteResponseKind::Canceled` could be generated at any point as well, but that is
1007    /// excluded from this function.
1008    pub fn generated_from(plan: &PlanKind) -> &'static [ExecuteResponseKind] {
1009        use ExecuteResponseKind::*;
1010        use PlanKind::*;
1011
1012        match plan {
1013            AbortTransaction => &[TransactionRolledBack],
1014            AlterClusterRename
1015            | AlterClusterSwap
1016            | AlterCluster
1017            | AlterClusterReplicaRename
1018            | AlterOwner
1019            | AlterItemRename
1020            | AlterRetainHistory
1021            | AlterSourceTimestampInterval
1022            | AlterNoop
1023            | AlterSchemaRename
1024            | AlterSchemaSwap
1025            | AlterSecret
1026            | AlterConnection
1027            | AlterSource
1028            | AlterSink
1029            | AlterTableAddColumn
1030            | AlterMaterializedViewApplyReplacement
1031            | AlterNetworkPolicy => &[AlteredObject],
1032            AlterDefaultPrivileges => &[AlteredDefaultPrivileges],
1033            AlterSetCluster => &[AlteredObject],
1034            AlterRole => &[AlteredRole],
1035            AlterSystemSet | AlterSystemReset | AlterSystemResetAll => {
1036                &[AlteredSystemConfiguration]
1037            }
1038            Close => &[ClosedCursor],
1039            PlanKind::CopyFrom => &[ExecuteResponseKind::CopyFrom, ExecuteResponseKind::Copied],
1040            PlanKind::CopyTo => &[ExecuteResponseKind::Copied],
1041            PlanKind::Comment => &[ExecuteResponseKind::Comment],
1042            CommitTransaction => &[TransactionCommitted, TransactionRolledBack],
1043            CreateConnection => &[CreatedConnection],
1044            CreateDatabase => &[CreatedDatabase],
1045            CreateSchema => &[CreatedSchema],
1046            CreateRole => &[CreatedRole],
1047            CreateCluster => &[CreatedCluster],
1048            CreateClusterReplica => &[CreatedClusterReplica],
1049            CreateSource | CreateSources => &[CreatedSource],
1050            CreateSecret => &[CreatedSecret],
1051            CreateSink => &[CreatedSink],
1052            CreateTable => &[CreatedTable],
1053            CreateView => &[CreatedView],
1054            CreateMaterializedView => &[CreatedMaterializedView],
1055            CreateIndex => &[CreatedIndex],
1056            CreateMetricSink => &[CreatedMetricSink],
1057            CreateType => &[CreatedType],
1058            PlanKind::Deallocate => &[ExecuteResponseKind::Deallocate],
1059            CreateNetworkPolicy => &[CreatedNetworkPolicy],
1060            Declare => &[DeclaredCursor],
1061            DiscardTemp => &[DiscardedTemp],
1062            DiscardAll => &[DiscardedAll],
1063            DropObjects => &[DroppedObject],
1064            DropOwned => &[DroppedOwned],
1065            PlanKind::EmptyQuery => &[ExecuteResponseKind::EmptyQuery],
1066            ExplainPlan | ExplainPushdown | ExplainTimestamp | Select | ShowAllVariables
1067            | ShowCreate | ShowColumns | ShowVariable | InspectShard | ExplainSinkSchema => &[
1068                ExecuteResponseKind::CopyTo,
1069                SendingRowsStreaming,
1070                SendingRowsImmediate,
1071            ],
1072            Execute | ReadThenWrite => &[
1073                Deleted,
1074                Inserted,
1075                SendingRowsStreaming,
1076                SendingRowsImmediate,
1077                Updated,
1078            ],
1079            PlanKind::Fetch => &[ExecuteResponseKind::Fetch],
1080            GrantPrivileges => &[GrantedPrivilege],
1081            GrantRole => &[GrantedRole],
1082            Insert => &[Inserted, SendingRowsImmediate],
1083            PlanKind::Prepare => &[ExecuteResponseKind::Prepare],
1084            PlanKind::Raise => &[ExecuteResponseKind::Raised],
1085            PlanKind::ReassignOwned => &[ExecuteResponseKind::ReassignOwned],
1086            RevokePrivileges => &[RevokedPrivilege],
1087            RevokeRole => &[RevokedRole],
1088            PlanKind::SetVariable | ResetVariable | PlanKind::SetTransaction => {
1089                &[ExecuteResponseKind::SetVariable]
1090            }
1091            PlanKind::Subscribe => &[Subscribing, ExecuteResponseKind::CopyTo],
1092            StartTransaction => &[StartedTransaction],
1093            SideEffectingFunc => &[SendingRowsStreaming, SendingRowsImmediate],
1094            ValidateConnection => &[ExecuteResponseKind::ValidatedConnection],
1095        }
1096    }
1097}
1098
1099/// This implementation is meant to ensure that we maintain updated information
1100/// about which types of `ExecuteResponse`s are permitted to be sent, which will
1101/// be a function of which plan we're executing.
1102impl Transmittable for ExecuteResponse {
1103    type Allowed = ExecuteResponseKind;
1104    fn to_allowed(&self) -> Self::Allowed {
1105        ExecuteResponseKind::from(self)
1106    }
1107}
mz_adapter/command.rs

mz_adapter/
command.rs
