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    /// RoleId for the user.
579    pub role_id: RoleId,
580    /// The role's superuser attribute in the Catalog.
581    /// This attribute is None for Cloud. Cloud is able
582    /// to derive the role's superuser status from
583    /// external_metadata_rx.
584    pub superuser_attribute: SuperuserAttribute,
585    /// A future that completes when all necessary Builtin Table writes have completed.
586    #[derivative(Debug = "ignore")]
587    pub write_notify: BuiltinTableAppendNotify,
588    /// Map of (name, VarInput::Flat) tuples of session default variables that should be set.
589    pub session_defaults: BTreeMap<String, OwnedVarInput>,
590    pub catalog: Arc<Catalog>,
591    pub storage_collections:
592        Arc<dyn mz_storage_client::storage_collections::StorageCollections + Send + Sync>,
593    pub transient_id_gen: Arc<TransientIdGen>,
594    pub optimizer_metrics: OptimizerMetrics,
595    pub persist_client: PersistClient,
596    pub statement_logging_frontend: StatementLoggingFrontend,
597    /// Semaphore for limiting concurrent OCC (optimistic concurrency control)
598    /// write operations.
599    pub occ_write_semaphore: Arc<Semaphore>,
600    /// Requests a group commit, which is how the frontend asks for the write
601    /// timeline to advance without having anything to write.
602    pub group_commit_notifier: crate::coord::appends::GroupCommitNotifier,
603    /// Whether the coordinator is in read-only mode (e.g. during 0dt upgrades).
604    /// The frontend path must reject mutations when this is true.
605    pub read_only: bool,
606}
607
608#[derive(Derivative)]
609#[derivative(Debug)]
610pub struct SASLChallengeResponse {
611    pub iteration_count: usize,
612    /// Base64-encoded salt for the SASL challenge.
613    pub salt: String,
614    pub nonce: String,
615}
616
617#[derive(Derivative)]
618#[derivative(Debug)]
619pub struct SASLVerifyProofResponse {
620    pub verifier: String,
621}
622
623// Facile implementation for `StartupResponse`, which does not use the `allowed`
624// feature of `ClientTransmitter`.
625impl Transmittable for StartupResponse {
626    type Allowed = bool;
627    fn to_allowed(&self) -> Self::Allowed {
628        true
629    }
630}
631
632/// The response to [`SessionClient::dump_catalog`](crate::SessionClient::dump_catalog).
633#[derive(Debug, Clone)]
634pub struct CatalogDump(String);
635
636impl CatalogDump {
637    pub fn new(raw: String) -> Self {
638        CatalogDump(raw)
639    }
640
641    pub fn into_string(self) -> String {
642        self.0
643    }
644}
645
646impl Transmittable for CatalogDump {
647    type Allowed = bool;
648    fn to_allowed(&self) -> Self::Allowed {
649        true
650    }
651}
652
653impl Transmittable for SystemVars {
654    type Allowed = bool;
655    fn to_allowed(&self) -> Self::Allowed {
656        true
657    }
658}
659
660/// The response to [`SessionClient::execute`](crate::SessionClient::execute).
661#[derive(EnumKind, Derivative)]
662#[derivative(Debug)]
663#[enum_kind(ExecuteResponseKind, derive(PartialOrd, Ord))]
664pub enum ExecuteResponse {
665    /// The default privileges were altered.
666    AlteredDefaultPrivileges,
667    /// The requested object was altered.
668    AlteredObject(ObjectType),
669    /// The role was altered.
670    AlteredRole,
671    /// The system configuration was altered.
672    AlteredSystemConfiguration,
673    /// The requested cursor was closed.
674    ClosedCursor,
675    /// The provided comment was created.
676    Comment,
677    /// The specified number of rows were copied into the requested output.
678    Copied(usize),
679    /// The response for a COPY TO STDOUT query.
680    CopyTo {
681        format: mz_sql::plan::CopyFormat,
682        resp: Box<ExecuteResponse>,
683    },
684    CopyFrom {
685        /// Table we're copying into.
686        target_id: CatalogItemId,
687        /// Human-readable full name of the target table.
688        target_name: String,
689        columns: Vec<ColumnIndex>,
690        params: CopyFormatParams<'static>,
691        ctx_extra: ExecuteContextGuard,
692    },
693    /// The requested connection was created.
694    CreatedConnection,
695    /// The requested database was created.
696    CreatedDatabase,
697    /// The requested schema was created.
698    CreatedSchema,
699    /// The requested role was created.
700    CreatedRole,
701    /// The requested cluster was created.
702    CreatedCluster,
703    /// The requested cluster replica was created.
704    CreatedClusterReplica,
705    /// The requested index was created.
706    CreatedIndex,
707    /// The requested metric sink was created.
708    CreatedMetricSink,
709    /// The requested introspection subscribe was created.
710    CreatedIntrospectionSubscribe,
711    /// The requested secret was created.
712    CreatedSecret,
713    /// The requested sink was created.
714    CreatedSink,
715    /// The requested source was created.
716    CreatedSource,
717    /// The requested table was created.
718    CreatedTable,
719    /// The requested view was created.
720    CreatedView,
721    /// The requested views were created.
722    CreatedViews,
723    /// The requested materialized view was created.
724    CreatedMaterializedView,
725    /// The requested type was created.
726    CreatedType,
727    /// The requested network policy was created.
728    CreatedNetworkPolicy,
729    /// The requested prepared statement was removed.
730    Deallocate { all: bool },
731    /// The requested cursor was declared.
732    DeclaredCursor,
733    /// The specified number of rows were deleted from the requested table.
734    Deleted(usize),
735    /// The temporary objects associated with the session have been discarded.
736    DiscardedTemp,
737    /// All state associated with the session has been discarded.
738    DiscardedAll {
739        /// Session parameters that changed because the session was reset.
740        params: BTreeMap<&'static str, String>,
741    },
742    /// The requested object was dropped.
743    DroppedObject(ObjectType),
744    /// The requested objects were dropped.
745    DroppedOwned,
746    /// The provided query was empty.
747    EmptyQuery,
748    /// Fetch results from a cursor.
749    Fetch {
750        /// The name of the cursor from which to fetch results.
751        name: String,
752        /// The number of results to fetch.
753        count: Option<FetchDirection>,
754        /// How long to wait for results to arrive.
755        timeout: ExecuteTimeout,
756        ctx_extra: ExecuteContextGuard,
757    },
758    /// The requested privilege was granted.
759    GrantedPrivilege,
760    /// The requested role was granted.
761    GrantedRole,
762    /// The specified number of rows were inserted into the requested table.
763    Inserted(usize),
764    /// The specified prepared statement was created.
765    Prepare,
766    /// A user-requested warning was raised.
767    Raised,
768    /// The requested objects were reassigned.
769    ReassignOwned,
770    /// The requested privilege was revoked.
771    RevokedPrivilege,
772    /// The requested role was revoked.
773    RevokedRole,
774    /// Rows will be delivered via the specified stream.
775    SendingRowsStreaming {
776        #[derivative(Debug = "ignore")]
777        rows: Pin<Box<dyn Stream<Item = PeekResponseUnary> + Send + Sync>>,
778        instance_id: ComputeInstanceId,
779        strategy: StatementExecutionStrategy,
780    },
781    /// Rows are known to be available immediately, and thus the execution is
782    /// considered ended in the coordinator.
783    SendingRowsImmediate {
784        #[derivative(Debug = "ignore")]
785        rows: Box<dyn RowIterator + Send + Sync>,
786    },
787    /// The specified variable was set to a new value.
788    SetVariable {
789        name: String,
790        /// Whether the operation was a `RESET` rather than a set.
791        reset: bool,
792    },
793    /// A new transaction was started.
794    StartedTransaction,
795    /// Updates to the requested source or view will be streamed to the
796    /// contained receiver.
797    Subscribing {
798        #[derivative(Debug = "ignore")]
799        rx: RowBatchStream,
800        ctx_extra: ExecuteContextGuard,
801        instance_id: ComputeInstanceId,
802    },
803    /// The active transaction committed.
804    TransactionCommitted {
805        /// Session parameters that changed because the transaction ended.
806        params: BTreeMap<&'static str, String>,
807    },
808    /// The active transaction rolled back.
809    TransactionRolledBack {
810        /// Session parameters that changed because the transaction ended.
811        params: BTreeMap<&'static str, String>,
812    },
813    /// The specified number of rows were updated in the requested table.
814    Updated(usize),
815    /// A connection was validated.
816    ValidatedConnection,
817}
818
819impl TryFrom<&Statement<Raw>> for ExecuteResponse {
820    type Error = ();
821
822    /// Returns Ok if this Statement always produces a single, trivial ExecuteResponse.
823    fn try_from(stmt: &Statement<Raw>) -> Result<Self, Self::Error> {
824        let resp_kinds = Plan::generated_from(&stmt.into())
825            .iter()
826            .map(ExecuteResponse::generated_from)
827            .flatten()
828            .cloned()
829            .collect::<BTreeSet<ExecuteResponseKind>>();
830        let resps = resp_kinds
831            .iter()
832            .map(|r| (*r).try_into())
833            .collect::<Result<Vec<ExecuteResponse>, _>>();
834        // Check if this statement's possible plans yield exactly one possible ExecuteResponse.
835        if let Ok(resps) = resps {
836            if resps.len() == 1 {
837                return Ok(resps.into_element());
838            }
839        }
840        let resp = match stmt {
841            Statement::DropObjects(DropObjectsStatement { object_type, .. }) => {
842                ExecuteResponse::DroppedObject((*object_type).into())
843            }
844            Statement::AlterObjectRename(AlterObjectRenameStatement { object_type, .. })
845            | Statement::AlterOwner(AlterOwnerStatement { object_type, .. }) => {
846                ExecuteResponse::AlteredObject((*object_type).into())
847            }
848            _ => return Err(()),
849        };
850        // Ensure that if the planner ever adds possible plans we complain here.
851        soft_assert_no_log!(
852            resp_kinds.len() == 1
853                && resp_kinds.first().expect("must exist") == &ExecuteResponseKind::from(&resp),
854            "ExecuteResponses out of sync with planner"
855        );
856        Ok(resp)
857    }
858}
859
860impl TryInto<ExecuteResponse> for ExecuteResponseKind {
861    type Error = ();
862
863    /// Attempts to convert into an ExecuteResponse. Returns an error if not possible without
864    /// actually executing a statement.
865    fn try_into(self) -> Result<ExecuteResponse, Self::Error> {
866        match self {
867            ExecuteResponseKind::AlteredDefaultPrivileges => {
868                Ok(ExecuteResponse::AlteredDefaultPrivileges)
869            }
870            ExecuteResponseKind::AlteredObject => Err(()),
871            ExecuteResponseKind::AlteredRole => Ok(ExecuteResponse::AlteredRole),
872            ExecuteResponseKind::AlteredSystemConfiguration => {
873                Ok(ExecuteResponse::AlteredSystemConfiguration)
874            }
875            ExecuteResponseKind::ClosedCursor => Ok(ExecuteResponse::ClosedCursor),
876            ExecuteResponseKind::Comment => Ok(ExecuteResponse::Comment),
877            ExecuteResponseKind::Copied => Err(()),
878            ExecuteResponseKind::CopyTo => Err(()),
879            ExecuteResponseKind::CopyFrom => Err(()),
880            ExecuteResponseKind::CreatedConnection => Ok(ExecuteResponse::CreatedConnection),
881            ExecuteResponseKind::CreatedDatabase => Ok(ExecuteResponse::CreatedDatabase),
882            ExecuteResponseKind::CreatedSchema => Ok(ExecuteResponse::CreatedSchema),
883            ExecuteResponseKind::CreatedRole => Ok(ExecuteResponse::CreatedRole),
884            ExecuteResponseKind::CreatedCluster => Ok(ExecuteResponse::CreatedCluster),
885            ExecuteResponseKind::CreatedClusterReplica => {
886                Ok(ExecuteResponse::CreatedClusterReplica)
887            }
888            ExecuteResponseKind::CreatedIndex => Ok(ExecuteResponse::CreatedIndex),
889            ExecuteResponseKind::CreatedMetricSink => Ok(ExecuteResponse::CreatedMetricSink),
890            ExecuteResponseKind::CreatedSecret => Ok(ExecuteResponse::CreatedSecret),
891            ExecuteResponseKind::CreatedSink => Ok(ExecuteResponse::CreatedSink),
892            ExecuteResponseKind::CreatedSource => Ok(ExecuteResponse::CreatedSource),
893            ExecuteResponseKind::CreatedTable => Ok(ExecuteResponse::CreatedTable),
894            ExecuteResponseKind::CreatedView => Ok(ExecuteResponse::CreatedView),
895            ExecuteResponseKind::CreatedViews => Ok(ExecuteResponse::CreatedViews),
896            ExecuteResponseKind::CreatedMaterializedView => {
897                Ok(ExecuteResponse::CreatedMaterializedView)
898            }
899            ExecuteResponseKind::CreatedNetworkPolicy => Ok(ExecuteResponse::CreatedNetworkPolicy),
900            ExecuteResponseKind::CreatedType => Ok(ExecuteResponse::CreatedType),
901            ExecuteResponseKind::Deallocate => Err(()),
902            ExecuteResponseKind::DeclaredCursor => Ok(ExecuteResponse::DeclaredCursor),
903            ExecuteResponseKind::Deleted => Err(()),
904            ExecuteResponseKind::DiscardedTemp => Ok(ExecuteResponse::DiscardedTemp),
905            ExecuteResponseKind::DiscardedAll => Err(()),
906            ExecuteResponseKind::DroppedObject => Err(()),
907            ExecuteResponseKind::DroppedOwned => Ok(ExecuteResponse::DroppedOwned),
908            ExecuteResponseKind::EmptyQuery => Ok(ExecuteResponse::EmptyQuery),
909            ExecuteResponseKind::Fetch => Err(()),
910            ExecuteResponseKind::GrantedPrivilege => Ok(ExecuteResponse::GrantedPrivilege),
911            ExecuteResponseKind::GrantedRole => Ok(ExecuteResponse::GrantedRole),
912            ExecuteResponseKind::Inserted => Err(()),
913            ExecuteResponseKind::Prepare => Ok(ExecuteResponse::Prepare),
914            ExecuteResponseKind::Raised => Ok(ExecuteResponse::Raised),
915            ExecuteResponseKind::ReassignOwned => Ok(ExecuteResponse::ReassignOwned),
916            ExecuteResponseKind::RevokedPrivilege => Ok(ExecuteResponse::RevokedPrivilege),
917            ExecuteResponseKind::RevokedRole => Ok(ExecuteResponse::RevokedRole),
918            ExecuteResponseKind::SetVariable => Err(()),
919            ExecuteResponseKind::StartedTransaction => Ok(ExecuteResponse::StartedTransaction),
920            ExecuteResponseKind::Subscribing => Err(()),
921            ExecuteResponseKind::TransactionCommitted => Err(()),
922            ExecuteResponseKind::TransactionRolledBack => Err(()),
923            ExecuteResponseKind::Updated => Err(()),
924            ExecuteResponseKind::ValidatedConnection => Ok(ExecuteResponse::ValidatedConnection),
925            ExecuteResponseKind::SendingRowsStreaming => Err(()),
926            ExecuteResponseKind::SendingRowsImmediate => Err(()),
927            ExecuteResponseKind::CreatedIntrospectionSubscribe => {
928                Ok(ExecuteResponse::CreatedIntrospectionSubscribe)
929            }
930        }
931    }
932}
933
934impl ExecuteResponse {
935    pub fn tag(&self) -> Option<String> {
936        use ExecuteResponse::*;
937        match self {
938            AlteredDefaultPrivileges => Some("ALTER DEFAULT PRIVILEGES".into()),
939            AlteredObject(o) => Some(format!("ALTER {}", o)),
940            AlteredRole => Some("ALTER ROLE".into()),
941            AlteredSystemConfiguration => Some("ALTER SYSTEM".into()),
942            ClosedCursor => Some("CLOSE CURSOR".into()),
943            Comment => Some("COMMENT".into()),
944            Copied(n) => Some(format!("COPY {}", n)),
945            CopyTo { .. } => None,
946            CopyFrom { .. } => None,
947            CreatedConnection { .. } => Some("CREATE CONNECTION".into()),
948            CreatedDatabase { .. } => Some("CREATE DATABASE".into()),
949            CreatedSchema { .. } => Some("CREATE SCHEMA".into()),
950            CreatedRole => Some("CREATE ROLE".into()),
951            CreatedCluster { .. } => Some("CREATE CLUSTER".into()),
952            CreatedClusterReplica { .. } => Some("CREATE CLUSTER REPLICA".into()),
953            CreatedIndex { .. } => Some("CREATE INDEX".into()),
954            CreatedMetricSink { .. } => Some("CREATE METRIC SINK".into()),
955            CreatedSecret { .. } => Some("CREATE SECRET".into()),
956            CreatedSink { .. } => Some("CREATE SINK".into()),
957            CreatedSource { .. } => Some("CREATE SOURCE".into()),
958            CreatedTable { .. } => Some("CREATE TABLE".into()),
959            CreatedView { .. } => Some("CREATE VIEW".into()),
960            CreatedViews { .. } => Some("CREATE VIEWS".into()),
961            CreatedMaterializedView { .. } => Some("CREATE MATERIALIZED VIEW".into()),
962            CreatedType => Some("CREATE TYPE".into()),
963            CreatedNetworkPolicy => Some("CREATE NETWORKPOLICY".into()),
964            Deallocate { all } => Some(format!("DEALLOCATE{}", if *all { " ALL" } else { "" })),
965            DeclaredCursor => Some("DECLARE CURSOR".into()),
966            Deleted(n) => Some(format!("DELETE {}", n)),
967            DiscardedTemp => Some("DISCARD TEMP".into()),
968            DiscardedAll { .. } => Some("DISCARD ALL".into()),
969            DroppedObject(o) => Some(format!("DROP {o}")),
970            DroppedOwned => Some("DROP OWNED".into()),
971            EmptyQuery => None,
972            Fetch { .. } => None,
973            GrantedPrivilege => Some("GRANT".into()),
974            GrantedRole => Some("GRANT ROLE".into()),
975            Inserted(n) => {
976                // "On successful completion, an INSERT command returns a
977                // command tag of the form `INSERT <oid> <count>`."
978                //     -- https://www.postgresql.org/docs/11/sql-insert.html
979                //
980                // OIDs are a PostgreSQL-specific historical quirk, but we
981                // can return a 0 OID to indicate that the table does not
982                // have OIDs.
983                Some(format!("INSERT 0 {}", n))
984            }
985            Prepare => Some("PREPARE".into()),
986            Raised => Some("RAISE".into()),
987            ReassignOwned => Some("REASSIGN OWNED".into()),
988            RevokedPrivilege => Some("REVOKE".into()),
989            RevokedRole => Some("REVOKE ROLE".into()),
990            SendingRowsStreaming { .. } | SendingRowsImmediate { .. } => None,
991            SetVariable { reset: true, .. } => Some("RESET".into()),
992            SetVariable { reset: false, .. } => Some("SET".into()),
993            StartedTransaction { .. } => Some("BEGIN".into()),
994            Subscribing { .. } => None,
995            TransactionCommitted { .. } => Some("COMMIT".into()),
996            TransactionRolledBack { .. } => Some("ROLLBACK".into()),
997            Updated(n) => Some(format!("UPDATE {}", n)),
998            ValidatedConnection => Some("VALIDATE CONNECTION".into()),
999            CreatedIntrospectionSubscribe => Some("CREATE INTROSPECTION SUBSCRIBE".into()),
1000        }
1001    }
1002
1003    /// Expresses which [`PlanKind`] generate which set of [`ExecuteResponseKind`].
1004    /// `ExecuteResponseKind::Canceled` could be generated at any point as well, but that is
1005    /// excluded from this function.
1006    pub fn generated_from(plan: &PlanKind) -> &'static [ExecuteResponseKind] {
1007        use ExecuteResponseKind::*;
1008        use PlanKind::*;
1009
1010        match plan {
1011            AbortTransaction => &[TransactionRolledBack],
1012            AlterClusterRename
1013            | AlterClusterSwap
1014            | AlterCluster
1015            | AlterClusterReplicaRename
1016            | AlterOwner
1017            | AlterItemRename
1018            | AlterRetainHistory
1019            | AlterSourceTimestampInterval
1020            | AlterNoop
1021            | AlterSchemaRename
1022            | AlterSchemaSwap
1023            | AlterSecret
1024            | AlterConnection
1025            | AlterSource
1026            | AlterSink
1027            | AlterTableAddColumn
1028            | AlterMaterializedViewApplyReplacement
1029            | AlterNetworkPolicy => &[AlteredObject],
1030            AlterDefaultPrivileges => &[AlteredDefaultPrivileges],
1031            AlterSetCluster => &[AlteredObject],
1032            AlterRole => &[AlteredRole],
1033            AlterSystemSet | AlterSystemReset | AlterSystemResetAll => {
1034                &[AlteredSystemConfiguration]
1035            }
1036            Close => &[ClosedCursor],
1037            PlanKind::CopyFrom => &[ExecuteResponseKind::CopyFrom, ExecuteResponseKind::Copied],
1038            PlanKind::CopyTo => &[ExecuteResponseKind::Copied],
1039            PlanKind::Comment => &[ExecuteResponseKind::Comment],
1040            CommitTransaction => &[TransactionCommitted, TransactionRolledBack],
1041            CreateConnection => &[CreatedConnection],
1042            CreateDatabase => &[CreatedDatabase],
1043            CreateSchema => &[CreatedSchema],
1044            CreateRole => &[CreatedRole],
1045            CreateCluster => &[CreatedCluster],
1046            CreateClusterReplica => &[CreatedClusterReplica],
1047            CreateSource | CreateSources => &[CreatedSource],
1048            CreateSecret => &[CreatedSecret],
1049            CreateSink => &[CreatedSink],
1050            CreateTable => &[CreatedTable],
1051            CreateView => &[CreatedView],
1052            CreateMaterializedView => &[CreatedMaterializedView],
1053            CreateIndex => &[CreatedIndex],
1054            CreateMetricSink => &[CreatedMetricSink],
1055            CreateType => &[CreatedType],
1056            PlanKind::Deallocate => &[ExecuteResponseKind::Deallocate],
1057            CreateNetworkPolicy => &[CreatedNetworkPolicy],
1058            Declare => &[DeclaredCursor],
1059            DiscardTemp => &[DiscardedTemp],
1060            DiscardAll => &[DiscardedAll],
1061            DropObjects => &[DroppedObject],
1062            DropOwned => &[DroppedOwned],
1063            PlanKind::EmptyQuery => &[ExecuteResponseKind::EmptyQuery],
1064            ExplainPlan | ExplainPushdown | ExplainTimestamp | Select | ShowAllVariables
1065            | ShowCreate | ShowColumns | ShowVariable | InspectShard | ExplainSinkSchema => &[
1066                ExecuteResponseKind::CopyTo,
1067                SendingRowsStreaming,
1068                SendingRowsImmediate,
1069            ],
1070            Execute | ReadThenWrite => &[
1071                Deleted,
1072                Inserted,
1073                SendingRowsStreaming,
1074                SendingRowsImmediate,
1075                Updated,
1076            ],
1077            PlanKind::Fetch => &[ExecuteResponseKind::Fetch],
1078            GrantPrivileges => &[GrantedPrivilege],
1079            GrantRole => &[GrantedRole],
1080            Insert => &[Inserted, SendingRowsImmediate],
1081            PlanKind::Prepare => &[ExecuteResponseKind::Prepare],
1082            PlanKind::Raise => &[ExecuteResponseKind::Raised],
1083            PlanKind::ReassignOwned => &[ExecuteResponseKind::ReassignOwned],
1084            RevokePrivileges => &[RevokedPrivilege],
1085            RevokeRole => &[RevokedRole],
1086            PlanKind::SetVariable | ResetVariable | PlanKind::SetTransaction => {
1087                &[ExecuteResponseKind::SetVariable]
1088            }
1089            PlanKind::Subscribe => &[Subscribing, ExecuteResponseKind::CopyTo],
1090            StartTransaction => &[StartedTransaction],
1091            SideEffectingFunc => &[SendingRowsStreaming, SendingRowsImmediate],
1092            ValidateConnection => &[ExecuteResponseKind::ValidatedConnection],
1093        }
1094    }
1095}
1096
1097/// This implementation is meant to ensure that we maintain updated information
1098/// about which types of `ExecuteResponse`s are permitted to be sent, which will
1099/// be a function of which plan we're executing.
1100impl Transmittable for ExecuteResponse {
1101    type Allowed = ExecuteResponseKind;
1102    fn to_allowed(&self) -> Self::Allowed {
1103        ExecuteResponseKind::from(self)
1104    }
1105}
mz_adapter/command.rs

mz_adapter/
command.rs
