mz_adapter/

session.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.

//! Per-connection configuration parameters and state.

#![warn(missing_docs)]

use std::collections::BTreeMap;
use std::collections::btree_map::Entry;
use std::future::Future;
use std::mem;
use std::net::IpAddr;
use std::sync::Arc;

use chrono::{DateTime, Utc};
use derivative::Derivative;
use mz_adapter_types::connection::ConnectionId;
use mz_auth::AuthenticatorKind;
use mz_auth::user::ExternalUserMetadata;
use mz_build_info::{BuildInfo, DUMMY_BUILD_INFO};
use mz_controller_types::ClusterId;
use mz_ore::metrics::{MetricsFutureExt, MetricsRegistry};
use mz_ore::now::{EpochMillis, NowFn};
use mz_pgwire_common::Format;
use mz_repr::role_id::RoleId;
use mz_repr::user::InternalUserMetadata;
use mz_repr::{CatalogItemId, Datum, Row, RowIterator, SqlScalarType, Timestamp};
use mz_sql::ast::{AstInfo, Raw, Statement, TransactionAccessMode};
use mz_sql::plan::{Params, PlanContext, QueryWhen, StatementDesc};
use mz_sql::session::metadata::SessionMetadata;
use mz_sql::session::user::{
    INTERNAL_USER_NAME_TO_DEFAULT_CLUSTER, RoleMetadata, SYSTEM_USER, User,
};
use mz_sql::session::vars::IsolationLevel;
pub use mz_sql::session::vars::{
    DEFAULT_DATABASE_NAME, EndTransactionAction, SERVER_MAJOR_VERSION, SERVER_MINOR_VERSION,
    SERVER_PATCH_VERSION, SessionVars, Var,
};
use mz_sql_parser::ast::TransactionIsolationLevel;
use mz_storage_client::client::TableData;
use mz_storage_types::sources::Timeline;
use qcell::{QCell, QCellOwner};
use timely::progress::Timestamp as _;
use tokio::sync::mpsc::{self, UnboundedSender};
use tokio::sync::watch;
use uuid::Uuid;

use crate::AdapterNotice;
use crate::ExecutionTime;
use crate::catalog::CatalogState;
use crate::client::RecordFirstRowStream;
use crate::coord::ExplainContext;
use crate::coord::appends::BuiltinTableAppendNotify;
use crate::coord::in_memory_oracle::InMemoryTimestampOracle;
use crate::coord::peek::PeekResponseUnary;
use crate::coord::timestamp_selection::{TimestampContext, TimestampDetermination};
use crate::error::AdapterError;
use crate::metrics::{Metrics, SessionMetrics};
use crate::notice::ExecutionTimeKind;
use crate::statement_logging::PreparedStatementLoggingInfo;
use mz_catalog::durable::Snapshot;

const DUMMY_CONNECTION_ID: ConnectionId = ConnectionId::Static(0);

/// A session holds per-connection state.
#[derive(Derivative)]
#[derivative(Debug)]
pub struct Session {
    conn_id: ConnectionId,
    /// A globally unique identifier for the session. Not to be confused
    /// with `conn_id`, which may be reused.
    uuid: Uuid,
    prepared_statements: BTreeMap<String, PreparedStatement>,
    portals: BTreeMap<String, Portal>,
    transaction: TransactionStatus,
    pcx: Option<PlanContext>,
    metrics: SessionMetrics,
    #[derivative(Debug = "ignore")]
    builtin_updates: Option<BuiltinTableAppendNotify>,

    /// The role metadata of the current session.
    ///
    /// Invariant: role_metadata must be `Some` after the user has
    /// successfully connected to and authenticated with Materialize.
    ///
    /// Prefer using this value over [`SessionConfig::user`].
    //
    // It would be better for this not to be an Option, but the
    // `Session` is initialized before the user has connected to
    // Materialize and is able to look up the `RoleMetadata`. The `Session`
    // is also used to return an error when no role exists and
    // therefore there is no valid `RoleMetadata`.
    role_metadata: Option<RoleMetadata>,
    client_ip: Option<IpAddr>,
    vars: SessionVars,
    notices_tx: mpsc::UnboundedSender<AdapterNotice>,
    notices_rx: mpsc::UnboundedReceiver<AdapterNotice>,
    next_transaction_id: TransactionId,
    secret_key: u32,
    external_metadata_rx: Option<watch::Receiver<ExternalUserMetadata>>,
    // Token allowing us to access `Arc<QCell<StatementLogging>>`
    // metadata. We want these to be reference-counted, because the same
    // statement might be referenced from multiple portals simultaneously.
    //
    // However, they can't be `Rc<RefCell<StatementLogging>>`, because
    // the `Session` is sent around to different threads.
    //
    // On the other hand, they don't need to be
    // `Arc<Mutex<StatementLogging>>`, because they will always be
    // accessed from the same thread that the `Session` is currently
    // on. We express this by gating access with this token.
    #[derivative(Debug = "ignore")]
    qcell_owner: QCellOwner,
    session_oracles: BTreeMap<Timeline, InMemoryTimestampOracle>,
    /// Incremented when session state that is relevant to prepared statement planning changes.
    /// Currently, only changes to `portals` are tracked. Changes to `prepared_statements` don't
    /// need to be tracked, because prepared statements can't depend on other prepared statements.
    /// TODO: We might want to track changes also to session variables.
    /// (`Catalog::transient_revision` similarly tracks changes on the catalog side.)
    state_revision: u64,
}

impl SessionMetadata for Session {
    fn conn_id(&self) -> &ConnectionId {
        &self.conn_id
    }

    fn client_ip(&self) -> Option<&IpAddr> {
        self.client_ip.as_ref()
    }

    fn pcx(&self) -> &PlanContext {
        &self
            .transaction()
            .inner()
            .expect("no active transaction")
            .pcx
    }

    fn role_metadata(&self) -> &RoleMetadata {
        self.role_metadata
            .as_ref()
            .expect("role_metadata invariant violated")
    }

    fn vars(&self) -> &SessionVars {
        &self.vars
    }
}

/// Data structure suitable for passing to other threads that need access to some common Session
/// properties.
#[derive(Debug)]
pub struct SessionMeta {
    conn_id: ConnectionId,
    client_ip: Option<IpAddr>,
    pcx: PlanContext,
    role_metadata: RoleMetadata,
    vars: SessionVars,
}

impl SessionMetadata for SessionMeta {
    fn vars(&self) -> &SessionVars {
        &self.vars
    }

    fn conn_id(&self) -> &ConnectionId {
        &self.conn_id
    }

    fn client_ip(&self) -> Option<&IpAddr> {
        self.client_ip.as_ref()
    }

    fn pcx(&self) -> &PlanContext {
        &self.pcx
    }

    fn role_metadata(&self) -> &RoleMetadata {
        &self.role_metadata
    }
}

/// Configures a new [`Session`].
#[derive(Debug, Clone)]
pub struct SessionConfig {
    /// The connection ID for the session.
    ///
    /// May be reused after the session terminates.
    pub conn_id: ConnectionId,
    /// A universally unique identifier for the session, across all processes,
    /// region, and all time.
    ///
    /// Must not be reused, even after the session terminates.
    pub uuid: Uuid,
    /// The peer address of the client
    pub client_ip: Option<IpAddr>,
    /// The name of the user associated with the session.
    pub user: String,
    /// An optional receiver that the session will periodically check for
    /// updates to a user's external metadata.
    pub external_metadata_rx: Option<watch::Receiver<ExternalUserMetadata>>,
    /// Helm chart version
    pub helm_chart_version: Option<String>,
    /// The authenticator that authenticated this user, if any.
    pub authenticator_kind: AuthenticatorKind,
    /// Groups from JWT claims for OIDC group-to-role sync.
    pub groups: Option<Vec<String>>,
}

impl Session {
    /// Creates a new session for the specified connection ID.
    pub(crate) fn new(
        build_info: &'static BuildInfo,
        config: SessionConfig,
        metrics: SessionMetrics,
    ) -> Session {
        assert_ne!(config.conn_id, DUMMY_CONNECTION_ID);
        Self::new_internal(build_info, config, metrics)
    }

    /// Returns a reference-less collection of data usable by other tasks that don't have ownership
    /// of the Session.
    pub fn meta(&self) -> SessionMeta {
        SessionMeta {
            conn_id: self.conn_id().clone(),
            client_ip: self.client_ip().copied(),
            pcx: self.pcx().clone(),
            role_metadata: self.role_metadata().clone(),
            vars: self.vars.clone(),
        }

        // TODO: soft_assert that these are the same as Session.
    }

