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_compute_types::dataflows::DataflowDescription;
25use mz_controller_types::ClusterId;
26use mz_expr::RowSetFinishing;
27use mz_ore::collections::CollectionExt;
28use mz_ore::soft_assert_no_log;
29use mz_ore::tracing::OpenTelemetryContext;
30use mz_persist_client::PersistClient;
31use mz_pgcopy::CopyFormatParams;
32use mz_repr::global_id::TransientIdGen;
33use mz_repr::role_id::RoleId;
34use mz_repr::{CatalogItemId, ColumnIndex, Diff, GlobalId, Row, RowIterator, SqlRelationType};
35use mz_sql::ast::{FetchDirection, Raw, Statement};
36use mz_sql::catalog::ObjectType;
37use mz_sql::optimizer_metrics::OptimizerMetrics;
38use mz_sql::plan;
39use mz_sql::plan::{ExecuteTimeout, Plan, PlanKind, SideEffectingFunc};
40use mz_sql::session::user::User;
41use mz_sql::session::vars::{OwnedVarInput, SystemVars};
42use mz_sql_parser::ast::{AlterObjectRenameStatement, AlterOwnerStatement, DropObjectsStatement};
43use mz_storage_types::sources::Timeline;
44use mz_timestamp_oracle::TimestampOracle;
45use tokio::sync::{Semaphore, mpsc, oneshot, watch};
46use uuid::Uuid;
47
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 the create-cluster / create-replica paths can resolve a
193    /// new object's scoped overrides synchronously, before the controller
194    /// installs it or its first dataflow is planned, rather than waiting for
195    /// the next sync tick. Sent by the sync loop whenever it (re)initializes the
196    /// frontend. See the scoped feature flags 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: DataflowDescription<mz_compute_types::plan::LirRelationExpr>,
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<DataflowDescription<mz_compute_types::plan::LirRelationExpr>>,
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.
386    /// This is used when `emit_timestamp_notice` is enabled.
387    ExplainTimestamp {
388        conn_id: ConnectionId,
389        session_wall_time: DateTime<Utc>,
390        cluster_id: ClusterId,
391        id_bundle: CollectionIdBundle,
392        determination: TimestampDetermination,
393        tx: oneshot::Sender<TimestampExplanation>,
394    },
395
396    /// Statement logging event from frontend peek sequencing.
397    /// No response channel needed - this is fire-and-forget.
398    FrontendStatementLogging(FrontendStatementLoggingEvent),
399
400    /// Registers a connection-scoped cancellation watch and returns a receiver
401    /// that becomes `true` when cancellation is requested for the connection.
402    ///
403    /// Registration always installs a fresh channel, so the caller cannot
404    /// observe a cancellation aimed at an earlier statement.
405    RegisterConnectionCancelWatch {
406        conn_id: ConnectionId,
407        tx: oneshot::Sender<watch::Receiver<bool>>,
408    },
409
410    /// Creates an internal subscribe, meaning one that writes no
411    /// `mz_subscriptions` row, and returns the response channel. Used by
412    /// frontend-sequenced read-then-write (DELETE/UPDATE/INSERT...SELECT)
413    /// operations via OCC.
414    CreateInternalSubscribe {
415        df_desc: Box<LirDataflowDescription>,
416        cluster_id: ComputeInstanceId,
417        replica_id: Option<ReplicaId>,
418        depends_on: BTreeSet<GlobalId>,
419        as_of: mz_repr::Timestamp,
420        arity: usize,
421        sink_id: GlobalId,
422        conn_id: ConnectionId,
423        session_uuid: Uuid,
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        /// Connection originating the write. Used so the coordinator can
441        /// cancel this pending write if the connection is cancelled before
442        /// the write commits.
443        conn_id: ConnectionId,
444        target_id: CatalogItemId,
445        target_global_id: GlobalId,
446        diffs: Vec<(Row, Diff)>,
447        write_ts: Option<mz_repr::Timestamp>,
448        tx: oneshot::Sender<WriteResult>,
449    },
450
451    /// Drops an internal subscribe. Fire-and-forget, the caller does not wait
452    /// for completion.
453    DropInternalSubscribe {
454        sink_id: GlobalId,
455    },
456}
457
458impl Command {
459    pub fn session(&self) -> Option<&Session> {
460        match self {
461            Command::Execute { session, .. }
462            | Command::Commit { session, .. }
463            | Command::StartCopyFromStdin { session, .. } => Some(session),
464            Command::CancelRequest { .. }
465            | Command::Startup { .. }
466            | Command::AuthenticatePassword { .. }
467            | Command::AuthenticateGetSASLChallenge { .. }
468            | Command::AuthenticateVerifySASLProof { .. }
469            | Command::CheckRoleCanLogin { .. }
470            | Command::CatalogSnapshot { .. }
471            | Command::PrivilegedCancelRequest { .. }
472            | Command::GetWebhook { .. }
473            | Command::Terminate { .. }
474            | Command::GetSystemVars { .. }
475            | Command::SetSystemVars { .. }
476            | Command::UpdateScopedSystemParameters { .. }
477            | Command::InstallScopedSystemParameterFrontend { .. }
478            | Command::RetireExecute { .. }
479            | Command::CheckConsistency { .. }
480            | Command::Dump { .. }
481            | Command::GetComputeInstanceClient { .. }
482            | Command::GetOracle { .. }
483            | Command::DetermineRealTimeRecentTimestamp { .. }
484            | Command::GetTransactionReadHoldsBundle { .. }
485            | Command::StoreTransactionReadHolds { .. }
486            | Command::ExecuteSlowPathPeek { .. }
487            | Command::ExecuteSubscribe { .. }
488            | Command::CopyToPreflight { .. }
489            | Command::ExecuteCopyTo { .. }
490            | Command::ExecuteSideEffectingFunc { .. }
491            | Command::LookupConnection { .. }
492            | Command::RegisterFrontendPeek { .. }
493            | Command::UnregisterFrontendPeek { .. }
494            | Command::ExplainTimestamp { .. }
495            | Command::FrontendStatementLogging(..)
496            | Command::InjectAuditEvents { .. }
497            | Command::RegisterConnectionCancelWatch { .. }
498            | Command::CreateInternalSubscribe { .. }
499            | Command::AttemptWrite { .. }
500            | Command::DropInternalSubscribe { .. } => None,
501        }
502    }
503
504    pub fn session_mut(&mut self) -> Option<&mut Session> {
505        match self {
506            Command::Execute { session, .. }
507            | Command::Commit { session, .. }
508            | Command::StartCopyFromStdin { session, .. } => Some(session),
509            Command::CancelRequest { .. }
510            | Command::Startup { .. }
511            | Command::AuthenticatePassword { .. }
512            | Command::AuthenticateGetSASLChallenge { .. }
513            | Command::AuthenticateVerifySASLProof { .. }
514            | Command::CheckRoleCanLogin { .. }
515            | Command::CatalogSnapshot { .. }
516            | Command::PrivilegedCancelRequest { .. }
517            | Command::GetWebhook { .. }
518            | Command::Terminate { .. }
519            | Command::GetSystemVars { .. }
520            | Command::SetSystemVars { .. }
521            | Command::UpdateScopedSystemParameters { .. }
522            | Command::InstallScopedSystemParameterFrontend { .. }
523            | Command::RetireExecute { .. }
524            | Command::CheckConsistency { .. }
525            | Command::Dump { .. }
526            | Command::GetComputeInstanceClient { .. }
527            | Command::GetOracle { .. }
528            | Command::DetermineRealTimeRecentTimestamp { .. }
529            | Command::GetTransactionReadHoldsBundle { .. }
530            | Command::StoreTransactionReadHolds { .. }
531            | Command::ExecuteSlowPathPeek { .. }
532            | Command::ExecuteSubscribe { .. }
533            | Command::CopyToPreflight { .. }
534            | Command::ExecuteCopyTo { .. }
535            | Command::ExecuteSideEffectingFunc { .. }
536            | Command::LookupConnection { .. }
537            | Command::RegisterFrontendPeek { .. }
538            | Command::UnregisterFrontendPeek { .. }
539            | Command::ExplainTimestamp { .. }
540            | Command::FrontendStatementLogging(..)
541            | Command::InjectAuditEvents { .. }
542            | Command::RegisterConnectionCancelWatch { .. }
543            | Command::CreateInternalSubscribe { .. }
544            | Command::AttemptWrite { .. }
545            | Command::DropInternalSubscribe { .. } => None,
546        }
547    }
548}
549
550#[derive(Debug)]
551pub struct Response<T> {
552    pub result: Result<T, AdapterError>,
553    pub session: Session,
554    pub otel_ctx: OpenTelemetryContext,
555}
556
557#[derive(Debug, Clone, Copy)]
558pub struct SuperuserAttribute(pub Option<bool>);
559
560/// The response to [`Client::startup`](crate::Client::startup).
561#[derive(Derivative)]
562#[derivative(Debug)]
563pub struct StartupResponse {
564    /// RoleId for the user.
565    pub role_id: RoleId,
566    /// The role's superuser attribute in the Catalog.
567    /// This attribute is None for Cloud. Cloud is able
568    /// to derive the role's superuser status from
569    /// external_metadata_rx.
570    pub superuser_attribute: SuperuserAttribute,
571    /// A future that completes when all necessary Builtin Table writes have completed.
572    #[derivative(Debug = "ignore")]
573    pub write_notify: BuiltinTableAppendNotify,
574    /// Map of (name, VarInput::Flat) tuples of session default variables that should be set.
575    pub session_defaults: BTreeMap<String, OwnedVarInput>,
576    pub catalog: Arc<Catalog>,
577    pub storage_collections:
578        Arc<dyn mz_storage_client::storage_collections::StorageCollections + Send + Sync>,
579    pub transient_id_gen: Arc<TransientIdGen>,
580    pub optimizer_metrics: OptimizerMetrics,
581    pub persist_client: PersistClient,
582    pub statement_logging_frontend: StatementLoggingFrontend,
583    /// Semaphore for limiting concurrent OCC (optimistic concurrency control)
584    /// write operations.
585    pub occ_write_semaphore: Arc<Semaphore>,
586    /// Whether frontend OCC read-then-write is enabled (determined once at
587    /// process startup).
588    pub frontend_read_then_write_enabled: bool,
589    /// Requests a group commit, which is how the frontend asks for the write
590    /// timeline to advance without having anything to write.
591    pub group_commit_notifier: crate::coord::appends::GroupCommitNotifier,
592    /// Whether the coordinator is in read-only mode (e.g. during 0dt upgrades).
593    /// The frontend path must reject mutations when this is true.
594    pub read_only: bool,
595}
596
597#[derive(Derivative)]
598#[derivative(Debug)]
599pub struct SASLChallengeResponse {
600    pub iteration_count: usize,
601    /// Base64-encoded salt for the SASL challenge.
602    pub salt: String,
603    pub nonce: String,
604}
605
606#[derive(Derivative)]
607#[derivative(Debug)]
608pub struct SASLVerifyProofResponse {
609    pub verifier: String,
610}
611
612// Facile implementation for `StartupResponse`, which does not use the `allowed`
613// feature of `ClientTransmitter`.
614impl Transmittable for StartupResponse {
615    type Allowed = bool;
616    fn to_allowed(&self) -> Self::Allowed {
617        true
618    }
619}
620
621/// The response to [`SessionClient::dump_catalog`](crate::SessionClient::dump_catalog).
622#[derive(Debug, Clone)]
623pub struct CatalogDump(String);
624
625impl CatalogDump {
626    pub fn new(raw: String) -> Self {
627        CatalogDump(raw)
628    }
629
630    pub fn into_string(self) -> String {
631        self.0
632    }
633}
634
635impl Transmittable for CatalogDump {
636    type Allowed = bool;
637    fn to_allowed(&self) -> Self::Allowed {
638        true
639    }
640}
641
642impl Transmittable for SystemVars {
643    type Allowed = bool;
644    fn to_allowed(&self) -> Self::Allowed {
645        true
646    }
647}
648
649/// The response to [`SessionClient::execute`](crate::SessionClient::execute).
650#[derive(EnumKind, Derivative)]
651#[derivative(Debug)]
652#[enum_kind(ExecuteResponseKind, derive(PartialOrd, Ord))]
653pub enum ExecuteResponse {
654    /// The default privileges were altered.
655    AlteredDefaultPrivileges,
656    /// The requested object was altered.
657    AlteredObject(ObjectType),
658    /// The role was altered.
659    AlteredRole,
660    /// The system configuration was altered.
661    AlteredSystemConfiguration,
662    /// The requested cursor was closed.
663    ClosedCursor,
664    /// The provided comment was created.
665    Comment,
666    /// The specified number of rows were copied into the requested output.
667    Copied(usize),
668    /// The response for a COPY TO STDOUT query.
669    CopyTo {
670        format: mz_sql::plan::CopyFormat,
671        resp: Box<ExecuteResponse>,
672    },
673    CopyFrom {
674        /// Table we're copying into.
675        target_id: CatalogItemId,
676        /// Human-readable full name of the target table.
677        target_name: String,
678        columns: Vec<ColumnIndex>,
679        params: CopyFormatParams<'static>,
680        ctx_extra: ExecuteContextGuard,
681    },
682    /// The requested connection was created.
683    CreatedConnection,
684    /// The requested database was created.
685    CreatedDatabase,
686    /// The requested schema was created.
687    CreatedSchema,
688    /// The requested role was created.
689    CreatedRole,
690    /// The requested cluster was created.
691    CreatedCluster,
692    /// The requested cluster replica was created.
693    CreatedClusterReplica,
694    /// The requested index was created.
695    CreatedIndex,
696    /// The requested metric sink was created.
697    CreatedMetricSink,
698    /// The requested introspection subscribe was created.
699    CreatedIntrospectionSubscribe,
700    /// The requested secret was created.
701    CreatedSecret,
702    /// The requested sink was created.
703    CreatedSink,
704    /// The requested source was created.
705    CreatedSource,
706    /// The requested table was created.
707    CreatedTable,
708    /// The requested view was created.
709    CreatedView,
710    /// The requested views were created.
711    CreatedViews,
712    /// The requested materialized view was created.
713    CreatedMaterializedView,
714    /// The requested type was created.
715    CreatedType,
716    /// The requested network policy was created.
717    CreatedNetworkPolicy,
718    /// The requested prepared statement was removed.
719    Deallocate { all: bool },
720    /// The requested cursor was declared.
721    DeclaredCursor,
722    /// The specified number of rows were deleted from the requested table.
723    Deleted(usize),
724    /// The temporary objects associated with the session have been discarded.
725    DiscardedTemp,
726    /// All state associated with the session has been discarded.
727    DiscardedAll,
728    /// The requested object was dropped.
729    DroppedObject(ObjectType),
730    /// The requested objects were dropped.
731    DroppedOwned,
732    /// The provided query was empty.
733    EmptyQuery,
734    /// Fetch results from a cursor.
735    Fetch {
736        /// The name of the cursor from which to fetch results.
737        name: String,
738        /// The number of results to fetch.
739        count: Option<FetchDirection>,
740        /// How long to wait for results to arrive.
741        timeout: ExecuteTimeout,
742        ctx_extra: ExecuteContextGuard,
743    },
744    /// The requested privilege was granted.
745    GrantedPrivilege,
746    /// The requested role was granted.
747    GrantedRole,
748    /// The specified number of rows were inserted into the requested table.
749    Inserted(usize),
750    /// The specified prepared statement was created.
751    Prepare,
752    /// A user-requested warning was raised.
753    Raised,
754    /// The requested objects were reassigned.
755    ReassignOwned,
756    /// The requested privilege was revoked.
757    RevokedPrivilege,
758    /// The requested role was revoked.
759    RevokedRole,
760    /// Rows will be delivered via the specified stream.
761    SendingRowsStreaming {
762        #[derivative(Debug = "ignore")]
763        rows: Pin<Box<dyn Stream<Item = PeekResponseUnary> + Send + Sync>>,
764        instance_id: ComputeInstanceId,
765        strategy: StatementExecutionStrategy,
766    },
767    /// Rows are known to be available immediately, and thus the execution is
768    /// considered ended in the coordinator.
769    SendingRowsImmediate {
770        #[derivative(Debug = "ignore")]
771        rows: Box<dyn RowIterator + Send + Sync>,
772    },
773    /// The specified variable was set to a new value.
774    SetVariable {
775        name: String,
776        /// Whether the operation was a `RESET` rather than a set.
777        reset: bool,
778    },
779    /// A new transaction was started.
780    StartedTransaction,
781    /// Updates to the requested source or view will be streamed to the
782    /// contained receiver.
783    Subscribing {
784        #[derivative(Debug = "ignore")]
785        rx: RowBatchStream,
786        ctx_extra: ExecuteContextGuard,
787        instance_id: ComputeInstanceId,
788    },
789    /// The active transaction committed.
790    TransactionCommitted {
791        /// Session parameters that changed because the transaction ended.
792        params: BTreeMap<&'static str, String>,
793    },
794    /// The active transaction rolled back.
795    TransactionRolledBack {
796        /// Session parameters that changed because the transaction ended.
797        params: BTreeMap<&'static str, String>,
798    },
799    /// The specified number of rows were updated in the requested table.
800    Updated(usize),
801    /// A connection was validated.
802    ValidatedConnection,
803}
804
805impl TryFrom<&Statement<Raw>> for ExecuteResponse {
806    type Error = ();
807
808    /// Returns Ok if this Statement always produces a single, trivial ExecuteResponse.
809    fn try_from(stmt: &Statement<Raw>) -> Result<Self, Self::Error> {
810        let resp_kinds = Plan::generated_from(&stmt.into())
811            .iter()
812            .map(ExecuteResponse::generated_from)
813            .flatten()
814            .cloned()
815            .collect::<BTreeSet<ExecuteResponseKind>>();
816        let resps = resp_kinds
817            .iter()
818            .map(|r| (*r).try_into())
819            .collect::<Result<Vec<ExecuteResponse>, _>>();
820        // Check if this statement's possible plans yield exactly one possible ExecuteResponse.
821        if let Ok(resps) = resps {
822            if resps.len() == 1 {
823                return Ok(resps.into_element());
824            }
825        }
826        let resp = match stmt {
827            Statement::DropObjects(DropObjectsStatement { object_type, .. }) => {
828                ExecuteResponse::DroppedObject((*object_type).into())
829            }
830            Statement::AlterObjectRename(AlterObjectRenameStatement { object_type, .. })
831            | Statement::AlterOwner(AlterOwnerStatement { object_type, .. }) => {
832                ExecuteResponse::AlteredObject((*object_type).into())
833            }
834            _ => return Err(()),
835        };
836        // Ensure that if the planner ever adds possible plans we complain here.
837        soft_assert_no_log!(
838            resp_kinds.len() == 1
839                && resp_kinds.first().expect("must exist") == &ExecuteResponseKind::from(&resp),
840            "ExecuteResponses out of sync with planner"
841        );
842        Ok(resp)
843    }
844}
845
846impl TryInto<ExecuteResponse> for ExecuteResponseKind {
847    type Error = ();
848
849    /// Attempts to convert into an ExecuteResponse. Returns an error if not possible without
850    /// actually executing a statement.
851    fn try_into(self) -> Result<ExecuteResponse, Self::Error> {
852        match self {
853            ExecuteResponseKind::AlteredDefaultPrivileges => {
854                Ok(ExecuteResponse::AlteredDefaultPrivileges)
855            }
856            ExecuteResponseKind::AlteredObject => Err(()),
857            ExecuteResponseKind::AlteredRole => Ok(ExecuteResponse::AlteredRole),
858            ExecuteResponseKind::AlteredSystemConfiguration => {
859                Ok(ExecuteResponse::AlteredSystemConfiguration)
860            }
861            ExecuteResponseKind::ClosedCursor => Ok(ExecuteResponse::ClosedCursor),
862            ExecuteResponseKind::Comment => Ok(ExecuteResponse::Comment),
863            ExecuteResponseKind::Copied => Err(()),
864            ExecuteResponseKind::CopyTo => Err(()),
865            ExecuteResponseKind::CopyFrom => Err(()),
866            ExecuteResponseKind::CreatedConnection => Ok(ExecuteResponse::CreatedConnection),
867            ExecuteResponseKind::CreatedDatabase => Ok(ExecuteResponse::CreatedDatabase),
868            ExecuteResponseKind::CreatedSchema => Ok(ExecuteResponse::CreatedSchema),
869            ExecuteResponseKind::CreatedRole => Ok(ExecuteResponse::CreatedRole),
870            ExecuteResponseKind::CreatedCluster => Ok(ExecuteResponse::CreatedCluster),
871            ExecuteResponseKind::CreatedClusterReplica => {
872                Ok(ExecuteResponse::CreatedClusterReplica)
873            }
874            ExecuteResponseKind::CreatedIndex => Ok(ExecuteResponse::CreatedIndex),
875            ExecuteResponseKind::CreatedMetricSink => Ok(ExecuteResponse::CreatedMetricSink),
876            ExecuteResponseKind::CreatedSecret => Ok(ExecuteResponse::CreatedSecret),
877            ExecuteResponseKind::CreatedSink => Ok(ExecuteResponse::CreatedSink),
878            ExecuteResponseKind::CreatedSource => Ok(ExecuteResponse::CreatedSource),
879            ExecuteResponseKind::CreatedTable => Ok(ExecuteResponse::CreatedTable),
880            ExecuteResponseKind::CreatedView => Ok(ExecuteResponse::CreatedView),
881            ExecuteResponseKind::CreatedViews => Ok(ExecuteResponse::CreatedViews),
882            ExecuteResponseKind::CreatedMaterializedView => {
883                Ok(ExecuteResponse::CreatedMaterializedView)
884            }
885            ExecuteResponseKind::CreatedNetworkPolicy => Ok(ExecuteResponse::CreatedNetworkPolicy),
886            ExecuteResponseKind::CreatedType => Ok(ExecuteResponse::CreatedType),
887            ExecuteResponseKind::Deallocate => Err(()),
888            ExecuteResponseKind::DeclaredCursor => Ok(ExecuteResponse::DeclaredCursor),
889            ExecuteResponseKind::Deleted => Err(()),
890            ExecuteResponseKind::DiscardedTemp => Ok(ExecuteResponse::DiscardedTemp),
891            ExecuteResponseKind::DiscardedAll => Ok(ExecuteResponse::DiscardedAll),
892            ExecuteResponseKind::DroppedObject => Err(()),
893            ExecuteResponseKind::DroppedOwned => Ok(ExecuteResponse::DroppedOwned),
894            ExecuteResponseKind::EmptyQuery => Ok(ExecuteResponse::EmptyQuery),
895            ExecuteResponseKind::Fetch => Err(()),
896            ExecuteResponseKind::GrantedPrivilege => Ok(ExecuteResponse::GrantedPrivilege),
897            ExecuteResponseKind::GrantedRole => Ok(ExecuteResponse::GrantedRole),
898            ExecuteResponseKind::Inserted => Err(()),
899            ExecuteResponseKind::Prepare => Ok(ExecuteResponse::Prepare),
900            ExecuteResponseKind::Raised => Ok(ExecuteResponse::Raised),
901            ExecuteResponseKind::ReassignOwned => Ok(ExecuteResponse::ReassignOwned),
902            ExecuteResponseKind::RevokedPrivilege => Ok(ExecuteResponse::RevokedPrivilege),
903            ExecuteResponseKind::RevokedRole => Ok(ExecuteResponse::RevokedRole),
904            ExecuteResponseKind::SetVariable => Err(()),
905            ExecuteResponseKind::StartedTransaction => Ok(ExecuteResponse::StartedTransaction),
906            ExecuteResponseKind::Subscribing => Err(()),
907            ExecuteResponseKind::TransactionCommitted => Err(()),
908            ExecuteResponseKind::TransactionRolledBack => Err(()),
909            ExecuteResponseKind::Updated => Err(()),
910            ExecuteResponseKind::ValidatedConnection => Ok(ExecuteResponse::ValidatedConnection),
911            ExecuteResponseKind::SendingRowsStreaming => Err(()),
912            ExecuteResponseKind::SendingRowsImmediate => Err(()),
913            ExecuteResponseKind::CreatedIntrospectionSubscribe => {
914                Ok(ExecuteResponse::CreatedIntrospectionSubscribe)
915            }
916        }
917    }
918}
919
920impl ExecuteResponse {
921    pub fn tag(&self) -> Option<String> {
922        use ExecuteResponse::*;
923        match self {
924            AlteredDefaultPrivileges => Some("ALTER DEFAULT PRIVILEGES".into()),
925            AlteredObject(o) => Some(format!("ALTER {}", o)),
926            AlteredRole => Some("ALTER ROLE".into()),
927            AlteredSystemConfiguration => Some("ALTER SYSTEM".into()),
928            ClosedCursor => Some("CLOSE CURSOR".into()),
929            Comment => Some("COMMENT".into()),
930            Copied(n) => Some(format!("COPY {}", n)),
931            CopyTo { .. } => None,
932            CopyFrom { .. } => None,
933            CreatedConnection { .. } => Some("CREATE CONNECTION".into()),
934            CreatedDatabase { .. } => Some("CREATE DATABASE".into()),
935            CreatedSchema { .. } => Some("CREATE SCHEMA".into()),
936            CreatedRole => Some("CREATE ROLE".into()),
937            CreatedCluster { .. } => Some("CREATE CLUSTER".into()),
938            CreatedClusterReplica { .. } => Some("CREATE CLUSTER REPLICA".into()),
939            CreatedIndex { .. } => Some("CREATE INDEX".into()),
940            CreatedMetricSink { .. } => Some("CREATE METRIC SINK".into()),
941            CreatedSecret { .. } => Some("CREATE SECRET".into()),
942            CreatedSink { .. } => Some("CREATE SINK".into()),
943            CreatedSource { .. } => Some("CREATE SOURCE".into()),
944            CreatedTable { .. } => Some("CREATE TABLE".into()),
945            CreatedView { .. } => Some("CREATE VIEW".into()),
946            CreatedViews { .. } => Some("CREATE VIEWS".into()),
947            CreatedMaterializedView { .. } => Some("CREATE MATERIALIZED VIEW".into()),
948            CreatedType => Some("CREATE TYPE".into()),
949            CreatedNetworkPolicy => Some("CREATE NETWORKPOLICY".into()),
950            Deallocate { all } => Some(format!("DEALLOCATE{}", if *all { " ALL" } else { "" })),
951            DeclaredCursor => Some("DECLARE CURSOR".into()),
952            Deleted(n) => Some(format!("DELETE {}", n)),
953            DiscardedTemp => Some("DISCARD TEMP".into()),
954            DiscardedAll => Some("DISCARD ALL".into()),
955            DroppedObject(o) => Some(format!("DROP {o}")),
956            DroppedOwned => Some("DROP OWNED".into()),
957            EmptyQuery => None,
958            Fetch { .. } => None,
959            GrantedPrivilege => Some("GRANT".into()),
960            GrantedRole => Some("GRANT ROLE".into()),
961            Inserted(n) => {
962                // "On successful completion, an INSERT command returns a
963                // command tag of the form `INSERT <oid> <count>`."
964                //     -- https://www.postgresql.org/docs/11/sql-insert.html
965                //
966                // OIDs are a PostgreSQL-specific historical quirk, but we
967                // can return a 0 OID to indicate that the table does not
968                // have OIDs.
969                Some(format!("INSERT 0 {}", n))
970            }
971            Prepare => Some("PREPARE".into()),
972            Raised => Some("RAISE".into()),
973            ReassignOwned => Some("REASSIGN OWNED".into()),
974            RevokedPrivilege => Some("REVOKE".into()),
975            RevokedRole => Some("REVOKE ROLE".into()),
976            SendingRowsStreaming { .. } | SendingRowsImmediate { .. } => None,
977            SetVariable { reset: true, .. } => Some("RESET".into()),
978            SetVariable { reset: false, .. } => Some("SET".into()),
979            StartedTransaction { .. } => Some("BEGIN".into()),
980            Subscribing { .. } => None,
981            TransactionCommitted { .. } => Some("COMMIT".into()),
982            TransactionRolledBack { .. } => Some("ROLLBACK".into()),
983            Updated(n) => Some(format!("UPDATE {}", n)),
984            ValidatedConnection => Some("VALIDATE CONNECTION".into()),
985            CreatedIntrospectionSubscribe => Some("CREATE INTROSPECTION SUBSCRIBE".into()),
986        }
987    }
988
989    /// Expresses which [`PlanKind`] generate which set of [`ExecuteResponseKind`].
990    /// `ExecuteResponseKind::Canceled` could be generated at any point as well, but that is
991    /// excluded from this function.
992    pub fn generated_from(plan: &PlanKind) -> &'static [ExecuteResponseKind] {
993        use ExecuteResponseKind::*;
994        use PlanKind::*;
995
996        match plan {
997            AbortTransaction => &[TransactionRolledBack],
998            AlterClusterRename
999            | AlterClusterSwap
1000            | AlterCluster
1001            | AlterClusterReplicaRename
1002            | AlterOwner
1003            | AlterItemRename
1004            | AlterRetainHistory
1005            | AlterSourceTimestampInterval
1006            | AlterNoop
1007            | AlterSchemaRename
1008            | AlterSchemaSwap
1009            | AlterSecret
1010            | AlterConnection
1011            | AlterSource
1012            | AlterSink
1013            | AlterTableAddColumn
1014            | AlterMaterializedViewApplyReplacement
1015            | AlterNetworkPolicy => &[AlteredObject],
1016            AlterDefaultPrivileges => &[AlteredDefaultPrivileges],
1017            AlterSetCluster => &[AlteredObject],
1018            AlterRole => &[AlteredRole],
1019            AlterSystemSet | AlterSystemReset | AlterSystemResetAll => {
1020                &[AlteredSystemConfiguration]
1021            }
1022            Close => &[ClosedCursor],
1023            PlanKind::CopyFrom => &[ExecuteResponseKind::CopyFrom, ExecuteResponseKind::Copied],
1024            PlanKind::CopyTo => &[ExecuteResponseKind::Copied],
1025            PlanKind::Comment => &[ExecuteResponseKind::Comment],
1026            CommitTransaction => &[TransactionCommitted, TransactionRolledBack],
1027            CreateConnection => &[CreatedConnection],
1028            CreateDatabase => &[CreatedDatabase],
1029            CreateSchema => &[CreatedSchema],
1030            CreateRole => &[CreatedRole],
1031            CreateCluster => &[CreatedCluster],
1032            CreateClusterReplica => &[CreatedClusterReplica],
1033            CreateSource | CreateSources => &[CreatedSource],
1034            CreateSecret => &[CreatedSecret],
1035            CreateSink => &[CreatedSink],
1036            CreateTable => &[CreatedTable],
1037            CreateView => &[CreatedView],
1038            CreateMaterializedView => &[CreatedMaterializedView],
1039            CreateIndex => &[CreatedIndex],
1040            CreateMetricSink => &[CreatedMetricSink],
1041            CreateType => &[CreatedType],
1042            PlanKind::Deallocate => &[ExecuteResponseKind::Deallocate],
1043            CreateNetworkPolicy => &[CreatedNetworkPolicy],
1044            Declare => &[DeclaredCursor],
1045            DiscardTemp => &[DiscardedTemp],
1046            DiscardAll => &[DiscardedAll],
1047            DropObjects => &[DroppedObject],
1048            DropOwned => &[DroppedOwned],
1049            PlanKind::EmptyQuery => &[ExecuteResponseKind::EmptyQuery],
1050            ExplainPlan | ExplainPushdown | ExplainTimestamp | Select | ShowAllVariables
1051            | ShowCreate | ShowColumns | ShowVariable | InspectShard | ExplainSinkSchema => &[
1052                ExecuteResponseKind::CopyTo,
1053                SendingRowsStreaming,
1054                SendingRowsImmediate,
1055            ],
1056            Execute | ReadThenWrite => &[
1057                Deleted,
1058                Inserted,
1059                SendingRowsStreaming,
1060                SendingRowsImmediate,
1061                Updated,
1062            ],
1063            PlanKind::Fetch => &[ExecuteResponseKind::Fetch],
1064            GrantPrivileges => &[GrantedPrivilege],
1065            GrantRole => &[GrantedRole],
1066            Insert => &[Inserted, SendingRowsImmediate],
1067            PlanKind::Prepare => &[ExecuteResponseKind::Prepare],
1068            PlanKind::Raise => &[ExecuteResponseKind::Raised],
1069            PlanKind::ReassignOwned => &[ExecuteResponseKind::ReassignOwned],
1070            RevokePrivileges => &[RevokedPrivilege],
1071            RevokeRole => &[RevokedRole],
1072            PlanKind::SetVariable | ResetVariable | PlanKind::SetTransaction => {
1073                &[ExecuteResponseKind::SetVariable]
1074            }
1075            PlanKind::Subscribe => &[Subscribing, ExecuteResponseKind::CopyTo],
1076            StartTransaction => &[StartedTransaction],
1077            SideEffectingFunc => &[SendingRowsStreaming, SendingRowsImmediate],
1078            ValidateConnection => &[ExecuteResponseKind::ValidatedConnection],
1079        }
1080    }
1081}
1082
1083/// This implementation is meant to ensure that we maintain updated information
1084/// about which types of `ExecuteResponse`s are permitted to be sent, which will
1085/// be a function of which plan we're executing.
1086impl Transmittable for ExecuteResponse {
1087    type Allowed = ExecuteResponseKind;
1088    fn to_allowed(&self) -> Self::Allowed {
1089        ExecuteResponseKind::from(self)
1090    }
1091}
mz_adapter/command.rs

mz_adapter/
command.rs
