mz_adapter/

command.rs

// Copyright Materialize, Inc. and contributors. All rights reserved.
//
// Use of this software is governed by the Business Source License
// included in the LICENSE file.
//
// As of the Change Date specified in that file, in accordance with
// the Business Source License, use of this software will be governed
// by the Apache License, Version 2.0.

use std::collections::{BTreeMap, BTreeSet};
use std::net::IpAddr;
use std::pin::Pin;
use std::sync::Arc;
use std::time::Duration;

use chrono::{DateTime, Utc};
use derivative::Derivative;
use enum_kinds::EnumKind;
use futures::Stream;
use mz_adapter_types::connection::{ConnectionId, ConnectionIdType};
use mz_auth::password::Password;
use mz_cluster_client::ReplicaId;
use mz_compute_types::ComputeInstanceId;
use mz_compute_types::dataflows::DataflowDescription;
use mz_controller_types::ClusterId;
use mz_expr::RowSetFinishing;
use mz_ore::collections::CollectionExt;
use mz_ore::soft_assert_no_log;
use mz_ore::tracing::OpenTelemetryContext;
use mz_persist_client::PersistClient;
use mz_pgcopy::CopyFormatParams;
use mz_repr::global_id::TransientIdGen;
use mz_repr::role_id::RoleId;
use mz_repr::{CatalogItemId, ColumnIndex, GlobalId, RowIterator, SqlRelationType};
use mz_sql::ast::{FetchDirection, Raw, Statement};
use mz_sql::catalog::ObjectType;
use mz_sql::optimizer_metrics::OptimizerMetrics;
use mz_sql::plan;
use mz_sql::plan::{ExecuteTimeout, Plan, PlanKind, SideEffectingFunc};
use mz_sql::session::user::User;
use mz_sql::session::vars::{OwnedVarInput, SystemVars};
use mz_sql_parser::ast::{AlterObjectRenameStatement, AlterOwnerStatement, DropObjectsStatement};
use mz_storage_types::sources::Timeline;
use mz_timestamp_oracle::TimestampOracle;
use tokio::sync::{mpsc, oneshot};
use uuid::Uuid;

use crate::catalog::Catalog;
use crate::config::{ScopedParameters, ScopedParametersScope, SystemParameterFrontend};
use crate::coord::appends::BuiltinTableAppendNotify;
use crate::coord::consistency::CoordinatorInconsistencies;
use crate::coord::peek::{PeekDataflowPlan, PeekResponseUnary};
use crate::coord::timestamp_selection::TimestampDetermination;
use crate::coord::{ExecuteContextExtra, ExecuteContextGuard};
use crate::error::AdapterError;
use crate::session::{EndTransactionAction, RowBatchStream, Session};
use crate::statement_logging::{
    FrontendStatementLoggingEvent, StatementEndedExecutionReason, StatementExecutionStrategy,
    StatementLoggingFrontend,
};
use crate::statement_logging::{StatementLoggingId, WatchSetCreation};
use crate::util::Transmittable;
use crate::webhook::AppendWebhookResponse;
use crate::{
    AdapterNotice, AppendWebhookError, CollectionIdBundle, ReadHolds, TimestampExplanation,
};

/// A handle for pgwire to stream raw byte chunks to the coordinator's
/// parallel background batch builder tasks during COPY FROM STDIN.
#[derive(Debug)]
pub struct CopyFromStdinWriter {
    /// Channels for distributing raw byte chunks (split at row boundaries)
    /// across parallel worker tasks. pgwire round-robins chunks across these.
    pub batch_txs: Vec<mpsc::Sender<Vec<u8>>>,
    /// Receives the final result (`Vec<ProtoBatch>` + total row count, or error)
    /// from the collector task. pgwire uses this to commit the batches.
    pub completion_rx: oneshot::Receiver<
        Result<(Vec<mz_persist_client::batch::ProtoBatch>, u64), crate::AdapterError>,
    >,
}

#[derive(Debug)]
pub struct CatalogSnapshot {
    pub catalog: Arc<Catalog>,
}

#[derive(Debug)]
pub enum Command {
    CatalogSnapshot {
        tx: oneshot::Sender<CatalogSnapshot>,
    },

    Startup {
        tx: oneshot::Sender<Result<StartupResponse, AdapterError>>,
        user: User,
        conn_id: ConnectionId,
        client_ip: Option<IpAddr>,
        secret_key: u32,
        uuid: Uuid,
        application_name: String,
        notice_tx: mpsc::UnboundedSender<AdapterNotice>,
    },

    AuthenticatePassword {
        tx: oneshot::Sender<Result<(), AdapterError>>,
        role_name: String,
        password: Option<Password>,
    },

    AuthenticateGetSASLChallenge {
        tx: oneshot::Sender<Result<SASLChallengeResponse, AdapterError>>,
        role_name: String,
        nonce: String,
    },

    AuthenticateVerifySASLProof {
        tx: oneshot::Sender<Result<SASLVerifyProofResponse, AdapterError>>,
        role_name: String,
        proof: String,
        auth_message: String,
        mock_hash: String,
    },

    CheckRoleCanLogin {
        tx: oneshot::Sender<Result<(), AdapterError>>,
        role_name: String,
    },

    Execute {
        portal_name: String,
        session: Session,
        tx: oneshot::Sender<Response<ExecuteResponse>>,
        outer_ctx_extra: Option<ExecuteContextGuard>,
    },