    /// Creates new statement logging metadata for a one-off
    /// statement.
    // Normally, such logging information would be created as part of
    // allocating a new prepared statement, and a refcounted handle
    // would be copied from that prepared statement to portals during
    // binding. However, we also support (via `Command::declare`)
    // binding a statement directly to a portal without creating an
    // intermediate prepared statement. Thus, for those cases, a
    // mechanism for generating the logging metadata directly is needed.
    pub(crate) fn mint_logging<A: AstInfo>(
        &self,
        raw_sql: String,
        stmt: Option<&Statement<A>>,
        now: EpochMillis,
    ) -> Arc<QCell<PreparedStatementLoggingInfo>> {
        Arc::new(QCell::new(
            &self.qcell_owner,
            PreparedStatementLoggingInfo::still_to_log(
                raw_sql,
                stmt,
                now,
                "".to_string(),
                self.uuid,
                false,
            ),
        ))
    }

    pub(crate) fn qcell_ro<'a, T2: 'a>(&'a self, cell: &'a Arc<QCell<T2>>) -> &'a T2 {
        self.qcell_owner.ro(&*cell)
    }

    pub(crate) fn qcell_rw<'a, T2: 'a>(&'a mut self, cell: &'a Arc<QCell<T2>>) -> &'a mut T2 {
        self.qcell_owner.rw(&*cell)
    }

    /// Returns a unique ID for the session.
    /// Not to be confused with `connection_id`, which can be reused.
    pub fn uuid(&self) -> Uuid {
        self.uuid
    }

    /// Creates a new dummy session.
    ///
    /// Dummy sessions are intended for use when executing queries on behalf of
    /// the system itself, rather than on behalf of a user.
    pub fn dummy() -> Session {
        let registry = MetricsRegistry::new();
        let metrics = Metrics::register_into(&registry);
        let metrics = metrics.session_metrics();
        let mut dummy = Self::new_internal(
            &DUMMY_BUILD_INFO,
            SessionConfig {
                conn_id: DUMMY_CONNECTION_ID,
                uuid: Uuid::new_v4(),
                user: SYSTEM_USER.name.clone(),
                client_ip: None,
                external_metadata_rx: None,
                helm_chart_version: None,
                authenticator_kind: AuthenticatorKind::None,
                groups: None,
            },
            metrics,
        );
        dummy.initialize_role_metadata(RoleId::User(0));
        dummy
    }

    fn new_internal(
        build_info: &'static BuildInfo,
        SessionConfig {
            conn_id,
            uuid,
            user,
            client_ip,
            mut external_metadata_rx,
            helm_chart_version,
            authenticator_kind,
            groups,
        }: SessionConfig,
        metrics: SessionMetrics,
    ) -> Session {
        let (notices_tx, notices_rx) = mpsc::unbounded_channel();
        let default_cluster = INTERNAL_USER_NAME_TO_DEFAULT_CLUSTER.get(&user);
        let user = User {
            name: user,
            internal_metadata: None,
            external_metadata: external_metadata_rx
                .as_mut()
                .map(|rx| rx.borrow_and_update().clone()),
            authenticator_kind: Some(authenticator_kind),
            groups,
        };
        let mut vars = SessionVars::new_unchecked(build_info, user, helm_chart_version);
        if let Some(default_cluster) = default_cluster {
            vars.set_cluster(default_cluster.clone());
        }
        Session {
            conn_id,
            uuid,
            transaction: TransactionStatus::Default,
            pcx: None,
            metrics,
            builtin_updates: None,
            prepared_statements: BTreeMap::new(),
            portals: BTreeMap::new(),
            role_metadata: None,
            client_ip,
            vars,
            notices_tx,
            notices_rx,
            next_transaction_id: 0,
            secret_key: rand::random(),
            external_metadata_rx,
            qcell_owner: QCellOwner::new(),
            session_oracles: BTreeMap::new(),
            state_revision: 0,
        }
    }

    /// Returns the secret key associated with the session.
    pub fn secret_key(&self) -> u32 {
        self.secret_key
    }

    fn new_pcx(&self, mut wall_time: DateTime<Utc>) -> PlanContext {
        if let Some(mock_time) = self.vars().unsafe_new_transaction_wall_time() {
            wall_time = *mock_time;
        }
        PlanContext::new(wall_time)
    }

    /// Starts an explicit transaction, or changes an implicit to an explicit
    /// transaction.
    pub fn start_transaction(
        &mut self,
        wall_time: DateTime<Utc>,
        access: Option<TransactionAccessMode>,
        isolation_level: Option<TransactionIsolationLevel>,
    ) -> Result<(), AdapterError> {
        // Check that current transaction state is compatible with new `access`
        if let Some(txn) = self.transaction.inner() {
            // `READ WRITE` prohibited if:
            // - Currently in `READ ONLY`
            // - Already performed a query
            let read_write_prohibited = match txn.ops {
                TransactionOps::Peeks { .. } | TransactionOps::Subscribe => {
                    txn.access == Some(TransactionAccessMode::ReadOnly)
                }
                TransactionOps::None
                | TransactionOps::Writes(_)
                | TransactionOps::SingleStatement { .. }
                | TransactionOps::DDL { .. } => false,
            };

            if read_write_prohibited && access == Some(TransactionAccessMode::ReadWrite) {
                return Err(AdapterError::ReadWriteUnavailable);
            }
        }

        match std::mem::take(&mut self.transaction) {
            TransactionStatus::Default => {
                let id = self.next_transaction_id;
                self.next_transaction_id = self.next_transaction_id.wrapping_add(1);
                self.transaction = TransactionStatus::InTransaction(Transaction {
                    completion_owner: TransactionCompletionOwner::Coordinator,
                    pcx: self.new_pcx(wall_time),
                    ops: TransactionOps::None,
                    pending_ddl_inspections: Vec::new(),
                    contains_explain: false,
                    access,
                    id,
                });
            }
            TransactionStatus::Started(mut txn)
            | TransactionStatus::InTransactionImplicit(mut txn)
            | TransactionStatus::InTransaction(mut txn) => {
                if access.is_some() {
                    txn.access = access;
                }
                self.transaction = TransactionStatus::InTransaction(txn);
            }
            TransactionStatus::Failed(_) => unreachable!(),
        };

        if let Some(isolation_level) = isolation_level {
            self.vars
                .set_local_transaction_isolation(isolation_level.into());
        }

        Ok(())
    }

    /// Starts either a single statement or implicit transaction based on the
    /// number of statements, but only if no transaction has been started already.
    pub fn start_transaction_implicit(&mut self, wall_time: DateTime<Utc>, stmts: usize) {
        if let TransactionStatus::Default = self.transaction {
            let id = self.next_transaction_id;
            self.next_transaction_id = self.next_transaction_id.wrapping_add(1);
            let txn = Transaction {
                completion_owner: TransactionCompletionOwner::Session,
                pcx: self.new_pcx(wall_time),
                ops: TransactionOps::None,
                pending_ddl_inspections: Vec::new(),
                contains_explain: false,
                access: None,
                id,
            };
            match stmts {
                1 => self.transaction = TransactionStatus::Started(txn),
                n if n > 1 => self.transaction = TransactionStatus::InTransactionImplicit(txn),
                _ => {}
            }
        }
    }

    /// Starts a single statement transaction, but only if no transaction has been started already.
    pub fn start_transaction_single_stmt(&mut self, wall_time: DateTime<Utc>) {
        self.start_transaction_implicit(wall_time, 1);
    }

    /// Clears a transaction, setting its state to Default and destroying all
    /// portals. Returned are:
    /// - sinks that were started in this transaction and need to be dropped
    /// - the cleared transaction so its operations can be handled
    ///
    /// The [Postgres protocol docs](https://www.postgresql.org/docs/current/protocol-flow.html#PROTOCOL-FLOW-EXT-QUERY) specify:
    /// > a named portal object lasts till the end of the current transaction
    /// and
    /// > An unnamed portal is destroyed at the end of the transaction
    #[must_use]
    pub fn clear_transaction(&mut self) -> TransactionStatus {
        self.portals.clear();
        self.pcx = None;
        self.state_revision += 1;
        mem::take(&mut self.transaction)
    }

    /// Transfers transaction completion to the coordinator. Call before handing it any
    /// transaction-scoped state, including on paths that can fail after installing resources.
    pub(crate) fn require_coordinator_completion(&mut self) {
        if let Some(txn) = self.transaction.inner_mut() {
            txn.completion_owner = TransactionCompletionOwner::Coordinator;
        }
    }

    /// Returns the obligations for a session-owned read-only completion.
    pub(crate) fn read_only_completion(
        &self,
        mut action: EndTransactionAction,
    ) -> Option<ReadOnlyCompletion> {
        let txn = self.transaction.inner()?;
        if txn.completion_owner != TransactionCompletionOwner::Session {
            return None;
        }
        let timestamp = match &txn.ops {
            TransactionOps::None => None,
            TransactionOps::Peeks {
                determination,
                requires_linearization,
                ..
            } => matches!(requires_linearization, RequireLinearization::Required)
                .then(|| determination.timestamp_context.clone()),
            TransactionOps::Writes(_)
            | TransactionOps::Subscribe
            | TransactionOps::SingleStatement { .. }
            | TransactionOps::DDL { .. } => return None,
        };
        if matches!(self.transaction, TransactionStatus::Failed(_)) {
            action = EndTransactionAction::Rollback;
        }
        Some(ReadOnlyCompletion {
            action,
            timestamp: if action == EndTransactionAction::Commit {
                timestamp
            } else {
                None
            },
        })
    }

    /// Marks the current transaction as failed.
    pub fn fail_transaction(mut self) -> Self {
        match self.transaction {
            TransactionStatus::Default => unreachable!(),
            TransactionStatus::Started(txn)
            | TransactionStatus::InTransactionImplicit(txn)
            | TransactionStatus::InTransaction(txn) => {
                self.transaction = TransactionStatus::Failed(txn);
            }
            TransactionStatus::Failed(_) => {}
        };
        self
    }

    /// Returns the current transaction status.
    pub fn transaction(&self) -> &TransactionStatus {
        &self.transaction
    }

    /// Returns the current transaction status.
    pub fn transaction_mut(&mut self) -> &mut TransactionStatus {
        &mut self.transaction
    }

    /// Returns the session's transaction code.
    pub fn transaction_code(&self) -> TransactionCode {
        self.transaction().into()
    }

    /// Retains a SHOW CREATE result for validation if this transaction commits DDL.
    pub(crate) fn retain_catalog_result(
        &mut self,
        command: Arc<crate::catalog::command::CatalogResultCommand>,
    ) {
        self.retain_ddl_inspection(DdlInspection::CatalogResult(command));
    }

    fn retain_ddl_inspection(&mut self, inspection: DdlInspection) {
        if !self.transaction().can_retain_ddl_inspections() {
            return;
        }
        let txn = self
            .transaction_mut()
            .inner_mut()
            .expect("active transaction");
        match &mut txn.ops {
            TransactionOps::DDL { commands, .. } => commands.push(inspection.into_command()),
            _ => txn.pending_ddl_inspections.push(inspection),
        }
    }

    /// Retains a successful statement description for validation at DDL commit.
    pub(crate) fn retain_description(
        &mut self,
        statement: mz_sql::ast::Statement<mz_sql::ast::Raw>,
        param_types: Vec<Option<mz_repr::SqlScalarType>>,
        expected: mz_sql::plan::StatementDesc,
    ) {
        if !self.transaction().can_retain_ddl_inspections()
            || (expected.relation_desc.is_none() && expected.param_types.is_empty())
        {
            return;
        }
        // DECLARE delegates description to its inner statement. EXECUTE and FETCH
        // copy session-owned schemas without describing a catalog query.
        let mut described = &statement;
        while let mz_sql::ast::Statement::Declare(declare) = described {
            described = &declare.stmt;
        }
        if matches!(
            described,
            mz_sql::ast::Statement::Execute(_) | mz_sql::ast::Statement::Fetch(_)
        ) {
            return;
        }
        let command = crate::catalog::command::DescriptionCommand::new(
            statement,
            param_types,
            self.meta(),
            expected,
        );
        self.retain_ddl_inspection(DdlInspection::Description(Arc::new(command)));
    }

    /// Adds operations, rejecting incompatible combinations of reads and writes.
    pub fn add_transaction_ops(&mut self, add_ops: TransactionOps) -> Result<(), AdapterError> {
        self.transaction.add_ops(add_ops)
    }

    /// Returns a channel on which to send notices to the session.
    pub fn retain_notice_transmitter(&self) -> UnboundedSender<AdapterNotice> {
        self.notices_tx.clone()
    }

    /// Adds a notice to the session.
    pub fn add_notice(&self, notice: AdapterNotice) {
        self.add_notices([notice])
    }

    /// Adds multiple notices to the session.
    pub fn add_notices(&self, notices: impl IntoIterator<Item = AdapterNotice>) {
        for notice in notices {
            let _ = self.notices_tx.send(notice);
        }
    }

    /// Awaits a possible notice.
    ///
    /// This method is cancel safe.
    pub async fn recv_notice(&mut self) -> AdapterNotice {
        // This method is cancel safe because recv is cancel safe.
        loop {
            let notice = self
                .notices_rx
                .recv()
                .await
                .expect("Session also holds a sender, so recv won't ever return None");
            match self.notice_filter(notice) {
                Some(notice) => return notice,
                None => continue,
            }
        }
    }

    /// Returns a draining iterator over the notices attached to the session.
    pub fn drain_notices(&mut self) -> Vec<AdapterNotice> {
        let mut notices = Vec::new();
        while let Ok(notice) = self.notices_rx.try_recv() {
            if let Some(notice) = self.notice_filter(notice) {
                notices.push(notice);
            }
        }
        notices
    }

    /// Builds a timing notice if enabled and admitted by `client_min_messages`.
    pub fn execution_time_notice(&self, time: ExecutionTime) -> Option<AdapterNotice> {
        if !self.vars.emit_execution_time_notice() {
            return None;
        }
        self.notice_filter(AdapterNotice::ExecutionTime(time))
    }

    /// Returns `time` with the kind `staged` if the current transaction holds work for its commit,
    /// and `time` unchanged otherwise.
    ///
    /// Must be called before the statement settles, because settling can commit that work. A
    /// statement that follows such work in the same transaction is marked too.
    pub fn maybe_mark_staged(&self, time: ExecutionTime) -> ExecutionTime {
        let staged = self
            .transaction()
            .inner()
            .is_some_and(|txn| txn.ops.stages_work());
        if staged {
            ExecutionTime {
                kind: ExecutionTimeKind::Staged,
                ..time
            }
        } else {
            time
        }
    }

    /// Returns Some if the notice should be reported, otherwise None.
    fn notice_filter(&self, notice: AdapterNotice) -> Option<AdapterNotice> {
        // Filter out low threshold severity.
        let minimum_client_severity = self.vars.client_min_messages();
        let sev = notice.severity();
        if !minimum_client_severity.should_output_to_client(&sev) {
            return None;
        }
        // Filter out notices for other clusters.
        if let AdapterNotice::ClusterReplicaStatusChanged { cluster, .. } = &notice {
            if cluster != self.vars.cluster() {
                return None;
            }
        }
        Some(notice)
    }

    /// Sets the transaction ops to `TransactionOps::None`. Must only be used after
    /// verifying that no transaction anomalies will occur if cleared.
    pub fn clear_transaction_ops(&mut self) {
        if let Some(txn) = self.transaction.inner_mut() {
            txn.ops = TransactionOps::None;
        }
    }

    /// Returns the transaction's read timestamp determination, if set.
    ///
    /// Returns `None` if there is no active transaction, or if the active
    /// transaction is not a read transaction.
    pub fn get_transaction_timestamp_determination(&self) -> Option<TimestampDetermination> {
        match self.transaction.inner() {
            Some(Transaction {
                pcx: _,
                ops: TransactionOps::Peeks { determination, .. },
                access: _,
                id: _,
                pending_ddl_inspections: _,
                contains_explain: _,
                completion_owner: _,
            }) => Some(determination.clone()),
            _ => None,
        }
    }

    /// Whether this session has a timestamp for a read transaction.
    pub fn contains_read_timestamp(&self) -> bool {
        matches!(
            self.transaction.inner(),
            Some(Transaction {
                pcx: _,
                ops: TransactionOps::Peeks {
                    determination: TimestampDetermination {
                        timestamp_context: TimestampContext::TimelineTimestamp { .. },
                        ..
                    },
                    ..
                },
                access: _,
                id: _,
                pending_ddl_inspections: _,
                contains_explain: _,
                completion_owner: _,
            })
        )
    }

    /// Registers the prepared statement under `name`.
    pub fn set_prepared_statement(
        &mut self,
        name: String,
        stmt: Option<Statement<Raw>>,
        raw_sql: String,
        desc: StatementDesc,
        state_revision: StateRevision,
        now: EpochMillis,
    ) {
        let logging = PreparedStatementLoggingInfo::still_to_log(
            raw_sql,
            stmt.as_ref(),
            now,
            name.clone(),
            self.uuid,
            false,
        );
        let statement = PreparedStatement {
            stmt,
            desc,
            state_revision,
            logging: Arc::new(QCell::new(&self.qcell_owner, logging)),
        };
        self.prepared_statements.insert(name, statement);
    }

    /// Removes the prepared statement associated with `name`.
    ///
    /// Returns whether a statement previously existed.
    pub fn remove_prepared_statement(&mut self, name: &str) -> bool {
        self.prepared_statements.remove(name).is_some()
    }

    /// Removes all prepared statements.
    pub fn remove_all_prepared_statements(&mut self) {
        self.prepared_statements.clear();
    }

    /// Retrieves the prepared statement associated with `name`.
    ///
    /// This is unverified and could be incorrect if the underlying catalog has
    /// changed.
    pub fn get_prepared_statement_unverified(&self, name: &str) -> Option<&PreparedStatement> {
        self.prepared_statements.get(name)
    }

    /// Retrieves the prepared statement associated with `name`.
    ///
    /// This is unverified and could be incorrect if the underlying catalog has
    /// changed.
    pub fn get_prepared_statement_mut_unverified(
        &mut self,
        name: &str,
    ) -> Option<&mut PreparedStatement> {
        self.prepared_statements.get_mut(name)
    }

    /// Returns the prepared statements for the session.
    pub fn prepared_statements(&self) -> &BTreeMap<String, PreparedStatement> {
        &self.prepared_statements
    }

    /// Returns the portals for the session.
    pub fn portals(&self) -> &BTreeMap<String, Portal> {
        &self.portals
    }

    /// Binds the specified portal to the specified prepared statement.
    ///
    /// If the prepared statement contains parameters, the values and types of
    /// those parameters must be provided in `params`. It is the caller's
    /// responsibility to ensure that the correct number of parameters is
    /// provided.
    ///
    /// The `results_formats` parameter sets the desired format of the results,
    /// and is stored on the portal.
    pub fn set_portal(
        &mut self,
        portal_name: String,
        desc: StatementDesc,
        stmt: Option<Statement<Raw>>,
        logging: Arc<QCell<PreparedStatementLoggingInfo>>,
        params: Vec<(Datum, SqlScalarType)>,
        result_formats: Vec<Format>,
        state_revision: StateRevision,
    ) -> Result<(), AdapterError> {
        // The empty portal can be silently replaced.
        if !portal_name.is_empty() && self.portals.contains_key(&portal_name) {
            return Err(AdapterError::DuplicateCursor(portal_name));
        }
        self.state_revision += 1;
        let param_types = desc.param_types.clone();
        self.portals.insert(
            portal_name,
            Portal {
                stmt: stmt.map(Arc::new),
                desc,
                state_revision,
                parameters: Params {
                    datums: Row::pack(params.iter().map(|(d, _t)| d)),
                    execute_types: params.into_iter().map(|(_d, t)| t).collect(),
                    expected_types: param_types,
                },
                result_formats,
                state: PortalState::NotStarted,
                logging,
                lifecycle_timestamps: None,
            },
        );
        Ok(())
    }

    /// Removes the specified portal.
    ///
    /// If there is no such portal, this method does nothing. Returns whether that portal existed.
    pub fn remove_portal(&mut self, portal_name: &str) -> bool {
        self.state_revision += 1;
        self.portals.remove(portal_name).is_some()
    }

    /// Retrieves a reference to the specified portal.
    ///
    /// If there is no such portal, returns `None`.
    pub fn get_portal_unverified(&self, portal_name: &str) -> Option<&Portal> {
        self.portals.get(portal_name)
    }

    /// Retrieves a mutable reference to the specified portal.
    ///
    /// If there is no such portal, returns `None`.
    ///
    /// Note: When using the returned `PortalRefMut`, there is no need to increment
    /// `Session::state_revision`, because the portal's meaning is not changed.
    pub fn get_portal_unverified_mut(&mut self, portal_name: &str) -> Option<PortalRefMut<'_>> {
        self.portals.get_mut(portal_name).map(|p| PortalRefMut {
            stmt: &p.stmt,
            desc: &p.desc,
            state_revision: &mut p.state_revision,
            parameters: &mut p.parameters,
            result_formats: &mut p.result_formats,
            logging: &mut p.logging,
            state: &mut p.state,
            lifecycle_timestamps: &mut p.lifecycle_timestamps,
        })
    }

    /// Creates and installs a new portal.
    pub fn create_new_portal(
        &mut self,
        stmt: Option<Statement<Raw>>,
        logging: Arc<QCell<PreparedStatementLoggingInfo>>,
        desc: StatementDesc,
        parameters: Params,
        result_formats: Vec<Format>,
        state_revision: StateRevision,
    ) -> Result<String, AdapterError> {
        self.state_revision += 1;

        // See: https://github.com/postgres/postgres/blob/84f5c2908dad81e8622b0406beea580e40bb03ac/src/backend/utils/mmgr/portalmem.c#L234
        for i in 0usize.. {
            let name = format!("<unnamed portal {}>", i);
            match self.portals.entry(name.clone()) {
                Entry::Occupied(_) => continue,
                Entry::Vacant(entry) => {
                    entry.insert(Portal {
                        stmt: stmt.map(Arc::new),
                        desc,
                        state_revision,
                        parameters,
                        result_formats,
                        state: PortalState::NotStarted,
                        logging,
                        lifecycle_timestamps: None,
                    });
                    return Ok(name);
                }
            }
        }

        coord_bail!("unable to create a new portal");
    }

    /// Resets the session to its initial state.
    ///
    /// Returns the parameters changed by [`SessionVars::reset_all`].
    pub fn reset(&mut self) -> BTreeMap<&'static str, String> {
        let _ = self.clear_transaction();
        self.prepared_statements.clear();
        self.vars.reset_all()
    }

    /// Returns the [application_name] that created this session.
    ///
    /// [application_name]: (https://www.postgresql.org/docs/current/runtime-config-logging.html#GUC-APPLICATION-NAME)
    pub fn application_name(&self) -> &str {
        self.vars.application_name()
    }

    /// Returns a reference to the variables in this session.
    pub fn vars(&self) -> &SessionVars {
        &self.vars
    }

    /// Returns a mutable reference to the variables in this session.
    pub fn vars_mut(&mut self) -> &mut SessionVars {
        &mut self.vars
    }

    /// Drains any external metadata updates and applies the changes from the latest update.
    pub fn apply_external_metadata_updates(&mut self) {
        // If no sender is registered then there isn't anything to do.
        let Some(rx) = &mut self.external_metadata_rx else {
            return;
        };

        // If the value hasn't changed then return.
        if !rx.has_changed().unwrap_or(false) {
            return;
        }

        // Update our metadata! Note the short critical section (just a clone) to avoid blocking
        // the sending side of this watch channel.
        let metadata = rx.borrow_and_update().clone();
        self.vars.set_external_user_metadata(metadata);
    }

    /// Applies the internal user metadata to the session.
    pub fn apply_internal_user_metadata(&mut self, metadata: InternalUserMetadata) {
        self.vars.set_internal_user_metadata(metadata);
    }

    /// Initializes the session's role metadata.
    pub fn initialize_role_metadata(&mut self, role_id: RoleId) {
        self.role_metadata = Some(RoleMetadata::new(role_id));
    }

    /// Ensures that a timestamp oracle exists for `timeline` and returns a mutable reference to
    /// the timestamp oracle.
    pub fn ensure_timestamp_oracle(&mut self, timeline: Timeline) -> &mut InMemoryTimestampOracle {
        self.session_oracles.entry(timeline).or_insert_with(|| {
            InMemoryTimestampOracle::new(Timestamp::minimum(), NowFn::from(Timestamp::minimum))
        })
    }

    /// Ensures that a timestamp oracle exists for reads and writes from/to a local input and
    /// returns a mutable reference to the timestamp oracle.
    pub fn ensure_local_timestamp_oracle(&mut self) -> &mut InMemoryTimestampOracle {
        self.ensure_timestamp_oracle(Timeline::EpochMilliseconds)
    }

    /// Returns a reference to the timestamp oracle for `timeline`.
    pub fn get_timestamp_oracle(&self, timeline: &Timeline) -> Option<&InMemoryTimestampOracle> {
        self.session_oracles.get(timeline)
    }

    /// If the current session is using the Strong Session Serializable isolation level advance the
    /// session local timestamp oracle to `write_ts`.
    pub fn apply_write(&mut self, timestamp: Timestamp) {
        if self.vars().transaction_isolation() == &IsolationLevel::StrongSessionSerializable {
            self.ensure_local_timestamp_oracle().apply_write(timestamp);
        }
    }

    /// Returns the [`SessionMetrics`] instance associated with this [`Session`].
    pub fn metrics(&self) -> &SessionMetrics {
        &self.metrics
    }

    /// Sets the `BuiltinTableAppendNotify` for this session.
    pub fn set_builtin_table_updates(&mut self, fut: BuiltinTableAppendNotify) {
        let prev = self.builtin_updates.replace(fut);
        mz_ore::soft_assert_or_log!(prev.is_none(), "replacing old builtin table notify");
    }

    /// Takes the stashed `BuiltinTableAppendNotify`, if one exists, and returns a [`Future`] that
    /// waits for the writes to complete.
    pub fn clear_builtin_table_updates(&mut self) -> Option<impl Future<Output = ()> + 'static> {
        if let Some(fut) = self.builtin_updates.take() {
            // Record how long we blocked for, if we blocked at all.
            let histogram = self
                .metrics()
                .session_startup_table_writes_seconds()
                .clone();
            Some(async move {
                fut.wall_time().observe(histogram).await;
            })
        } else {
            None
        }
    }

    /// Return the state_revision of the session, which can be used by dependent objects for knowing
    /// when to re-plan due to session state changes.
    pub fn state_revision(&self) -> u64 {
        self.state_revision
    }
}

