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::active_compute_sink::ActiveSubscribeOwner;
49use crate::catalog::Catalog;
50use crate::config::{ScopedParameters, ScopedParametersScope, SystemParameterFrontend};
51use crate::coord::appends::{BuiltinTableAppendNotify, WriteResult};
52use crate::coord::consistency::CoordinatorInconsistencies;
53use crate::coord::peek::{PeekDataflowPlan, PeekResponseUnary};
54use crate::coord::timestamp_selection::TimestampDetermination;
55use crate::coord::{ExecuteContextExtra, ExecuteContextGuard};
56use crate::error::AdapterError;
57use crate::optimize::LirDataflowDescription;
58use crate::session::{EndTransactionAction, RowBatchStream, Session};
59use crate::statement_logging::{
60    FrontendStatementLoggingEvent, StatementEndedExecutionReason, StatementExecutionStrategy,
61    StatementLoggingFrontend,
62};
63use crate::statement_logging::{StatementLoggingId, WatchSetCreation};
64use crate::util::Transmittable;
65use crate::webhook::AppendWebhookResponse;
66use crate::{
67    AdapterNotice, AppendWebhookError, CollectionIdBundle, ReadHolds, TimestampExplanation,
68};
69
70/// A handle for pgwire to stream raw byte chunks to the coordinator's
71/// parallel background batch builder tasks during COPY FROM STDIN.
72#[derive(Debug)]
73pub struct CopyFromStdinWriter {
74    /// Channels for distributing raw byte chunks (split at row boundaries)
75    /// across parallel worker tasks. pgwire round-robins chunks across these.
76    pub batch_txs: Vec<mpsc::Sender<Vec<u8>>>,
77    /// Receives the final result (`Vec<ProtoBatch>` + total row count, or error)
78    /// from the collector task. pgwire uses this to commit the batches.
79    pub completion_rx: oneshot::Receiver<
80        Result<(Vec<mz_persist_client::batch::ProtoBatch>, u64), crate::AdapterError>,
81    >,
82}
83
84#[derive(Debug)]
85pub struct CatalogSnapshot {
86    pub catalog: Arc<Catalog>,
87}
88
89#[derive(Debug)]
90pub enum Command {
91    CatalogSnapshot {
92        tx: oneshot::Sender<CatalogSnapshot>,
93    },
94
95    Startup {
96        tx: oneshot::Sender<Result<StartupResponse, AdapterError>>,
97        user: User,
98        conn_id: ConnectionId,
99        client_ip: Option<IpAddr>,
100        secret_key: u32,
101        uuid: Uuid,
102        application_name: String,
103        notice_tx: mpsc::UnboundedSender<AdapterNotice>,
104    },
105
106    AuthenticatePassword {
107        tx: oneshot::Sender<Result<(), AdapterError>>,
108        role_name: String,
109        password: Option<Password>,
110    },
111
112    AuthenticateGetSASLChallenge {
113        tx: oneshot::Sender<Result<SASLChallengeResponse, AdapterError>>,
114        role_name: String,
115        nonce: String,
116    },
117
118    AuthenticateVerifySASLProof {
119        tx: oneshot::Sender<Result<SASLVerifyProofResponse, AdapterError>>,
120        role_name: String,
121        proof: String,
122        auth_message: String,
123        mock_hash: String,
124    },
125
126    CheckRoleCanLogin {
127        tx: oneshot::Sender<Result<(), AdapterError>>,
128        role_name: String,
129    },
130
131    Execute {
132        portal_name: String,
133        session: Session,
134        tx: oneshot::Sender<Response<ExecuteResponse>>,
135        /// The end-of-execution obligation of the statement this execution
136        /// serves, if any. `Coordinator::handle_execute` arms it again on
137        /// receipt. `None` means no log entry exists yet, so the coordinator
138        /// begins one.
139        outer_ctx_extra: Option<ExecuteContextExtra>,
140    },
141
142    /// Attempts to commit or abort the session's transaction. Guarantees that the Coordinator's
143    /// transaction state has been cleared, even if the commit or abort fails. (A failure can
144    /// happen, for example, if the session's role id has been dropped which will prevent
145    /// sequence_end_transaction from running.)
146    Commit {
147        action: EndTransactionAction,
148        session: Session,
149        tx: oneshot::Sender<Response<ExecuteResponse>>,
150    },
151
152    CancelRequest {
153        conn_id: ConnectionIdType,
154        secret_key: u32,
155    },
156
157    PrivilegedCancelRequest {
158        conn_id: ConnectionId,
159    },
160
161    GetWebhook {
162        database: String,
163        schema: String,
164        name: String,
165        tx: oneshot::Sender<Result<AppendWebhookResponse, AppendWebhookError>>,
166    },
167
168    GetSystemVars {
169        tx: oneshot::Sender<SystemVars>,
170    },
171
172    SetSystemVars {
173        vars: BTreeMap<String, String>,
174        conn_id: ConnectionId,
175        tx: oneshot::Sender<Result<(), AdapterError>>,
176    },
177
178    /// Replace the scoped feature-flag overrides (the complete desired state).
179    /// Computed by the system-parameter sync loop from continuous LaunchDarkly
180    /// evaluation. The coordinator stores this working copy and reconciles it
181    /// into the per-scope resolution boundaries (the compute controller's
182    /// per-replica dyncfg layer for `replica`-scoped parameters). See the
183    /// scoped feature flags design.
184    UpdateScopedSystemParameters {
185        overrides: ScopedParameters,
186        /// Bounds which objects' durable rows the reconcile may prune. See
187        /// [`crate::catalog::Op::UpdateScopedSystemParameters`].
188        prune_scope: ScopedParametersScope,
189        tx: oneshot::Sender<()>,
190    },
191
192    /// Install (or replace) the shared system-parameter frontend on the
193    /// coordinator, so catalog transactions can resolve a new object's scoped
194    /// overrides before its replica is provisioned or its first dataflow is
195    /// planned, rather than waiting for the next sync tick. Sent by the sync loop
196    /// whenever it (re)initializes the frontend. See the scoped feature flags
197    /// design.
198    InstallScopedSystemParameterFrontend {
199        frontend: Arc<SystemParameterFrontend>,
200    },
201
202    InjectAuditEvents {
203        events: Vec<crate::catalog::InjectedAuditEvent>,
204        conn_id: ConnectionId,
205        tx: oneshot::Sender<Result<(), AdapterError>>,
206    },
207
208    Terminate {
209        conn_id: ConnectionId,
210        tx: Option<oneshot::Sender<Result<(), AdapterError>>>,
211    },
212
213    /// Sets up a streaming COPY FROM STDIN operation. The coordinator
214    /// creates parallel background batch builder tasks and returns a
215    /// [`CopyFromStdinWriter`] that pgwire uses to stream raw byte chunks.
216    StartCopyFromStdin {
217        target_id: CatalogItemId,
218        target_name: String,
219        columns: Vec<ColumnIndex>,
220        /// The row description for the target table (used for constraint checks).
221        row_desc: mz_repr::RelationDesc,
222        /// Copy format parameters (text/csv/binary) for decoding raw bytes.
223        params: mz_pgcopy::CopyFormatParams<'static>,
224        session: Session,
225        tx: oneshot::Sender<Response<CopyFromStdinWriter>>,
226    },
227
228    /// Performs any cleanup and logging actions necessary for
229    /// finalizing a statement execution.
230    ///
231    /// Only used for cases that terminate in the protocol layer and
232    /// otherwise have no reason to hand control back to the coordinator.
233    /// In other cases, we piggy-back on another command.
234    RetireExecute {
235        data: ExecuteContextExtra,
236        reason: StatementEndedExecutionReason,
237    },
238
239    CheckConsistency {
240        tx: oneshot::Sender<Result<(), CoordinatorInconsistencies>>,
241    },
242
243    Dump {
244        tx: oneshot::Sender<Result<serde_json::Value, anyhow::Error>>,
245    },
246
247    GetComputeInstanceClient {
248        instance_id: ComputeInstanceId,
249        tx: oneshot::Sender<
250            Result<
251                mz_compute_client::controller::instance_client::InstanceClient,
252                mz_compute_client::controller::error::InstanceMissing,
253            >,
254        >,
255    },
256
257    GetOracle {
258        timeline: Timeline,
259        tx: oneshot::Sender<
260            Result<Arc<dyn TimestampOracle<mz_repr::Timestamp> + Send + Sync>, AdapterError>,
261        >,
262    },
263
264    DetermineRealTimeRecentTimestamp {
265        source_ids: BTreeSet<GlobalId>,
266        real_time_recency_timeout: Duration,
267        tx: oneshot::Sender<Result<Option<mz_repr::Timestamp>, AdapterError>>,
268    },
269
270    GetTransactionReadHoldsBundle {
271        conn_id: ConnectionId,
272        tx: oneshot::Sender<Option<ReadHolds>>,
273    },
274
275    /// _Merges_ the given read holds into the given connection's stored transaction read holds.
276    StoreTransactionReadHolds {
277        conn_id: ConnectionId,
278        read_holds: ReadHolds,
279        tx: oneshot::Sender<()>,
280    },
281
282    ExecuteSlowPathPeek {
283        dataflow_plan: Box<PeekDataflowPlan>,
284        determination: TimestampDetermination,
285        finishing: RowSetFinishing,
286        compute_instance: ComputeInstanceId,
287        target_replica: Option<ReplicaId>,
288        intermediate_result_type: SqlRelationType,
289        source_ids: BTreeSet<GlobalId>,
290        conn_id: ConnectionId,
291        max_result_size: u64,
292        max_query_result_size: Option<u64>,
293        /// If statement logging is enabled, contains all info needed for installing watch sets
294        /// and logging the statement execution.
295        watch_set: Option<WatchSetCreation>,
296        tx: oneshot::Sender<Result<ExecuteResponse, AdapterError>>,
297    },
298
299    ExecuteSubscribe {
300        df_desc: DataflowDescription<mz_compute_types::plan::LirRelationExpr>,
301        dependency_ids: BTreeSet<GlobalId>,
302        cluster_id: ComputeInstanceId,
303        replica_id: Option<ReplicaId>,
304        conn_id: ConnectionId,
305        session_uuid: Uuid,
306        read_holds: ReadHolds,
307        plan: plan::SubscribePlan,
308        statement_logging_id: Option<StatementLoggingId>,
309        tx: oneshot::Sender<Result<ExecuteResponse, AdapterError>>,
310    },
311
312    /// Preflight check for COPY TO S3 operation. This runs the slow S3 operations
313    /// (loading SDK config, checking bucket path, verifying permissions, uploading sentinel)
314    /// in a background task to avoid blocking the coordinator.
315    CopyToPreflight {
316        /// The S3 connection info needed for preflight checks.
317        s3_sink_connection: mz_compute_types::sinks::CopyToS3OneshotSinkConnection,
318        /// The sink ID for logging and S3 key management.
319        sink_id: GlobalId,
320        /// Response channel for the preflight result.
321        tx: oneshot::Sender<Result<(), AdapterError>>,
322    },
323
324    ExecuteCopyTo {
325        df_desc: Box<DataflowDescription<mz_compute_types::plan::LirRelationExpr>>,
326        compute_instance: ComputeInstanceId,
327        target_replica: Option<ReplicaId>,
328        source_ids: BTreeSet<GlobalId>,
329        conn_id: ConnectionId,
330        /// If statement logging is enabled, contains all info needed for installing watch sets
331        /// and logging the statement execution.
332        watch_set: Option<WatchSetCreation>,
333        tx: oneshot::Sender<Result<ExecuteResponse, AdapterError>>,
334    },
335
336    /// Execute a side-effecting function from the frontend peek path.
337    ExecuteSideEffectingFunc {
338        plan: SideEffectingFunc,
339        conn_id: ConnectionId,
340        tx: oneshot::Sender<Result<ExecuteResponse, AdapterError>>,
341    },
342
343    /// Look up an active connection by its raw connection ID, returning its
344    /// `ConnectionId` handle and authenticated role, or `None` if there is no
345    /// such connection.
346    ///
347    /// While the caller holds the returned `ConnectionId` handle, the raw
348    /// connection ID cannot be reused by a new connection. Frontend peek
349    /// sequencing relies on this to ensure that the connection it performs an
350    /// RBAC check against is the same one that a subsequent
351    /// `ExecuteSideEffectingFunc` acts on.
352    LookupConnection {
353        connection_id: u32,
354        tx: oneshot::Sender<Option<(ConnectionId, RoleId)>>,
355    },
356
357    /// Register a pending peek initiated by frontend sequencing. This is needed for:
358    /// - statement logging
359    /// - query cancellation
360    RegisterFrontendPeek {
361        uuid: Uuid,
362        conn_id: ConnectionId,
363        cluster_id: mz_controller_types::ClusterId,
364        depends_on: BTreeSet<GlobalId>,
365        is_fast_path: bool,
366        /// If statement logging is enabled, contains all info needed for installing watch sets
367        /// and logging the statement execution.
368        watch_set: Option<WatchSetCreation>,
369        tx: oneshot::Sender<Result<(), AdapterError>>,
370    },
371
372    /// Unregister and retire a pending peek that was registered but then
373    /// failed to issue, ending its statement-logging execution with the given
374    /// reason.
375    ///
376    /// Registration handed ownership of end-of-execution logging to the
377    /// coordinator, so the frontend must not log the end itself. If a
378    /// concurrent teardown (e.g. a `DROP CLUSTER`) already retired the peek
379    /// and logged its end, this is a no-op.
380    UnregisterFrontendPeek {
381        uuid: Uuid,
382        reason: StatementEndedExecutionReason,
383        tx: oneshot::Sender<()>,
384    },
385
386    /// Generate a timestamp explanation.
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<TimestampExplanation>,
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    /// Whether frontend OCC read-then-write is enabled (determined once at
601    /// process startup).
602    pub frontend_read_then_write_enabled: bool,
603    /// Requests a group commit, which is how the frontend asks for the write
604    /// timeline to advance without having anything to write.
605    pub group_commit_notifier: crate::coord::appends::GroupCommitNotifier,
606    /// Whether the coordinator is in read-only mode (e.g. during 0dt upgrades).
607    /// The frontend path must reject mutations when this is true.
608    pub read_only: bool,
609}
610
611#[derive(Derivative)]
612#[derivative(Debug)]
613pub struct SASLChallengeResponse {
614    pub iteration_count: usize,
615    /// Base64-encoded salt for the SASL challenge.
616    pub salt: String,
617    pub nonce: String,
618}
619
620#[derive(Derivative)]
621#[derivative(Debug)]
622pub struct SASLVerifyProofResponse {
623    pub verifier: String,
624}
625
626// Facile implementation for `StartupResponse`, which does not use the `allowed`
627// feature of `ClientTransmitter`.
628impl Transmittable for StartupResponse {
629    type Allowed = bool;
630    fn to_allowed(&self) -> Self::Allowed {
631        true
632    }
633}
634
635/// The response to [`SessionClient::dump_catalog`](crate::SessionClient::dump_catalog).
636#[derive(Debug, Clone)]
637pub struct CatalogDump(String);
638
639impl CatalogDump {
640    pub fn new(raw: String) -> Self {
641        CatalogDump(raw)
642    }
643
644    pub fn into_string(self) -> String {
645        self.0
646    }
647}
648
649impl Transmittable for CatalogDump {
650    type Allowed = bool;
651    fn to_allowed(&self) -> Self::Allowed {
652        true
653    }
654}
655
656impl Transmittable for SystemVars {
657    type Allowed = bool;
658    fn to_allowed(&self) -> Self::Allowed {
659        true
660    }
661}
662
663/// The response to [`SessionClient::execute`](crate::SessionClient::execute).
664#[derive(EnumKind, Derivative)]
665#[derivative(Debug)]
666#[enum_kind(ExecuteResponseKind, derive(PartialOrd, Ord))]
667pub enum ExecuteResponse {
668    /// The default privileges were altered.
669    AlteredDefaultPrivileges,
670    /// The requested object was altered.
671    AlteredObject(ObjectType),
672    /// The role was altered.
673    AlteredRole,
674    /// The system configuration was altered.
675    AlteredSystemConfiguration,
676    /// The requested cursor was closed.
677    ClosedCursor,
678    /// The provided comment was created.
679    Comment,
680    /// The specified number of rows were copied into the requested output.
681    Copied(usize),
682    /// The response for a COPY TO STDOUT query.
683    CopyTo {
684        format: mz_sql::plan::CopyFormat,
685        resp: Box<ExecuteResponse>,
686    },
687    CopyFrom {
688        /// Table we're copying into.
689        target_id: CatalogItemId,
690        /// Human-readable full name of the target table.
691        target_name: String,
692        columns: Vec<ColumnIndex>,
693        params: CopyFormatParams<'static>,
694        ctx_extra: ExecuteContextGuard,
695    },
696    /// The requested connection was created.
697    CreatedConnection,
698    /// The requested database was created.
699    CreatedDatabase,
700    /// The requested schema was created.
701    CreatedSchema,
702    /// The requested role was created.
703    CreatedRole,
704    /// The requested cluster was created.
705    CreatedCluster,
706    /// The requested cluster replica was created.
707    CreatedClusterReplica,
708    /// The requested index was created.
709    CreatedIndex,
710    /// The requested metric sink was created.
711    CreatedMetricSink,
712    /// The requested introspection subscribe was created.
713    CreatedIntrospectionSubscribe,
714    /// The requested secret was created.
715    CreatedSecret,
716    /// The requested sink was created.
717    CreatedSink,
718    /// The requested source was created.
719    CreatedSource,
720    /// The requested table was created.
721    CreatedTable,
722    /// The requested view was created.
723    CreatedView,
724    /// The requested views were created.
725    CreatedViews,
726    /// The requested materialized view was created.
727    CreatedMaterializedView,
728    /// The requested type was created.
729    CreatedType,
730    /// The requested network policy was created.
731    CreatedNetworkPolicy,
732    /// The requested prepared statement was removed.
733    Deallocate { all: bool },
734    /// The requested cursor was declared.
735    DeclaredCursor,
736    /// The specified number of rows were deleted from the requested table.
737    Deleted(usize),
738    /// The temporary objects associated with the session have been discarded.
739    DiscardedTemp,
740    /// All state associated with the session has been discarded.
741    DiscardedAll,
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 => Ok(ExecuteResponse::DiscardedAll),
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