    /// Attempts to commit or abort the session's transaction. Guarantees that the Coordinator's
    /// transaction state has been cleared, even if the commit or abort fails. (A failure can
    /// happen, for example, if the session's role id has been dropped which will prevent
    /// sequence_end_transaction from running.)
    Commit {
        action: EndTransactionAction,
        session: Session,
        tx: oneshot::Sender<Response<ExecuteResponse>>,
    },

    CancelRequest {
        conn_id: ConnectionIdType,
        secret_key: u32,
    },

    PrivilegedCancelRequest {
        conn_id: ConnectionId,
    },

    GetWebhook {
        database: String,
        schema: String,
        name: String,
        tx: oneshot::Sender<Result<AppendWebhookResponse, AppendWebhookError>>,
    },

    GetSystemVars {
        tx: oneshot::Sender<SystemVars>,
    },

    SetSystemVars {
        vars: BTreeMap<String, String>,
        conn_id: ConnectionId,
        tx: oneshot::Sender<Result<(), AdapterError>>,
    },

    /// Replace the scoped feature-flag overrides (the complete desired state).
    /// Computed by the system-parameter sync loop from continuous LaunchDarkly
    /// evaluation. The coordinator stores this working copy and reconciles it
    /// into the per-scope resolution boundaries (the compute controller's
    /// per-replica dyncfg layer for `replica`-scoped parameters). See the
    /// scoped feature flags design.
    UpdateScopedSystemParameters {
        overrides: ScopedParameters,
        /// Bounds which objects' durable rows the reconcile may prune. See
        /// [`crate::catalog::Op::UpdateScopedSystemParameters`].
        prune_scope: Option<ScopedParametersScope>,
        tx: oneshot::Sender<()>,
    },

    /// Install (or replace) the shared system-parameter frontend on the
    /// coordinator, so the create-cluster / create-replica paths can resolve a
    /// new object's scoped overrides synchronously, before the controller
    /// installs it or its first dataflow is planned, rather than waiting for
    /// the next sync tick. Sent by the sync loop whenever it (re)initializes the
    /// frontend. See the scoped feature flags design.
    InstallScopedSystemParameterFrontend {
        frontend: Arc<SystemParameterFrontend>,
    },

    InjectAuditEvents {
        events: Vec<crate::catalog::InjectedAuditEvent>,
        conn_id: ConnectionId,
        tx: oneshot::Sender<Result<(), AdapterError>>,
    },

    Terminate {
        conn_id: ConnectionId,
        tx: Option<oneshot::Sender<Result<(), AdapterError>>>,
    },

    /// Sets up a streaming COPY FROM STDIN operation. The coordinator
    /// creates parallel background batch builder tasks and returns a
    /// [`CopyFromStdinWriter`] that pgwire uses to stream raw byte chunks.
    StartCopyFromStdin {
        target_id: CatalogItemId,
        target_name: String,
        columns: Vec<ColumnIndex>,
        /// The row description for the target table (used for constraint checks).
        row_desc: mz_repr::RelationDesc,
        /// Copy format parameters (text/csv/binary) for decoding raw bytes.
        params: mz_pgcopy::CopyFormatParams<'static>,
        session: Session,
        tx: oneshot::Sender<Response<CopyFromStdinWriter>>,
    },

    /// Performs any cleanup and logging actions necessary for
    /// finalizing a statement execution.
    ///
    /// Only used for cases that terminate in the protocol layer and
    /// otherwise have no reason to hand control back to the coordinator.
    /// In other cases, we piggy-back on another command.
    RetireExecute {
        data: ExecuteContextExtra,
        reason: StatementEndedExecutionReason,
    },

    CheckConsistency {
        tx: oneshot::Sender<Result<(), CoordinatorInconsistencies>>,
    },

    Dump {
        tx: oneshot::Sender<Result<serde_json::Value, anyhow::Error>>,
    },

    GetComputeInstanceClient {
        instance_id: ComputeInstanceId,
        tx: oneshot::Sender<
            Result<
                mz_compute_client::controller::instance_client::InstanceClient,
                mz_compute_client::controller::error::InstanceMissing,
            >,
        >,
    },

    GetOracle {
        timeline: Timeline,
        tx: oneshot::Sender<
            Result<Arc<dyn TimestampOracle<mz_repr::Timestamp> + Send + Sync>, AdapterError>,
        >,
    },

    DetermineRealTimeRecentTimestamp {
        source_ids: BTreeSet<GlobalId>,
        real_time_recency_timeout: Duration,
        tx: oneshot::Sender<Result<Option<mz_repr::Timestamp>, AdapterError>>,
    },

    GetTransactionReadHoldsBundle {
        conn_id: ConnectionId,
        tx: oneshot::Sender<Option<ReadHolds>>,
    },

    /// _Merges_ the given read holds into the given connection's stored transaction read holds.
    StoreTransactionReadHolds {
        conn_id: ConnectionId,
        read_holds: ReadHolds,
        tx: oneshot::Sender<()>,
    },