/// A prepared statement.
#[derive(Derivative, Clone)]
#[derivative(Debug)]
pub struct PreparedStatement {
    stmt: Option<Statement<Raw>>,
    desc: StatementDesc,
    /// The most recent state revision that has verified this statement.
    pub state_revision: StateRevision,
    #[derivative(Debug = "ignore")]
    logging: Arc<QCell<PreparedStatementLoggingInfo>>,
}

impl PreparedStatement {
    /// Returns the AST associated with this prepared statement,
    /// if the prepared statement was not the empty query.
    pub fn stmt(&self) -> Option<&Statement<Raw>> {
        self.stmt.as_ref()
    }

    /// Returns the description of the prepared statement.
    pub fn desc(&self) -> &StatementDesc {
        &self.desc
    }

    /// Returns a handle to the metadata for statement logging.
    pub fn logging(&self) -> &Arc<QCell<PreparedStatementLoggingInfo>> {
        &self.logging
    }
}

/// A portal represents the execution state of a running or runnable query.
#[derive(Derivative)]
#[derivative(Debug)]
pub struct Portal {
    /// The statement that is bound to this portal.
    pub stmt: Option<Arc<Statement<Raw>>>,
    /// The statement description.
    pub desc: StatementDesc,
    /// The most recent state revision that has verified this portal.
    pub state_revision: StateRevision,
    /// The bound values for the parameters in the prepared statement, if any.
    pub parameters: Params,
    /// The desired output format for each column in the result set.
    pub result_formats: Vec<Format>,
    /// A handle to metadata needed for statement logging.
    #[derivative(Debug = "ignore")]
    pub logging: Arc<QCell<PreparedStatementLoggingInfo>>,
    /// The execution state of the portal.
    #[derivative(Debug = "ignore")]
    pub state: PortalState,
    /// Statement lifecycle timestamps coming from `mz-pgwire`.
    pub lifecycle_timestamps: Option<LifecycleTimestamps>,
}

/// A mutable reference to a portal, capturing its state and associated metadata. Importantly, it
/// does _not_ give _mutable_ access to `stmt` and `desc`, which means that you do not need to
/// increment `Session::state_revision` when modifying fields through a `PortalRefMut`, because the
/// portal's meaning is not changed.
pub struct PortalRefMut<'a> {
    /// The statement that is bound to this portal.
    pub stmt: &'a Option<Arc<Statement<Raw>>>,
    /// The statement description.
    pub desc: &'a StatementDesc,
    /// The most recent state revision that has verified this portal.
    pub state_revision: &'a mut StateRevision,
    /// The bound values for the parameters in the prepared statement, if any.
    pub parameters: &'a mut Params,
    /// The desired output format for each column in the result set.
    pub result_formats: &'a mut Vec<Format>,
    /// A handle to metadata needed for statement logging.
    pub logging: &'a mut Arc<QCell<PreparedStatementLoggingInfo>>,
    /// The execution state of the portal.
    pub state: &'a mut PortalState,
    /// Statement lifecycle timestamps coming from `mz-pgwire`.
    pub lifecycle_timestamps: &'a mut Option<LifecycleTimestamps>,
}

/// Points to a revision of catalog state and session state. When the current revisions are not the
/// same as the revisions when a prepared statement or a portal was described, we need to check
/// whether the description is still valid.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct StateRevision {
    /// A revision of the catalog.
    pub catalog_revision: u64,
    /// A revision of the session state.
    pub session_state_revision: u64,
}