    ExecuteSlowPathPeek {
        dataflow_plan: Box<PeekDataflowPlan>,
        determination: TimestampDetermination,
        finishing: RowSetFinishing,
        compute_instance: ComputeInstanceId,
        target_replica: Option<ReplicaId>,
        intermediate_result_type: SqlRelationType,
        source_ids: BTreeSet<GlobalId>,
        conn_id: ConnectionId,
        max_result_size: u64,
        max_query_result_size: Option<u64>,
        /// If statement logging is enabled, contains all info needed for installing watch sets
        /// and logging the statement execution.
        watch_set: Option<WatchSetCreation>,
        tx: oneshot::Sender<Result<ExecuteResponse, AdapterError>>,
    },

    ExecuteSubscribe {
        df_desc: DataflowDescription<mz_compute_types::plan::LirRelationExpr>,
        dependency_ids: BTreeSet<GlobalId>,
        cluster_id: ComputeInstanceId,
        replica_id: Option<ReplicaId>,
        conn_id: ConnectionId,
        session_uuid: Uuid,
        read_holds: ReadHolds,
        plan: plan::SubscribePlan,
        statement_logging_id: Option<StatementLoggingId>,
        tx: oneshot::Sender<Result<ExecuteResponse, AdapterError>>,
    },

    /// Preflight check for COPY TO S3 operation. This runs the slow S3 operations
    /// (loading SDK config, checking bucket path, verifying permissions, uploading sentinel)
    /// in a background task to avoid blocking the coordinator.
    CopyToPreflight {
        /// The S3 connection info needed for preflight checks.
        s3_sink_connection: mz_compute_types::sinks::CopyToS3OneshotSinkConnection,
        /// The sink ID for logging and S3 key management.
        sink_id: GlobalId,
        /// Response channel for the preflight result.
        tx: oneshot::Sender<Result<(), AdapterError>>,
    },

    ExecuteCopyTo {
        df_desc: Box<DataflowDescription<mz_compute_types::plan::LirRelationExpr>>,
        compute_instance: ComputeInstanceId,
        target_replica: Option<ReplicaId>,
        source_ids: BTreeSet<GlobalId>,
        conn_id: ConnectionId,
        /// If statement logging is enabled, contains all info needed for installing watch sets
        /// and logging the statement execution.
        watch_set: Option<WatchSetCreation>,
        tx: oneshot::Sender<Result<ExecuteResponse, AdapterError>>,
    },

    /// Execute a side-effecting function from the frontend peek path.
    ExecuteSideEffectingFunc {
        plan: SideEffectingFunc,
        conn_id: ConnectionId,
        tx: oneshot::Sender<Result<ExecuteResponse, AdapterError>>,
    },

    /// Look up an active connection by its raw connection ID, returning its
    /// `ConnectionId` handle and authenticated role, or `None` if there is no
    /// such connection.
    ///
    /// While the caller holds the returned `ConnectionId` handle, the raw
    /// connection ID cannot be reused by a new connection. Frontend peek
    /// sequencing relies on this to ensure that the connection it performs an
    /// RBAC check against is the same one that a subsequent
    /// `ExecuteSideEffectingFunc` acts on.
    LookupConnection {
        connection_id: u32,
        tx: oneshot::Sender<Option<(ConnectionId, RoleId)>>,
    },

    /// Register a pending peek initiated by frontend sequencing. This is needed for:
    /// - statement logging
    /// - query cancellation
    RegisterFrontendPeek {
        uuid: Uuid,
        conn_id: ConnectionId,
        cluster_id: mz_controller_types::ClusterId,
        depends_on: BTreeSet<GlobalId>,
        is_fast_path: bool,
        /// If statement logging is enabled, contains all info needed for installing watch sets
        /// and logging the statement execution.
        watch_set: Option<WatchSetCreation>,
        tx: oneshot::Sender<Result<(), AdapterError>>,
    },

    /// Unregister and retire a pending peek that was registered but then
    /// failed to issue, ending its statement-logging execution with the given
    /// reason.
    ///
    /// Registration handed ownership of end-of-execution logging to the
    /// coordinator, so the frontend must not log the end itself. If a
    /// concurrent teardown (e.g. a `DROP CLUSTER`) already retired the peek
    /// and logged its end, this is a no-op.
    UnregisterFrontendPeek {
        uuid: Uuid,
        reason: StatementEndedExecutionReason,
        tx: oneshot::Sender<()>,
    },

    /// Generate a timestamp explanation.
    /// This is used when `emit_timestamp_notice` is enabled.
    ExplainTimestamp {
        conn_id: ConnectionId,
        session_wall_time: DateTime<Utc>,
        cluster_id: ClusterId,
        id_bundle: CollectionIdBundle,
        determination: TimestampDetermination,
        tx: oneshot::Sender<TimestampExplanation>,
    },

    /// Statement logging event from frontend peek sequencing.
    /// No response channel needed - this is fire-and-forget.
    FrontendStatementLogging(FrontendStatementLoggingEvent),
}

impl Command {
    pub fn session(&self) -> Option<&Session> {
        match self {
            Command::Execute { session, .. }
            | Command::Commit { session, .. }
            | Command::StartCopyFromStdin { session, .. } => Some(session),
            Command::CancelRequest { .. }
            | Command::Startup { .. }
            | Command::AuthenticatePassword { .. }
            | Command::AuthenticateGetSASLChallenge { .. }
            | Command::AuthenticateVerifySASLProof { .. }
            | Command::CheckRoleCanLogin { .. }
            | Command::CatalogSnapshot { .. }
            | Command::PrivilegedCancelRequest { .. }
            | Command::GetWebhook { .. }
            | Command::Terminate { .. }
            | Command::GetSystemVars { .. }
            | Command::SetSystemVars { .. }
            | Command::UpdateScopedSystemParameters { .. }
            | Command::InstallScopedSystemParameterFrontend { .. }
            | Command::RetireExecute { .. }
            | Command::CheckConsistency { .. }
            | Command::Dump { .. }
            | Command::GetComputeInstanceClient { .. }
            | Command::GetOracle { .. }
            | Command::DetermineRealTimeRecentTimestamp { .. }
            | Command::GetTransactionReadHoldsBundle { .. }
            | Command::StoreTransactionReadHolds { .. }
            | Command::ExecuteSlowPathPeek { .. }
            | Command::ExecuteSubscribe { .. }
            | Command::CopyToPreflight { .. }
            | Command::ExecuteCopyTo { .. }
            | Command::ExecuteSideEffectingFunc { .. }
            | Command::LookupConnection { .. }
            | Command::RegisterFrontendPeek { .. }
            | Command::UnregisterFrontendPeek { .. }
            | Command::ExplainTimestamp { .. }
            | Command::FrontendStatementLogging(..)
            | Command::InjectAuditEvents { .. } => None,
        }
    }

    pub fn session_mut(&mut self) -> Option<&mut Session> {
        match self {
            Command::Execute { session, .. }
            | Command::Commit { session, .. }
            | Command::StartCopyFromStdin { session, .. } => Some(session),
            Command::CancelRequest { .. }
            | Command::Startup { .. }
            | Command::AuthenticatePassword { .. }
            | Command::AuthenticateGetSASLChallenge { .. }
            | Command::AuthenticateVerifySASLProof { .. }
            | Command::CheckRoleCanLogin { .. }
            | Command::CatalogSnapshot { .. }
            | Command::PrivilegedCancelRequest { .. }
            | Command::GetWebhook { .. }
            | Command::Terminate { .. }
            | Command::GetSystemVars { .. }
            | Command::SetSystemVars { .. }
            | Command::UpdateScopedSystemParameters { .. }
            | Command::InstallScopedSystemParameterFrontend { .. }
            | Command::RetireExecute { .. }
            | Command::CheckConsistency { .. }
            | Command::Dump { .. }
            | Command::GetComputeInstanceClient { .. }
            | Command::GetOracle { .. }
            | Command::DetermineRealTimeRecentTimestamp { .. }
            | Command::GetTransactionReadHoldsBundle { .. }
            | Command::StoreTransactionReadHolds { .. }
            | Command::ExecuteSlowPathPeek { .. }
            | Command::ExecuteSubscribe { .. }
            | Command::CopyToPreflight { .. }
            | Command::ExecuteCopyTo { .. }
            | Command::ExecuteSideEffectingFunc { .. }
            | Command::LookupConnection { .. }
            | Command::RegisterFrontendPeek { .. }
            | Command::UnregisterFrontendPeek { .. }
            | Command::ExplainTimestamp { .. }
            | Command::FrontendStatementLogging(..)
            | Command::InjectAuditEvents { .. } => None,
        }
    }
}

#[derive(Debug)]
pub struct Response<T> {
    pub result: Result<T, AdapterError>,
    pub session: Session,
    pub otel_ctx: OpenTelemetryContext,
}

#[derive(Debug, Clone, Copy)]
pub struct SuperuserAttribute(pub Option<bool>);

/// The response to [`Client::startup`](crate::Client::startup).
#[derive(Derivative)]
#[derivative(Debug)]
pub struct StartupResponse {
    /// RoleId for the user.
    pub role_id: RoleId,
    /// The role's superuser attribute in the Catalog.
    /// This attribute is None for Cloud. Cloud is able
    /// to derive the role's superuser status from
    /// external_metadata_rx.
    pub superuser_attribute: SuperuserAttribute,
    /// A future that completes when all necessary Builtin Table writes have completed.
    #[derivative(Debug = "ignore")]
    pub write_notify: BuiltinTableAppendNotify,
    /// Map of (name, VarInput::Flat) tuples of session default variables that should be set.
    pub session_defaults: BTreeMap<String, OwnedVarInput>,
    pub catalog: Arc<Catalog>,
    pub storage_collections:
        Arc<dyn mz_storage_client::storage_collections::StorageCollections + Send + Sync>,
    pub transient_id_gen: Arc<TransientIdGen>,
    pub optimizer_metrics: OptimizerMetrics,
    pub persist_client: PersistClient,
    pub statement_logging_frontend: StatementLoggingFrontend,
}

#[derive(Derivative)]
#[derivative(Debug)]
pub struct SASLChallengeResponse {
    pub iteration_count: usize,
    /// Base64-encoded salt for the SASL challenge.
    pub salt: String,
    pub nonce: String,
}

#[derive(Derivative)]
#[derivative(Debug)]
pub struct SASLVerifyProofResponse {
    pub verifier: String,
}

// Facile implementation for `StartupResponse`, which does not use the `allowed`
// feature of `ClientTransmitter`.
impl Transmittable for StartupResponse {
    type Allowed = bool;
    fn to_allowed(&self) -> Self::Allowed {
        true
    }
}