/// Execution states of a portal.
pub enum PortalState {
    /// Portal not yet started.
    NotStarted,
    /// Portal is a rows-returning statement in progress with 0 or more rows
    /// remaining.
    InProgress(Option<InProgressRows>),
    /// Portal has completed and should not be re-executed. If the optional string
    /// is present, it is returned as a CommandComplete tag, otherwise an error
    /// is sent.
    Completed(Option<String>),
}

/// State of an in-progress, rows-returning portal.
pub struct InProgressRows {
    /// The current batch of rows.
    pub current: Option<Box<dyn RowIterator + Send + Sync>>,
    /// A stream from which to fetch more row batches.
    pub remaining: RecordFirstRowStream,
}

impl InProgressRows {
    /// Creates a new InProgressRows from a batch stream.
    pub fn new(remaining: RecordFirstRowStream) -> Self {
        Self {
            current: None,
            remaining,
        }
    }

    /// Determines whether the underlying stream has ended and there are also no more rows
    /// stashed in `current`.
    pub fn no_more_rows(&self) -> bool {
        self.remaining.no_more_rows && self.current.is_none()
    }
}

/// A stream of batched rows delivered to a client writer.
///
/// This is a boxed stream, not a bare channel receiver, so the coordinator can
/// insert an adapter between itself and the client writer. The SUBSCRIBE path
/// wraps the receiver in a byte-accounting `map` whose closure type cannot be
/// named, so it has to be boxed. Consumers box the receiver into a
/// `RecordFirstRowStream` regardless, so this adds no allocation.
pub type RowBatchStream = Box<dyn futures::Stream<Item = PeekResponseUnary> + Unpin + Send + Sync>;

/// Part of statement lifecycle. These are timestamps that come from the Adapter frontend
/// (`mz-pgwire`) part of the lifecycle.
#[derive(Debug, Clone)]
pub struct LifecycleTimestamps {
    /// When the query was received. More specifically, when the tokio recv returned.
    /// For a Simple Query, this is for the whole query, for the Extended Query flow, this is only
    /// for `FrontendMessage::Execute`. (This means that this is after parsing for the
    /// Extended Query flow.)
    pub received: EpochMillis,
}

impl LifecycleTimestamps {
    /// Creates a new `LifecycleTimestamps`.
    pub fn new(received: EpochMillis) -> Self {
        Self { received }
    }
}

/// The transaction status of a session.
///
/// PostgreSQL's transaction states are in backend/access/transam/xact.c.
#[derive(Debug)]
pub enum TransactionStatus {
    /// Idle. Matches `TBLOCK_DEFAULT`.
    Default,
    /// Running a transaction started outside of any `BEGIN`. Matches
    /// `TBLOCK_STARTED`. In PostgreSQL, when using the extended query protocol, this
    /// may be upgraded into multi-statement implicit query (see [`Self::InTransactionImplicit`]).
    /// Additionally, some statements may trigger an eager commit of the implicit transaction,
    /// see: <https://git.postgresql.org/gitweb/?p=postgresql.git&a=commitdiff&h=f92944137>. In
    /// Materialize we never upgrade it. We eagerly commit it unless it can take on further
    /// statements of the same pipeline, see [`Self::may_span_pipeline`].
    Started(Transaction),
    /// Currently in a transaction issued from a `BEGIN`. Matches `TBLOCK_INPROGRESS`.
    InTransaction(Transaction),
    /// Currently in an implicit transaction started from a multi-statement query
    /// with more than 1 statements. Matches `TBLOCK_IMPLICIT_INPROGRESS`.
    InTransactionImplicit(Transaction),
    /// In a failed transaction. Matches `TBLOCK_ABORT`.
    Failed(Transaction),
}

impl TransactionStatus {
    /// Extracts the inner transaction ops if not failed.
    pub fn into_ops(self) -> Option<TransactionOps> {
        match self {
            TransactionStatus::Default | TransactionStatus::Failed(_) => None,
            TransactionStatus::Started(txn)
            | TransactionStatus::InTransaction(txn)
            | TransactionStatus::InTransactionImplicit(txn) => Some(txn.ops),
        }
    }

    /// Exposes the inner transaction.
    pub fn inner(&self) -> Option<&Transaction> {
        match self {
            TransactionStatus::Default => None,
            TransactionStatus::Started(txn)
            | TransactionStatus::InTransaction(txn)
            | TransactionStatus::InTransactionImplicit(txn)
            | TransactionStatus::Failed(txn) => Some(txn),
        }
    }

    /// Exposes the inner transaction.
    pub fn inner_mut(&mut self) -> Option<&mut Transaction> {
        match self {
            TransactionStatus::Default => None,
            TransactionStatus::Started(txn)
            | TransactionStatus::InTransaction(txn)
            | TransactionStatus::InTransactionImplicit(txn)
            | TransactionStatus::Failed(txn) => Some(txn),
        }
    }

    /// Whether inspections must be retained for a possible DDL commit.
    ///
    /// This includes observations made before the first DDL statement. Once the
    /// transaction cannot enter DDL mode, retaining their statements and session
    /// metadata would consume memory without any commit-time validation.
    fn can_retain_ddl_inspections(&self) -> bool {
        match self {
            Self::Started(txn) | Self::InTransaction(txn) | Self::InTransactionImplicit(txn) => {
                !txn.contains_explain
                    && match txn.ops {
                        TransactionOps::None | TransactionOps::DDL { .. } => true,
                        TransactionOps::Peeks { .. }
                        | TransactionOps::Writes(_)
                        | TransactionOps::Subscribe
                        | TransactionOps::SingleStatement { .. } => false,
                    }
            }
            Self::Default | Self::Failed(_) => false,
        }
    }