/// The response to [`SessionClient::dump_catalog`](crate::SessionClient::dump_catalog).
#[derive(Debug, Clone)]
pub struct CatalogDump(String);

impl CatalogDump {
    pub fn new(raw: String) -> Self {
        CatalogDump(raw)
    }

    pub fn into_string(self) -> String {
        self.0
    }
}

impl Transmittable for CatalogDump {
    type Allowed = bool;
    fn to_allowed(&self) -> Self::Allowed {
        true
    }
}

impl Transmittable for SystemVars {
    type Allowed = bool;
    fn to_allowed(&self) -> Self::Allowed {
        true
    }
}

/// The response to [`SessionClient::execute`](crate::SessionClient::execute).
#[derive(EnumKind, Derivative)]
#[derivative(Debug)]
#[enum_kind(ExecuteResponseKind, derive(PartialOrd, Ord))]
pub enum ExecuteResponse {
    /// The default privileges were altered.
    AlteredDefaultPrivileges,
    /// The requested object was altered.
    AlteredObject(ObjectType),
    /// The role was altered.
    AlteredRole,
    /// The system configuration was altered.
    AlteredSystemConfiguration,
    /// The requested cursor was closed.
    ClosedCursor,
    /// The provided comment was created.
    Comment,
    /// The specified number of rows were copied into the requested output.
    Copied(usize),
    /// The response for a COPY TO STDOUT query.
    CopyTo {
        format: mz_sql::plan::CopyFormat,
        resp: Box<ExecuteResponse>,
    },
    CopyFrom {
        /// Table we're copying into.
        target_id: CatalogItemId,
        /// Human-readable full name of the target table.
        target_name: String,
        columns: Vec<ColumnIndex>,
        params: CopyFormatParams<'static>,
        ctx_extra: ExecuteContextGuard,
    },
    /// The requested connection was created.
    CreatedConnection,
    /// The requested database was created.
    CreatedDatabase,
    /// The requested schema was created.
    CreatedSchema,
    /// The requested role was created.
    CreatedRole,
    /// The requested cluster was created.
    CreatedCluster,
    /// The requested cluster replica was created.
    CreatedClusterReplica,
    /// The requested index was created.
    CreatedIndex,
    /// The requested introspection subscribe was created.
    CreatedIntrospectionSubscribe,
    /// The requested secret was created.
    CreatedSecret,
    /// The requested sink was created.
    CreatedSink,
    /// The requested source was created.
    CreatedSource,
    /// The requested table was created.
    CreatedTable,
    /// The requested view was created.
    CreatedView,
    /// The requested views were created.
    CreatedViews,
    /// The requested materialized view was created.
    CreatedMaterializedView,
    /// The requested type was created.
    CreatedType,
    /// The requested network policy was created.
    CreatedNetworkPolicy,
    /// The requested prepared statement was removed.
    Deallocate { all: bool },
    /// The requested cursor was declared.
    DeclaredCursor,
    /// The specified number of rows were deleted from the requested table.
    Deleted(usize),
    /// The temporary objects associated with the session have been discarded.
    DiscardedTemp,
    /// All state associated with the session has been discarded.
    DiscardedAll,
    /// The requested object was dropped.
    DroppedObject(ObjectType),
    /// The requested objects were dropped.
    DroppedOwned,
    /// The provided query was empty.
    EmptyQuery,
    /// Fetch results from a cursor.
    Fetch {
        /// The name of the cursor from which to fetch results.
        name: String,
        /// The number of results to fetch.
        count: Option<FetchDirection>,
        /// How long to wait for results to arrive.
        timeout: ExecuteTimeout,
        ctx_extra: ExecuteContextGuard,
    },
    /// The requested privilege was granted.
    GrantedPrivilege,
    /// The requested role was granted.
    GrantedRole,
    /// The specified number of rows were inserted into the requested table.
    Inserted(usize),
    /// The specified prepared statement was created.
    Prepare,
    /// A user-requested warning was raised.
    Raised,
    /// The requested objects were reassigned.
    ReassignOwned,
    /// The requested privilege was revoked.
    RevokedPrivilege,
    /// The requested role was revoked.
    RevokedRole,
    /// Rows will be delivered via the specified stream.
    SendingRowsStreaming {
        #[derivative(Debug = "ignore")]
        rows: Pin<Box<dyn Stream<Item = PeekResponseUnary> + Send + Sync>>,
        instance_id: ComputeInstanceId,
        strategy: StatementExecutionStrategy,
    },
    /// Rows are known to be available immediately, and thus the execution is
    /// considered ended in the coordinator.
    SendingRowsImmediate {
        #[derivative(Debug = "ignore")]
        rows: Box<dyn RowIterator + Send + Sync>,
    },
    /// The specified variable was set to a new value.
    SetVariable {
        name: String,
        /// Whether the operation was a `RESET` rather than a set.
        reset: bool,
    },
    /// A new transaction was started.
    StartedTransaction,
    /// Updates to the requested source or view will be streamed to the
    /// contained receiver.
    Subscribing {
        rx: RowBatchStream,
        ctx_extra: ExecuteContextGuard,
        instance_id: ComputeInstanceId,
    },
    /// The active transaction committed.
    TransactionCommitted {
        /// Session parameters that changed because the transaction ended.
        params: BTreeMap<&'static str, String>,
    },
    /// The active transaction rolled back.
    TransactionRolledBack {
        /// Session parameters that changed because the transaction ended.
        params: BTreeMap<&'static str, String>,
    },
    /// The specified number of rows were updated in the requested table.
    Updated(usize),
    /// A connection was validated.
    ValidatedConnection,
}

impl TryFrom<&Statement<Raw>> for ExecuteResponse {
    type Error = ();

    /// Returns Ok if this Statement always produces a single, trivial ExecuteResponse.
    fn try_from(stmt: &Statement<Raw>) -> Result<Self, Self::Error> {
        let resp_kinds = Plan::generated_from(&stmt.into())
            .iter()
            .map(ExecuteResponse::generated_from)
            .flatten()
            .cloned()
            .collect::<BTreeSet<ExecuteResponseKind>>();
        let resps = resp_kinds
            .iter()
            .map(|r| (*r).try_into())
            .collect::<Result<Vec<ExecuteResponse>, _>>();
        // Check if this statement's possible plans yield exactly one possible ExecuteResponse.
        if let Ok(resps) = resps {
            if resps.len() == 1 {
                return Ok(resps.into_element());
            }
        }
        let resp = match stmt {
            Statement::DropObjects(DropObjectsStatement { object_type, .. }) => {
                ExecuteResponse::DroppedObject((*object_type).into())
            }
            Statement::AlterObjectRename(AlterObjectRenameStatement { object_type, .. })
            | Statement::AlterOwner(AlterOwnerStatement { object_type, .. }) => {
                ExecuteResponse::AlteredObject((*object_type).into())
            }
            _ => return Err(()),
        };
        // Ensure that if the planner ever adds possible plans we complain here.
        soft_assert_no_log!(
            resp_kinds.len() == 1
                && resp_kinds.first().expect("must exist") == &ExecuteResponseKind::from(&resp),
            "ExecuteResponses out of sync with planner"
        );
        Ok(resp)
    }
}

impl TryInto<ExecuteResponse> for ExecuteResponseKind {
    type Error = ();

    /// Attempts to convert into an ExecuteResponse. Returns an error if not possible without
    /// actually executing a statement.
    fn try_into(self) -> Result<ExecuteResponse, Self::Error> {
        match self {
            ExecuteResponseKind::AlteredDefaultPrivileges => {
                Ok(ExecuteResponse::AlteredDefaultPrivileges)
            }
            ExecuteResponseKind::AlteredObject => Err(()),
            ExecuteResponseKind::AlteredRole => Ok(ExecuteResponse::AlteredRole),
            ExecuteResponseKind::AlteredSystemConfiguration => {
                Ok(ExecuteResponse::AlteredSystemConfiguration)
            }
            ExecuteResponseKind::ClosedCursor => Ok(ExecuteResponse::ClosedCursor),
            ExecuteResponseKind::Comment => Ok(ExecuteResponse::Comment),
            ExecuteResponseKind::Copied => Err(()),
            ExecuteResponseKind::CopyTo => Err(()),
            ExecuteResponseKind::CopyFrom => Err(()),
            ExecuteResponseKind::CreatedConnection => Ok(ExecuteResponse::CreatedConnection),
            ExecuteResponseKind::CreatedDatabase => Ok(ExecuteResponse::CreatedDatabase),
            ExecuteResponseKind::CreatedSchema => Ok(ExecuteResponse::CreatedSchema),
            ExecuteResponseKind::CreatedRole => Ok(ExecuteResponse::CreatedRole),
            ExecuteResponseKind::CreatedCluster => Ok(ExecuteResponse::CreatedCluster),
            ExecuteResponseKind::CreatedClusterReplica => {
                Ok(ExecuteResponse::CreatedClusterReplica)
            }
            ExecuteResponseKind::CreatedIndex => Ok(ExecuteResponse::CreatedIndex),
            ExecuteResponseKind::CreatedSecret => Ok(ExecuteResponse::CreatedSecret),
            ExecuteResponseKind::CreatedSink => Ok(ExecuteResponse::CreatedSink),
            ExecuteResponseKind::CreatedSource => Ok(ExecuteResponse::CreatedSource),
            ExecuteResponseKind::CreatedTable => Ok(ExecuteResponse::CreatedTable),
            ExecuteResponseKind::CreatedView => Ok(ExecuteResponse::CreatedView),
            ExecuteResponseKind::CreatedViews => Ok(ExecuteResponse::CreatedViews),
            ExecuteResponseKind::CreatedMaterializedView => {
                Ok(ExecuteResponse::CreatedMaterializedView)
            }
            ExecuteResponseKind::CreatedNetworkPolicy => Ok(ExecuteResponse::CreatedNetworkPolicy),
            ExecuteResponseKind::CreatedType => Ok(ExecuteResponse::CreatedType),
            ExecuteResponseKind::Deallocate => Err(()),
            ExecuteResponseKind::DeclaredCursor => Ok(ExecuteResponse::DeclaredCursor),
            ExecuteResponseKind::Deleted => Err(()),
            ExecuteResponseKind::DiscardedTemp => Ok(ExecuteResponse::DiscardedTemp),
            ExecuteResponseKind::DiscardedAll => Ok(ExecuteResponse::DiscardedAll),
            ExecuteResponseKind::DroppedObject => Err(()),
            ExecuteResponseKind::DroppedOwned => Ok(ExecuteResponse::DroppedOwned),
            ExecuteResponseKind::EmptyQuery => Ok(ExecuteResponse::EmptyQuery),
            ExecuteResponseKind::Fetch => Err(()),
            ExecuteResponseKind::GrantedPrivilege => Ok(ExecuteResponse::GrantedPrivilege),
            ExecuteResponseKind::GrantedRole => Ok(ExecuteResponse::GrantedRole),
            ExecuteResponseKind::Inserted => Err(()),
            ExecuteResponseKind::Prepare => Ok(ExecuteResponse::Prepare),
            ExecuteResponseKind::Raised => Ok(ExecuteResponse::Raised),
            ExecuteResponseKind::ReassignOwned => Ok(ExecuteResponse::ReassignOwned),
            ExecuteResponseKind::RevokedPrivilege => Ok(ExecuteResponse::RevokedPrivilege),
            ExecuteResponseKind::RevokedRole => Ok(ExecuteResponse::RevokedRole),
            ExecuteResponseKind::SetVariable => Err(()),
            ExecuteResponseKind::StartedTransaction => Ok(ExecuteResponse::StartedTransaction),
            ExecuteResponseKind::Subscribing => Err(()),
            ExecuteResponseKind::TransactionCommitted => Err(()),
            ExecuteResponseKind::TransactionRolledBack => Err(()),
            ExecuteResponseKind::Updated => Err(()),
            ExecuteResponseKind::ValidatedConnection => Ok(ExecuteResponse::ValidatedConnection),
            ExecuteResponseKind::SendingRowsStreaming => Err(()),
            ExecuteResponseKind::SendingRowsImmediate => Err(()),
            ExecuteResponseKind::CreatedIntrospectionSubscribe => {
                Ok(ExecuteResponse::CreatedIntrospectionSubscribe)
            }
        }
    }
}