    /// Records EXPLAIN execution, rejecting transactions that already contain DDL.
    pub(crate) fn record_explain(&mut self) -> Result<(), AdapterError> {
        let txn = self.inner_mut().expect("active transaction");
        if matches!(txn.ops, TransactionOps::DDL { .. }) {
            return Err(AdapterError::DDLOnlyTransaction);
        }
        // EXPLAIN outcomes are not replayed at commit, so this transaction must
        // not subsequently enter DDL mode and rebind its writes.
        txn.contains_explain = true;
        txn.pending_ddl_inspections.clear();
        Ok(())
    }

    /// Whether the transaction's ops are DDL.
    pub fn is_ddl(&self) -> bool {
        match self {
            TransactionStatus::Default => false,
            TransactionStatus::Started(txn)
            | TransactionStatus::InTransaction(txn)
            | TransactionStatus::InTransactionImplicit(txn)
            | TransactionStatus::Failed(txn) => {
                matches!(txn.ops, TransactionOps::DDL { .. })
            }
        }
    }

    /// Expresses whether or not the transaction was implicitly started.
    /// However, its negation does not imply explicitly started.
    pub fn is_implicit(&self) -> bool {
        match self {
            TransactionStatus::Started(_) | TransactionStatus::InTransactionImplicit(_) => true,
            TransactionStatus::Default
            | TransactionStatus::InTransaction(_)
            | TransactionStatus::Failed(_) => false,
        }
    }

    /// Whether this implicit transaction may stay open for the rest of its
    /// extended-protocol pipeline, so that the pipeline commits or rolls back as a
    /// unit like PostgreSQL's.
    ///
    /// Only writes may, because they are merely staged. A read already pinned its
    /// timestamp without timedomain read holds, so a second read could neither join
    /// that timestamp nor pick its own.
    pub fn may_span_pipeline(&self) -> bool {
        match self {
            TransactionStatus::Started(txn) => match &txn.ops {
                TransactionOps::Writes(_) => true,
                TransactionOps::None
                | TransactionOps::Peeks { .. }
                | TransactionOps::Subscribe
                | TransactionOps::SingleStatement { .. }
                | TransactionOps::DDL { .. } => false,
            },
            TransactionStatus::Default
            | TransactionStatus::InTransaction(_)
            | TransactionStatus::InTransactionImplicit(_)
            | TransactionStatus::Failed(_) => false,
        }
    }

    /// Whether the transaction may contain multiple statements.
    pub fn is_in_multi_statement_transaction(&self) -> bool {
        match self {
            TransactionStatus::InTransaction(_) | TransactionStatus::InTransactionImplicit(_) => {
                true
            }
            TransactionStatus::Default
            | TransactionStatus::Started(_)
            | TransactionStatus::Failed(_) => false,
        }
    }

    /// Whether a statement other than the current one can belong to this
    /// transaction, so committing here would commit more than this statement.
    ///
    /// This is [`Self::is_in_multi_statement_transaction`] widened to cover the
    /// `Started` trap. An extended-protocol pipeline stays `Started` from its
    /// first statement until `Sync`, so a `Started` transaction already holding
    /// ops has a pipeline accumulating in it, and a statement running now runs
    /// alongside those ops.
    ///
    /// Callers must evaluate this before the current statement stages ops of
    /// its own. Afterwards `contains_ops` reports the statement's own ops and
    /// every statement looks like it shares a transaction.
    ///
    /// Reports true for a `Failed` transaction that holds ops, since
    /// `contains_ops` reads through to the inner transaction in that state.
    /// Callers that treat a failed transaction as one a statement may run in
    /// need their own check, because pgwire admits only `COMMIT` and `ROLLBACK`
    /// once a transaction has failed and so nothing else can observe the
    /// difference.
    pub fn may_share_transaction_with_other_statements(&self) -> bool {
        self.is_in_multi_statement_transaction() || self.contains_ops()
    }

    /// Whether we are in a multi-statement transaction, AND the query is immediate.
    pub fn in_immediate_multi_stmt_txn(&self, when: &QueryWhen) -> bool {
        self.is_in_multi_statement_transaction() && when == &QueryWhen::Immediately
    }

    /// The timeline of the transaction, if one exists.
    pub fn timeline(&self) -> Option<Timeline> {
        match self {
            TransactionStatus::Default => None,
            TransactionStatus::Started(txn)
            | TransactionStatus::InTransaction(txn)
            | TransactionStatus::InTransactionImplicit(txn)
            | TransactionStatus::Failed(txn) => txn.timeline(),
        }
    }

    /// The cluster of the transaction, if one exists.
    pub fn cluster(&self) -> Option<ClusterId> {
        match self {
            TransactionStatus::Default => None,
            TransactionStatus::Started(txn)
            | TransactionStatus::InTransaction(txn)
            | TransactionStatus::InTransactionImplicit(txn)
            | TransactionStatus::Failed(txn) => txn.cluster(),
        }
    }

    /// Snapshot of the catalog that reflects DDL operations run in this transaction.
    pub fn catalog_state(&self) -> Option<&CatalogState> {
        match self.inner() {
            Some(Transaction {
                ops: TransactionOps::DDL { state, .. },
                ..
            }) => Some(state),
            _ => None,
        }
    }

    /// Reports whether any operations have been executed as part of this transaction
    pub fn contains_ops(&self) -> bool {
        match self.inner() {
            Some(txn) => txn.contains_ops(),
            None => false,
        }
    }

    /// Checks whether the current state of this transaction allows writes
    /// (adding write ops).
    pub fn allows_writes(&self) -> bool {
        match self {
            TransactionStatus::Started(Transaction { ops, access, .. })
            | TransactionStatus::InTransaction(Transaction { ops, access, .. })
            | TransactionStatus::InTransactionImplicit(Transaction { ops, access, .. }) => {
                match ops {
                    TransactionOps::None => access != &Some(TransactionAccessMode::ReadOnly),
                    TransactionOps::Peeks { determination, .. } => {
                        // We can switch a peek-only transaction to a write
                        // transaction only if the peeks thus far are constant
                        // (i.e. they do not have a timestamp) and the
                        // transaction is not explicitly marked read-only.
                        access != &Some(TransactionAccessMode::ReadOnly)
                            && !determination.timestamp_context.contains_timestamp()
                    }
                    TransactionOps::Subscribe => false,
                    TransactionOps::Writes(_) => true,
                    TransactionOps::SingleStatement { .. } => false,
                    TransactionOps::DDL { .. } => false,
                }
            }
            TransactionStatus::Default | TransactionStatus::Failed(_) => {
                unreachable!()
            }
        }
    }