impl ExecuteResponse {
    pub fn tag(&self) -> Option<String> {
        use ExecuteResponse::*;
        match self {
            AlteredDefaultPrivileges => Some("ALTER DEFAULT PRIVILEGES".into()),
            AlteredObject(o) => Some(format!("ALTER {}", o)),
            AlteredRole => Some("ALTER ROLE".into()),
            AlteredSystemConfiguration => Some("ALTER SYSTEM".into()),
            ClosedCursor => Some("CLOSE CURSOR".into()),
            Comment => Some("COMMENT".into()),
            Copied(n) => Some(format!("COPY {}", n)),
            CopyTo { .. } => None,
            CopyFrom { .. } => None,
            CreatedConnection { .. } => Some("CREATE CONNECTION".into()),
            CreatedDatabase { .. } => Some("CREATE DATABASE".into()),
            CreatedSchema { .. } => Some("CREATE SCHEMA".into()),
            CreatedRole => Some("CREATE ROLE".into()),
            CreatedCluster { .. } => Some("CREATE CLUSTER".into()),
            CreatedClusterReplica { .. } => Some("CREATE CLUSTER REPLICA".into()),
            CreatedIndex { .. } => Some("CREATE INDEX".into()),
            CreatedSecret { .. } => Some("CREATE SECRET".into()),
            CreatedSink { .. } => Some("CREATE SINK".into()),
            CreatedSource { .. } => Some("CREATE SOURCE".into()),
            CreatedTable { .. } => Some("CREATE TABLE".into()),
            CreatedView { .. } => Some("CREATE VIEW".into()),
            CreatedViews { .. } => Some("CREATE VIEWS".into()),
            CreatedMaterializedView { .. } => Some("CREATE MATERIALIZED VIEW".into()),
            CreatedType => Some("CREATE TYPE".into()),
            CreatedNetworkPolicy => Some("CREATE NETWORKPOLICY".into()),
            Deallocate { all } => Some(format!("DEALLOCATE{}", if *all { " ALL" } else { "" })),
            DeclaredCursor => Some("DECLARE CURSOR".into()),
            Deleted(n) => Some(format!("DELETE {}", n)),
            DiscardedTemp => Some("DISCARD TEMP".into()),
            DiscardedAll => Some("DISCARD ALL".into()),
            DroppedObject(o) => Some(format!("DROP {o}")),
            DroppedOwned => Some("DROP OWNED".into()),
            EmptyQuery => None,
            Fetch { .. } => None,
            GrantedPrivilege => Some("GRANT".into()),
            GrantedRole => Some("GRANT ROLE".into()),
            Inserted(n) => {
                // "On successful completion, an INSERT command returns a
                // command tag of the form `INSERT <oid> <count>`."
                //     -- https://www.postgresql.org/docs/11/sql-insert.html
                //
                // OIDs are a PostgreSQL-specific historical quirk, but we
                // can return a 0 OID to indicate that the table does not
                // have OIDs.
                Some(format!("INSERT 0 {}", n))
            }
            Prepare => Some("PREPARE".into()),
            Raised => Some("RAISE".into()),
            ReassignOwned => Some("REASSIGN OWNED".into()),
            RevokedPrivilege => Some("REVOKE".into()),
            RevokedRole => Some("REVOKE ROLE".into()),
            SendingRowsStreaming { .. } | SendingRowsImmediate { .. } => None,
            SetVariable { reset: true, .. } => Some("RESET".into()),
            SetVariable { reset: false, .. } => Some("SET".into()),
            StartedTransaction { .. } => Some("BEGIN".into()),
            Subscribing { .. } => None,
            TransactionCommitted { .. } => Some("COMMIT".into()),
            TransactionRolledBack { .. } => Some("ROLLBACK".into()),
            Updated(n) => Some(format!("UPDATE {}", n)),
            ValidatedConnection => Some("VALIDATE CONNECTION".into()),
            CreatedIntrospectionSubscribe => Some("CREATE INTROSPECTION SUBSCRIBE".into()),
        }
    }

    /// Expresses which [`PlanKind`] generate which set of [`ExecuteResponseKind`].
    /// `ExecuteResponseKind::Canceled` could be generated at any point as well, but that is
    /// excluded from this function.
    pub fn generated_from(plan: &PlanKind) -> &'static [ExecuteResponseKind] {
        use ExecuteResponseKind::*;
        use PlanKind::*;

        match plan {
            AbortTransaction => &[TransactionRolledBack],
            AlterClusterRename
            | AlterClusterSwap
            | AlterCluster
            | AlterClusterReplicaRename
            | AlterOwner
            | AlterItemRename
            | AlterRetainHistory
            | AlterSourceTimestampInterval
            | AlterNoop
            | AlterSchemaRename
            | AlterSchemaSwap
            | AlterSecret
            | AlterConnection
            | AlterSource
            | AlterSink
            | AlterTableAddColumn
            | AlterMaterializedViewApplyReplacement
            | AlterNetworkPolicy => &[AlteredObject],
            AlterDefaultPrivileges => &[AlteredDefaultPrivileges],
            AlterSetCluster => &[AlteredObject],
            AlterRole => &[AlteredRole],
            AlterSystemSet | AlterSystemReset | AlterSystemResetAll => {
                &[AlteredSystemConfiguration]
            }
            Close => &[ClosedCursor],
            PlanKind::CopyFrom => &[ExecuteResponseKind::CopyFrom, ExecuteResponseKind::Copied],
            PlanKind::CopyTo => &[ExecuteResponseKind::Copied],
            PlanKind::Comment => &[ExecuteResponseKind::Comment],
            CommitTransaction => &[TransactionCommitted, TransactionRolledBack],
            CreateConnection => &[CreatedConnection],
            CreateDatabase => &[CreatedDatabase],
            CreateSchema => &[CreatedSchema],
            CreateRole => &[CreatedRole],
            CreateCluster => &[CreatedCluster],
            CreateClusterReplica => &[CreatedClusterReplica],
            CreateSource | CreateSources => &[CreatedSource],
            CreateSecret => &[CreatedSecret],
            CreateSink => &[CreatedSink],
            CreateTable => &[CreatedTable],
            CreateView => &[CreatedView],
            CreateMaterializedView => &[CreatedMaterializedView],
            CreateIndex => &[CreatedIndex],
            CreateType => &[CreatedType],
            PlanKind::Deallocate => &[ExecuteResponseKind::Deallocate],
            CreateNetworkPolicy => &[CreatedNetworkPolicy],
            Declare => &[DeclaredCursor],
            DiscardTemp => &[DiscardedTemp],
            DiscardAll => &[DiscardedAll],
            DropObjects => &[DroppedObject],
            DropOwned => &[DroppedOwned],
            PlanKind::EmptyQuery => &[ExecuteResponseKind::EmptyQuery],
            ExplainPlan | ExplainPushdown | ExplainTimestamp | Select | ShowAllVariables
            | ShowCreate | ShowColumns | ShowVariable | InspectShard | ExplainSinkSchema => &[
                ExecuteResponseKind::CopyTo,
                SendingRowsStreaming,
                SendingRowsImmediate,
            ],
            Execute | ReadThenWrite => &[
                Deleted,
                Inserted,
                SendingRowsStreaming,
                SendingRowsImmediate,
                Updated,
            ],
            PlanKind::Fetch => &[ExecuteResponseKind::Fetch],
            GrantPrivileges => &[GrantedPrivilege],
            GrantRole => &[GrantedRole],
            Insert => &[Inserted, SendingRowsImmediate],
            PlanKind::Prepare => &[ExecuteResponseKind::Prepare],
            PlanKind::Raise => &[ExecuteResponseKind::Raised],
            PlanKind::ReassignOwned => &[ExecuteResponseKind::ReassignOwned],
            RevokePrivileges => &[RevokedPrivilege],
            RevokeRole => &[RevokedRole],
            PlanKind::SetVariable | ResetVariable | PlanKind::SetTransaction => {
                &[ExecuteResponseKind::SetVariable]
            }
            PlanKind::Subscribe => &[Subscribing, ExecuteResponseKind::CopyTo],
            StartTransaction => &[StartedTransaction],
            SideEffectingFunc => &[SendingRowsStreaming, SendingRowsImmediate],
            ValidateConnection => &[ExecuteResponseKind::ValidatedConnection],
        }
    }
}

/// This implementation is meant to ensure that we maintain updated information
/// about which types of `ExecuteResponse`s are permitted to be sent, which will
/// be a function of which plan we're executing.
impl Transmittable for ExecuteResponse {
    type Allowed = ExecuteResponseKind;
    fn to_allowed(&self) -> Self::Allowed {
        ExecuteResponseKind::from(self)
    }
}