    /// Adds operations to the current transaction. An error is produced if they cannot be merged
    /// (i.e., a timestamp-dependent read cannot be merged to an insert).
    ///
    /// The `DDL` variant is an exception and does not merge operations, but instead overwrites the
    /// old ops with the new ops. This is correct because it is only used in conjunction with the
    /// Dry Run catalog op which returns an error containing all of the ops, and those ops are
    /// passed to this function which then overwrites.
    ///
    /// # Panics
    /// If the operations are compatible but the operation metadata doesn't match. Such as reads at
    /// different timestamps, reads on different timelines, reads on different clusters, etc. It's
    /// up to the caller to make sure these are aligned.
    pub fn add_ops(&mut self, mut add_ops: TransactionOps) -> Result<(), AdapterError> {
        if matches!(add_ops, TransactionOps::DDL { .. })
            && self.inner().is_some_and(|txn| txn.contains_explain)
        {
            return Err(AdapterError::DDLOnlyTransaction);
        }
        match self {
            TransactionStatus::Started(Transaction {
                ops,
                access,
                pending_ddl_inspections,
                ..
            })
            | TransactionStatus::InTransaction(Transaction {
                ops,
                access,
                pending_ddl_inspections,
                ..
            })
            | TransactionStatus::InTransactionImplicit(Transaction {
                ops,
                access,
                pending_ddl_inspections,
                ..
            }) => {
                match ops {
                    TransactionOps::None => {
                        if matches!(access, Some(TransactionAccessMode::ReadOnly))
                            && matches!(add_ops, TransactionOps::Writes(_))
                        {
                            return Err(AdapterError::ReadOnlyTransaction);
                        }
                        match &mut add_ops {
                            TransactionOps::DDL { commands, .. } => {
                                // Pre-DDL inspections belong before every mutation, without
                                // restricting the transaction until it actually enters DDL mode.
                                commands.splice(
                                    0..0,
                                    pending_ddl_inspections
                                        .drain(..)
                                        .map(DdlInspection::into_command),
                                );
                            }
                            TransactionOps::None => {}
                            TransactionOps::Peeks { .. }
                            | TransactionOps::Writes(_)
                            | TransactionOps::Subscribe
                            | TransactionOps::SingleStatement { .. } => {
                                pending_ddl_inspections.clear();
                            }
                        }
                        *ops = add_ops;
                    }
                    TransactionOps::Peeks {
                        determination,
                        cluster_id,
                        requires_linearization,
                    } => match add_ops {
                        TransactionOps::Peeks {
                            determination: add_timestamp_determination,
                            cluster_id: add_cluster_id,
                            requires_linearization: add_requires_linearization,
                        } => {
                            assert_eq!(*cluster_id, add_cluster_id);
                            match (
                                &determination.timestamp_context,
                                &add_timestamp_determination.timestamp_context,
                            ) {
                                (
                                    TimestampContext::TimelineTimestamp {
                                        timeline: txn_timeline,
                                        chosen_ts: txn_ts,
                                        oracle_ts: _,
                                    },
                                    TimestampContext::TimelineTimestamp {
                                        timeline: add_timeline,
                                        chosen_ts: add_ts,
                                        oracle_ts: _,
                                    },
                                ) => {
                                    assert_eq!(txn_timeline, add_timeline);
                                    assert_eq!(txn_ts, add_ts);
                                }
                                (TimestampContext::NoTimestamp, _) => {
                                    *determination = add_timestamp_determination
                                }
                                (_, TimestampContext::NoTimestamp) => {}
                            };
                            if matches!(requires_linearization, RequireLinearization::NotRequired)
                                && matches!(
                                    add_requires_linearization,
                                    RequireLinearization::Required
                                )
                            {
                                *requires_linearization = add_requires_linearization;
                            }
                        }
                        // If the peeks thus far are constant (i.e. they do not
                        // have a timestamp), writes can follow and switch the
                        // transaction to a write transaction. But a read-only
                        // transaction must still reject the write. Without that
                        // check the write would silently turn a read-only
                        // transaction into a write transaction, and a later
                        // write would then trip the assert in the `Writes` arm
                        // below.
                        writes @ TransactionOps::Writes(..)
                            if !determination.timestamp_context.contains_timestamp() =>
                        {
                            if matches!(access, Some(TransactionAccessMode::ReadOnly)) {
                                return Err(AdapterError::ReadOnlyTransaction);
                            }
                            *ops = writes;
                        }
                        _ => return Err(AdapterError::ReadOnlyTransaction),
                    },
                    TransactionOps::Subscribe => {
                        return Err(AdapterError::SubscribeOnlyTransaction);
                    }
                    TransactionOps::Writes(txn_writes) => match add_ops {
                        TransactionOps::Writes(mut add_writes) => {
                            // We should have already checked the access above, but make sure we don't miss
                            // it anyway.
                            assert!(!matches!(access, Some(TransactionAccessMode::ReadOnly)));
                            txn_writes.append(&mut add_writes);
                        }
                        // Iff peeks do not have a timestamp (i.e. they are
                        // constant), we can permit them.
                        TransactionOps::Peeks { determination, .. }
                            if !determination.timestamp_context.contains_timestamp() => {}
                        _ => {
                            return Err(AdapterError::WriteOnlyTransaction);
                        }
                    },
                    TransactionOps::SingleStatement { .. } => {
                        return Err(AdapterError::SingleStatementTransaction);
                    }
                    TransactionOps::DDL {
                        commands,
                        state,
                        snapshot,
                    } => match add_ops {
                        TransactionOps::DDL {
                            commands: new_commands,
                            state: new_state,
                            snapshot: new_snapshot,
                        } => {
                            if !new_commands.is_empty() {
                                commands.extend(new_commands);
                                *state = new_state;
                                *snapshot = new_snapshot;
                            }
                        }
                        _ => return Err(AdapterError::DDLOnlyTransaction),
                    },
                }
            }
            TransactionStatus::Default | TransactionStatus::Failed(_) => {
                unreachable!()
            }
        }
        Ok(())
    }
}

/// An abstraction allowing us to identify different transactions.
pub type TransactionId = u64;

impl Default for TransactionStatus {
    fn default() -> Self {
        TransactionStatus::Default
    }
}

/// A catalog result or statement description that must remain valid at DDL commit.
///
/// Retaining inspections before the first DDL statement prevents commit-time
/// replanning from contradicting results already exposed or stored for the client.
/// Buffering them does not itself enter DDL mode. Transactions that never enter
/// DDL mode discard them without validation.
#[derive(Debug)]
enum DdlInspection {
    CatalogResult(Arc<crate::catalog::command::CatalogResultCommand>),
    Description(Arc<crate::catalog::command::DescriptionCommand>),
}

impl DdlInspection {
    fn into_command(self) -> crate::catalog::command::CatalogCommand {
        match self {
            Self::CatalogResult(command) => {
                crate::catalog::command::CatalogCommand::CatalogResult(command)
            }
            Self::Description(command) => {
                crate::catalog::command::CatalogCommand::Description(command)
            }
        }
    }
}

/// State data for transactions.
#[derive(Debug)]
pub struct Transaction {
    completion_owner: TransactionCompletionOwner,
    /// Plan context.
    pub pcx: PlanContext,
    /// Transaction operations.
    pub ops: TransactionOps,
    /// Catalog results and descriptions awaiting a possible transition to DDL mode.
    pending_ddl_inspections: Vec<DdlInspection>,
    contains_explain: bool,
    /// Uniquely identifies the transaction on a per connection basis.
    /// Two transactions started from separate connections may share the
    /// same ID.
    /// If all IDs have been exhausted, this will wrap around back to 0.
    pub id: TransactionId,
    /// Access mode (read only, read write).
    access: Option<TransactionAccessMode>,
}

#[derive(Debug, PartialEq)]
enum TransactionCompletionOwner {
    Session,
    Coordinator,
}

/// Obligations that can be satisfied without transferring the session to the coordinator.
pub(crate) struct ReadOnlyCompletion {
    pub action: EndTransactionAction,
    pub timestamp: Option<TimestampContext>,
}

impl Transaction {
    /// The timeline of the transaction, if one exists.
    fn timeline(&self) -> Option<Timeline> {
        match &self.ops {
            TransactionOps::Peeks {
                determination:
                    TimestampDetermination {
                        timestamp_context: TimestampContext::TimelineTimestamp { timeline, .. },
                        ..
                    },
                ..
            } => Some(timeline.clone()),
            TransactionOps::Peeks { .. }
            | TransactionOps::None
            | TransactionOps::Subscribe
            | TransactionOps::Writes(_)
            | TransactionOps::SingleStatement { .. }
            | TransactionOps::DDL { .. } => None,
        }
    }

    /// The cluster of the transaction, if one exists.
    pub fn cluster(&self) -> Option<ClusterId> {
        match &self.ops {
            TransactionOps::Peeks { cluster_id, .. } => Some(cluster_id.clone()),
            TransactionOps::None
            | TransactionOps::Subscribe
            | TransactionOps::Writes(_)
            | TransactionOps::SingleStatement { .. }
            | TransactionOps::DDL { .. } => None,
        }
    }

    /// Reports whether any operations have been executed as part of this transaction
    fn contains_ops(&self) -> bool {
        !matches!(self.ops, TransactionOps::None)
    }
}

/// A transaction's status code.
#[derive(Debug, Clone, Copy)]
pub enum TransactionCode {
    /// Not currently in a transaction
    Idle,
    /// Currently in a transaction
    InTransaction,
    /// Currently in a transaction block which is failed
    Failed,
}

impl From<TransactionCode> for u8 {
    fn from(code: TransactionCode) -> Self {
        match code {
            TransactionCode::Idle => b'I',
            TransactionCode::InTransaction => b'T',
            TransactionCode::Failed => b'E',
        }
    }
}

impl From<TransactionCode> for String {
    fn from(code: TransactionCode) -> Self {
        char::from(u8::from(code)).to_string()
    }
}

impl From<&TransactionStatus> for TransactionCode {
    /// Convert from the Session's version
    fn from(status: &TransactionStatus) -> TransactionCode {
        match status {
            TransactionStatus::Default => TransactionCode::Idle,
            TransactionStatus::Started(_) => TransactionCode::InTransaction,
            TransactionStatus::InTransaction(_) => TransactionCode::InTransaction,
            TransactionStatus::InTransactionImplicit(_) => TransactionCode::InTransaction,
            TransactionStatus::Failed(_) => TransactionCode::Failed,
        }
    }
}

/// The type of operation being performed by the transaction.
///
/// This is needed because we currently do not allow mixing reads and writes in
/// a transaction. Use this to record what we have done, and what may need to
/// happen at commit.
#[derive(Derivative)]
#[derivative(Debug)]
pub enum TransactionOps {
    /// The transaction has been initiated, but no statement has yet been executed
    /// in it.
    None,
    /// This transaction has had a peek (`SELECT`, `SUBSCRIBE`). If the inner value
    /// is has a timestamp, it must only do other peeks. However, if it doesn't
    /// have a timestamp (i.e. the values are constants), the transaction can still
    /// perform writes.
    Peeks {
        /// The timestamp and timestamp related metadata for the peek.
        determination: TimestampDetermination,
        /// The cluster used to execute peeks.
        cluster_id: ClusterId,
        /// Whether this peek needs to be linearized.
        requires_linearization: RequireLinearization,
    },
    /// This transaction has done a `SUBSCRIBE` and must do nothing else.
    Subscribe,
    /// This transaction has had a write (`INSERT`, `UPDATE`, `DELETE`) and must
    /// only do other writes, or reads whose timestamp is None (i.e. constants).
    Writes(Vec<WriteOp>),
    /// This transaction has a prospective statement that will execute during commit.
    SingleStatement {
        /// The prospective statement.
        stmt: Arc<Statement<Raw>>,
        /// The statement params.
        params: mz_sql::plan::Params,
    },
    /// This transaction has run some _simple_ DDL and must do nothing else. Any statement/plan that
    /// uses this must return false in `must_serialize_ddl()` because this is serialized instead in
    /// `sequence_plan()` during `COMMIT`.
    DDL {
        /// Commands in statement order, retaining outcomes already exposed to the client.
        commands: Vec<crate::catalog::command::CatalogCommand>,
        /// Speculative state after the last accepted command.
        state: CatalogState,
        /// Durable baseline corresponding to `state`.
        snapshot: Snapshot,
    },
}

impl TransactionOps {
    /// Reports whether the transaction holds work for its commit: writes, DDL that the commit
    /// applies, or a statement that runs at commit. A write that affected no rows counts too,
    /// because it leaves a `Writes` op behind.
    pub fn stages_work(&self) -> bool {
        matches!(
            self,
            TransactionOps::Writes(_)
                | TransactionOps::DDL { .. }
                | TransactionOps::SingleStatement { .. }
        )
    }
}

impl Default for TransactionOps {
    fn default() -> Self {
        Self::None
    }
}

/// An `INSERT` waiting to be committed.
#[derive(Debug, Clone, PartialEq)]
pub struct WriteOp {
    /// The target table.
    pub id: CatalogItemId,
    /// The data rows.
    pub rows: TableData,
}

/// Whether a transaction requires linearization.
#[derive(Debug)]
pub enum RequireLinearization {
    /// Linearization is required.
    Required,
    /// Linearization is not required.
    NotRequired,
}

impl From<&ExplainContext> for RequireLinearization {
    fn from(ctx: &ExplainContext) -> Self {
        match ctx {
            ExplainContext::None | ExplainContext::PlanInsightsNotice(_) => {
                RequireLinearization::Required
            }
            _ => RequireLinearization::NotRequired,
        }
    }
}
#[cfg(test)]
mod tests {
    use super::*;
    use mz_sql::session::vars::{SystemVars, VarInput};
    use timely::progress::Antichain;

    fn read_session() -> Session {
        let mut session = Session::dummy();
        session.start_transaction_single_stmt(Utc::now());
        session
            .add_transaction_ops(TransactionOps::Peeks {
                determination: TimestampDetermination {
                    timestamp_context: TimestampContext::TimelineTimestamp {
                        timeline: Timeline::EpochMilliseconds,
                        chosen_ts: 10.into(),
                        oracle_ts: Some(10.into()),
                    },
                    since: Antichain::from_elem(0.into()),
                    upper: Antichain::new(),
                    largest_not_in_advance_of_upper: Timestamp::MAX,
                    oracle_read_ts: Some(10.into()),
                    session_oracle_read_ts: None,
                    real_time_recency_ts: None,
                    constraints: Default::default(),
                },
                cluster_id: ClusterId::User(1),
                requires_linearization: RequireLinearization::Required,
            })
            .expect("read transaction");
        session
    }

    #[mz_ore::test]
    fn read_completion_tracks_ownership_not_just_ops() {
        let mut session = read_session();
        let completion = session
            .read_only_completion(EndTransactionAction::Commit)
            .expect("local read");
        assert_eq!(completion.action, EndTransactionAction::Commit);
        assert!(completion.timestamp.is_some());
        session.require_coordinator_completion();
        assert!(
            session
                .read_only_completion(EndTransactionAction::Commit)
                .is_none()
        );
        let _ = session.clear_transaction();
        session.start_transaction_single_stmt(Utc::now());
        assert!(
            session
                .read_only_completion(EndTransactionAction::Commit)
                .is_some()
        );
    }

    #[mz_ore::test]
    fn completion_skips_timestamp_for_failed_and_rolled_back_reads() {
        let session = read_session();
        let rollback = session
            .read_only_completion(EndTransactionAction::Rollback)
            .expect("local read");
        assert!(rollback.timestamp.is_none());
        let session = session.fail_transaction();
        let failed = session
            .read_only_completion(EndTransactionAction::Commit)
            .expect("failed local read");
        assert_eq!(failed.action, EndTransactionAction::Rollback);
        assert!(failed.timestamp.is_none());
    }

    #[mz_ore::test]
    fn completion_rejects_resource_bearing_transactions() {
        for ops in [
            TransactionOps::Subscribe,
            TransactionOps::Writes(Vec::new()),
        ] {
            let mut session = Session::dummy();
            session.start_transaction_single_stmt(Utc::now());
            session.add_transaction_ops(ops).expect("first operation");
            assert!(
                session
                    .read_only_completion(EndTransactionAction::Commit)
                    .is_none()
            );
            assert!(
                session
                    .fail_transaction()
                    .read_only_completion(EndTransactionAction::Rollback)
                    .is_none()
            );
        }
    }

    #[mz_ore::test]
    fn read_completion_preserves_transactional_parameter_changes() {
        for action in [EndTransactionAction::Commit, EndTransactionAction::Rollback] {
            let mut session = read_session();
            session
                .vars_mut()
                .set_default("application_name", VarInput::Flat("session"))
                .expect("default");
            session
                .vars_mut()
                .set(
                    &SystemVars::new(),
                    "application_name",
                    VarInput::Flat("local"),
                    true,
                )
                .expect("local setting");
            let completion = session.read_only_completion(action).expect("local read");
            let _ = session.clear_transaction();
            let params = session.vars_mut().end_transaction(completion.action);
            assert_eq!(
                params.get("application_name").map(String::as_str),
                Some("session")
            );
            assert_eq!(session.vars().application_name(), "session");
            assert!(matches!(session.transaction(), TransactionStatus::Default));
        }
    }

    #[mz_ore::test]
    fn ddl_inspections_are_released_when_ddl_becomes_impossible() {
        use mz_sql::plan::{Params, StatementDesc};
        use timely::progress::Antichain;

        let statement = mz_sql_parser::parser::parse_statements("SELECT 1")
            .expect("valid SQL")
            .remove(0)
            .ast;
        let modes = [
            TransactionOps::Peeks {
                determination: TimestampDetermination {
                    timestamp_context: TimestampContext::NoTimestamp,
                    since: Antichain::new(),
                    upper: Antichain::new(),
                    largest_not_in_advance_of_upper: mz_repr::Timestamp::MIN,
                    oracle_read_ts: None,
                    session_oracle_read_ts: None,
                    real_time_recency_ts: None,
                    constraints: Default::default(),
                },
                cluster_id: mz_controller_types::ClusterId::User(1),
                requires_linearization: RequireLinearization::NotRequired,
            },
            TransactionOps::Writes(vec![]),
            TransactionOps::Subscribe,
            TransactionOps::SingleStatement {
                stmt: Arc::new(statement.clone()),
                params: Params::empty(),
            },
        ];
        // None exercises EXPLAIN, which leaves the operation mode unchanged.
        for mode in modes.into_iter().map(Some).chain([None]) {
            let mut session = Session::dummy();
            session.start_transaction_implicit(Utc::now(), 2);
            let command = Arc::new(crate::catalog::command::DescriptionCommand::new(
                statement.clone(),
                vec![],
                session.meta(),
                StatementDesc::new(Some(mz_repr::RelationDesc::empty())),
            ));
            let retained = Arc::downgrade(&command);
            session.retain_ddl_inspection(DdlInspection::Description(command));
            assert!(retained.upgrade().is_some(), "DDL is still possible");
            if let Some(mode) = mode {
                session.add_transaction_ops(mode).expect("valid transition");
            } else {
                session
                    .transaction_mut()
                    .record_explain()
                    .expect("EXPLAIN allowed");
            }
            assert!(retained.upgrade().is_none(), "release earlier inspections");
            let command = Arc::new(crate::catalog::command::DescriptionCommand::new(
                statement.clone(),
                vec![],
                session.meta(),
                StatementDesc::new(Some(mz_repr::RelationDesc::empty())),
            ));
            let discarded = Arc::downgrade(&command);
            session.retain_ddl_inspection(DdlInspection::Description(command));
            assert!(
                discarded.upgrade().is_none(),
                "do not retain later inspections"
            );
        }
    }
}