mz_adapter/

client.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;
use std::fmt::{Debug, Display, Formatter};
use std::future::Future;
use std::pin::{self, Pin};
use std::sync::Arc;
use std::time::{Duration, Instant};

use anyhow::bail;
use chrono::{DateTime, Utc};
use derivative::Derivative;
use futures::{FutureExt, Stream, StreamExt};
use itertools::Itertools;
use mz_adapter_types::connection::{ConnectionId, ConnectionIdType};
use mz_auth::password::Password;
use mz_auth::{Authenticated, AuthenticatorKind};
use mz_build_info::BuildInfo;
use mz_compute_types::ComputeInstanceId;
use mz_expr::UnmaterializableFunc;
use mz_expr::{CollectionPlan, RowSetFinishing};
use mz_ore::channel::OneshotReceiverExt;
use mz_ore::collections::CollectionExt;
use mz_ore::id_gen::{IdAllocator, IdAllocatorInnerBitSet, MAX_ORG_ID, org_id_conn_bits};
use mz_ore::instrument;
use mz_ore::now::{EpochMillis, NowFn, to_datetime};
use mz_ore::str::StrExt;
use mz_ore::task::AbortOnDropHandle;
use mz_ore::thread::JoinOnDropHandle;
use mz_ore::tracing::OpenTelemetryContext;
use mz_repr::user::InternalUserMetadata;
use mz_repr::{CatalogItemId, ColumnIndex, SqlScalarType};
use mz_sql::ast::ConstantVisitor;
use mz_sql::ast::{Raw, Statement};
use mz_sql::catalog::{EnvironmentId, SessionCatalog};
use mz_sql::plan::{MutationKind, Plan, ReadThenWritePlan};
use mz_sql::session::hint::ApplicationNameHint;
use mz_sql::session::metadata::SessionMetadata;
use mz_sql::session::user::SUPPORT_USER;
use mz_sql::session::vars::{CLUSTER, OwnedVarInput, SystemVars, Var};
use mz_sql_parser::ast::display::AstDisplay;
use mz_sql_parser::ast::{InsertStatement, StatementKind};
use mz_sql_parser::parser::{ParserStatementError, StatementParseResult};
use prometheus::Histogram;
use serde_json::json;
use tokio::sync::{mpsc, oneshot};
use tracing::{debug, error};
use uuid::Uuid;

use crate::catalog::Catalog;
use crate::command::{
    CatalogDump, CatalogSnapshot, Command, CopyFromStdinWriter, ExecuteResponse, Response,
    SASLChallengeResponse, SASLVerifyProofResponse, SuperuserAttribute,
};
use crate::config::{ScopedParameters, ScopedParametersScope, SystemParameterFrontend};
use crate::coord::read_then_write::DependencyPolicy;
use crate::coord::{Coordinator, ExecuteContextGuard};
use crate::error::AdapterError;
use crate::frontend_read_then_write::{
    FrontendWriteAttemptState, FrontendWriteCancellation, contains_mz_now,
    validate_selection_dependencies,
};
use crate::metrics::Metrics;
use crate::optimize::dataflows::{EvalTime, ExprPrepOneShot};
use crate::optimize::{self, Optimize, OptimizerError};
use crate::peek_client::{CoordinatorClient, ExecutionLogging, TakeOver};
use crate::session::{
    EndTransactionAction, PreparedStatement, Session, SessionConfig, StateRevision, TransactionId,
    TransactionStatus,
};
use crate::statement_logging::{StatementEndedExecutionReason, StatementExecutionStrategy};
use crate::telemetry::{self, EventDetails, SegmentClientExt, StatementFailureType};
use crate::webhook::AppendWebhookResponse;
use crate::{AdapterNotice, AppendWebhookError, PeekClient, PeekResponseUnary, StartupResponse};

/// A handle to a running coordinator.
///
/// The coordinator runs on its own thread. Dropping the handle will wait for
/// the coordinator's thread to exit, which will only occur after all
/// outstanding [`Client`]s for the coordinator have dropped.
pub struct Handle {
    pub(crate) session_id: Uuid,
    pub(crate) start_instant: Instant,
    pub(crate) _thread: JoinOnDropHandle<()>,
}

impl Handle {
    /// Returns the session ID associated with this coordinator.
    ///
    /// The session ID is generated on coordinator boot. It lasts for the
    /// lifetime of the coordinator. Restarting the coordinator will result
    /// in a new session ID.
    pub fn session_id(&self) -> Uuid {
        self.session_id
    }

    /// Returns the instant at which the coordinator booted.
    pub fn start_instant(&self) -> Instant {
        self.start_instant
    }
}

/// A coordinator client.
///
/// A coordinator client is a simple handle to a communication channel with the
/// coordinator. It can be cheaply cloned.
///
/// Clients keep the coordinator alive. The coordinator will not exit until all
/// outstanding clients have dropped.
#[derive(Debug, Clone)]
pub struct Client {
    build_info: &'static BuildInfo,
    inner_cmd_tx: mpsc::UnboundedSender<(OpenTelemetryContext, Command)>,
    id_alloc: IdAllocator<IdAllocatorInnerBitSet>,
    now: NowFn,
    metrics: Metrics,
    environment_id: EnvironmentId,
    segment_client: Option<mz_segment::Client>,
}

impl Client {
    pub(crate) fn new(
        build_info: &'static BuildInfo,
        cmd_tx: mpsc::UnboundedSender<(OpenTelemetryContext, Command)>,
        metrics: Metrics,
        now: NowFn,
        environment_id: EnvironmentId,
        segment_client: Option<mz_segment::Client>,
    ) -> Client {
        // Connection ids are 32 bits and have 3 parts.
        // 1. MSB bit is always 0 because these are interpreted as an i32, and it is possible some
        //    driver will not handle a negative id since postgres has never produced one because it
        //    uses process ids.
        // 2. Next 12 bits are the lower 12 bits of the org id. This allows balancerd to route
        //    incoming cancel messages to a subset of the environments.
        // 3. Last 19 bits are random.
        let env_lower = org_id_conn_bits(&environment_id.organization_id());
        Client {
            build_info,
            inner_cmd_tx: cmd_tx,
            id_alloc: IdAllocator::new(1, MAX_ORG_ID, env_lower),
            now,
            metrics,
            environment_id,
            segment_client,
        }
    }

    /// Allocates a client for an incoming connection.
    pub fn new_conn_id(&self) -> Result<ConnectionId, AdapterError> {
        self.id_alloc.alloc().ok_or(AdapterError::IdExhaustionError)
    }

    /// Creates a new session associated with this client for the given user.
    ///
    /// It is the caller's responsibility to have authenticated the user.
    /// We pass in an Authenticated marker as a guardrail to ensure the
    /// user has authenticated with an authenticator before creating a session.
    pub fn new_session(&self, config: SessionConfig, _authenticated: Authenticated) -> Session {
        // We use the system clock to determine when a session connected to Materialize. This is not
        // intended to be 100% accurate and correct, so we don't burden the timestamp oracle with
        // generating a more correct timestamp.
        Session::new(self.build_info, config, self.metrics().session_metrics())
    }

    /// Used by [mz_auth::AuthenticatorKind::Password]
    /// to verify the provided user's password against the
    /// stored credentials in the catalog.
    pub async fn authenticate(
        &self,
        user: &String,
        password: &Password,
    ) -> Result<Authenticated, AdapterError> {
        let (tx, rx) = oneshot::channel();
        self.send(Command::AuthenticatePassword {
            role_name: user.to_string(),
            password: Some(password.clone()),
            tx,
        });
        rx.await.expect("sender dropped")?;
        Ok(Authenticated)
    }

    /// Used by [mz_auth::AuthenticatorKind::Sasl] for SASL-SCRAM authentication.
    /// This is used prior to [Client::verify_sasl_proof].
    pub async fn generate_sasl_challenge(
        &self,
        user: &String,
        client_nonce: &String,
    ) -> Result<SASLChallengeResponse, AdapterError> {
        let (tx, rx) = oneshot::channel();
        self.send(Command::AuthenticateGetSASLChallenge {
            role_name: user.to_string(),
            nonce: client_nonce.to_string(),
            tx,
        });
        let response = rx.await.expect("sender dropped")?;
        Ok(response)
    }

    /// Used by [mz_auth::AuthenticatorKind::Sasl] for SASL-SCRAM authentication.
    /// This is used after [Client::generate_sasl_challenge].
    pub async fn verify_sasl_proof(
        &self,
        user: &String,
        proof: &String,
        nonce: &String,
        mock_hash: &String,
    ) -> Result<(SASLVerifyProofResponse, Authenticated), AdapterError> {
        let (tx, rx) = oneshot::channel();
        self.send(Command::AuthenticateVerifySASLProof {
            role_name: user.to_string(),
            proof: proof.to_string(),
            auth_message: nonce.to_string(),
            mock_hash: mock_hash.to_string(),
            tx,
        });
        let response = rx.await.expect("sender dropped")?;
        Ok((response, Authenticated))
    }

    /// Checks if a role exists and has the `LOGIN` attribute.
    pub async fn role_can_login(&self, role_name: &str) -> Result<(), AdapterError> {
        let (tx, rx) = oneshot::channel();
        self.send(Command::CheckRoleCanLogin {
            role_name: role_name.to_string(),
            tx,
        });
        rx.await.expect("sender dropped")
    }

    /// Upgrades this client to a session client.
    ///
    /// A session is a connection that has successfully negotiated parameters,
    /// like the user. Most coordinator operations are available only after
    /// upgrading a connection to a session.
    ///
    /// Returns a new client that is bound to the session and a response
    /// containing various details about the startup.
    #[mz_ore::instrument(level = "debug")]
    pub async fn startup(&self, session: Session) -> Result<SessionClient, AdapterError> {
        let user = session.user().clone();
        let conn_id = session.conn_id().clone();
        let secret_key = session.secret_key();
        let uuid = session.uuid();
        let client_ip = session.client_ip();
        let application_name = session.application_name().into();
        let notice_tx = session.retain_notice_transmitter();

        let (tx, rx) = oneshot::channel();

        // ~~SPOOKY ZONE~~
        //
        // This guard prevents a race where the startup command finishes, but the Future returned
        // by this function is concurrently dropped, so we never create a `SessionClient` and thus
        // never cleanup the initialized Session.
        //
        // NOTE: Terminate must only be sent for a successful startup. On a failed startup the
        // Coordinator never registered the connection, so a Terminate for it would trip the
        // Coordinator's "unknown connection" assertion. There is also nothing that needs
        // cleaning up, see the invariant on `Coordinator::handle_startup_inner`.
        let rx = rx.with_guard(|resp: Result<StartupResponse, _>| {
            if resp.is_ok() {
                self.send(Command::Terminate {
                    conn_id: conn_id.clone(),
                    tx: None,
                });
            }
        });

        self.send(Command::Startup {
            tx,
            user,
            conn_id: conn_id.clone(),
            secret_key,
            uuid,
            client_ip: client_ip.copied(),
            application_name,
            notice_tx,
        });

        // When startup fails, no need to call terminate (handle_startup does this). Delay creating
        // the client until after startup to sidestep the panic in its `Drop` implementation.
        let response = rx.await.expect("sender dropped")?;

        // Create the client as soon as startup succeeds (before any await points) so its `Drop` can
        // handle termination.
        // Build the PeekClient with controller handles returned from startup.
        let StartupResponse {
            role_id,
            write_notify,
            session_defaults,
            catalog,
            storage_collections,
            transient_id_gen,
            optimizer_metrics,
            persist_client,
            statement_logging_frontend,
            superuser_attribute,
            occ_write_semaphore,
            group_commit_notifier,
            read_only,
        } = response;

        let peek_client = PeekClient::new(
            CoordinatorClient::Session(self.clone()),
            &catalog,
            storage_collections,
            transient_id_gen,
            optimizer_metrics,
            persist_client,
            statement_logging_frontend,
            occ_write_semaphore,
            group_commit_notifier,
            read_only,
        );

        let mut client = SessionClient {
            inner: Some(self.clone()),
            session: Some(session),
            timeouts: Timeout::new(),
            environment_id: self.environment_id.clone(),
            segment_client: self.segment_client.clone(),
            peek_client,
        };

        let session = client.session();

        // Apply the superuser attribute to the session's user if
        // it exists.
        if let SuperuserAttribute(Some(superuser)) = superuser_attribute {
            session.apply_internal_user_metadata(InternalUserMetadata { superuser });
        }

        session.initialize_role_metadata(role_id);
        let vars_mut = session.vars_mut();
        for (name, val) in session_defaults {
            if let Err(err) = vars_mut.set_default(&name, val.borrow()) {
                // Note: erroring here is unexpected, but we don't want to panic if somehow our
                // assumptions are wrong.
                tracing::error!("failed to set peristed default, {err:?}");
            }
        }
        session
            .vars_mut()
            .end_transaction(EndTransactionAction::Commit);

        // Stash the future that notifies us of builtin table writes completing, we'll block on
        // this future before allowing queries from this session against relevant relations.
        //
        // Note: We stash the future as opposed to waiting on it here to prevent blocking session
        // creation on builtin table updates. This improves the latency for session creation and
        // reduces scheduling load on any dataflows that read from these builtin relations, since
        // it allows updates to be batched.
        session.set_builtin_table_updates(write_notify);

        let catalog = catalog.for_session(session);

        let cluster_active = session.vars().cluster().to_string();
        if session.vars().welcome_message() {
            let cluster_info = if catalog.resolve_cluster(Some(&cluster_active)).is_err() {
                format!("{cluster_active} (does not exist)")
            } else {
                cluster_active.to_string()
            };

            // Emit a welcome message, optimized for readability by humans using
            // interactive tools. If you change the message, make sure that it
            // formats nicely in both `psql` and the console's SQL shell.
            session.add_notice(AdapterNotice::Welcome(format!(
                "connected to Materialize v{}
  Environment ID: {}
  Region: {}
  User: {}
  Cluster: {}
  Database: {}
  {}
  Session UUID: {}

Issue a SQL query to get started. Need help?
  View documentation: https://materialize.com/s/docs
  Join our Slack community: https://materialize.com/s/chat
    ",
                session.vars().build_info().semver_version(),
                self.environment_id,
                self.environment_id.region(),
                session.vars().user().name,
                cluster_info,
                session.vars().database(),
                match session.vars().search_path() {
                    [schema] => format!("Schema: {}", schema),
                    schemas => format!(
                        "Search path: {}",
                        schemas.iter().map(|id| id.to_string()).join(", ")
                    ),
                },
                session.uuid(),
            )));
        }

        if session.vars().current_object_missing_warnings() {
            if catalog.active_database().is_none() {
                let db = session.vars().database().into();
                session.add_notice(AdapterNotice::UnknownSessionDatabase(db));
            }
        }

        // Users stub their toe on their default cluster not existing, so we provide a notice to
        // help guide them on what do to.
        let cluster_var = session
            .vars()
            .inspect(CLUSTER.name())
            .expect("cluster should exist");
        if session.vars().current_object_missing_warnings()
            && catalog.resolve_cluster(Some(&cluster_active)).is_err()
        {
            let cluster_notice = 'notice: {
                if cluster_var.inspect_session_value().is_some() {
                    break 'notice Some(AdapterNotice::DefaultClusterDoesNotExist {
                        name: cluster_active,
                        kind: "session",
                        suggested_action: "Pick an extant cluster with SET CLUSTER = name. Run SHOW CLUSTERS to see available clusters.".into(),
                    });
                }

                let role_default = catalog.get_role(catalog.active_role_id());
                let role_cluster = match role_default.vars().get(CLUSTER.name()) {
                    Some(OwnedVarInput::Flat(name)) => Some(name),
                    None => None,
                    // This is unexpected!
                    Some(v @ OwnedVarInput::SqlSet(_)) => {
                        tracing::warn!(?v, "SqlSet found for cluster Role Default");
                        break 'notice None;
                    }
                };

                let alter_role = "with `ALTER ROLE <role> SET cluster TO <cluster>;`";
                match role_cluster {
                    // If there is no default, suggest a Role default.
                    None => Some(AdapterNotice::DefaultClusterDoesNotExist {
                        name: cluster_active,
                        kind: "system",
                        suggested_action: format!(
                            "Set a default cluster for the current role {alter_role}."
                        ),
                    }),
                    // If the default does not exist, suggest to change it.
                    Some(_) => Some(AdapterNotice::DefaultClusterDoesNotExist {
                        name: cluster_active,
                        kind: "role",
                        suggested_action: format!(
                            "Change the default cluster for the current role {alter_role}."
                        ),
                    }),
                }
            };

            if let Some(notice) = cluster_notice {
                session.add_notice(notice);
            }
        }

        Ok(client)
    }

    /// Cancels the query currently running on the specified connection.
    pub fn cancel_request(&self, conn_id: ConnectionIdType, secret_key: u32) {
        self.send(Command::CancelRequest {
            conn_id,
            secret_key,
        });
    }

    /// Executes a single SQL statement that returns rows as the
    /// `mz_support` user.
    pub async fn support_execute_one(
        &self,
        sql: &str,
    ) -> Result<Pin<Box<dyn Stream<Item = PeekResponseUnary> + Send>>, anyhow::Error> {
        // Connect to the coordinator.
        let conn_id = self.new_conn_id()?;
        let session = self.new_session(
            SessionConfig {
                conn_id,
                uuid: Uuid::new_v4(),
                user: SUPPORT_USER.name.clone(),
                client_ip: None,
                external_metadata_rx: None,
                helm_chart_version: None,
                authenticator_kind: AuthenticatorKind::None,
                groups: None,
            },
            Authenticated,
        );
        let mut session_client = self.startup(session).await?;

        // Parse the SQL statement.
        let stmts = mz_sql::parse::parse(sql)?;
        if stmts.len() != 1 {
            bail!("must supply exactly one query");
        }
        let StatementParseResult { ast: stmt, sql } = stmts.into_element();

        const EMPTY_PORTAL: &str = "";
        session_client.start_transaction(Some(1))?;
        session_client
            .declare(EMPTY_PORTAL.into(), stmt, sql.to_string())
            .await?;

        let execute_result = session_client
            .execute(EMPTY_PORTAL.into(), futures::future::pending(), None)
            .await?;
        match execute_result {
            (ExecuteResponse::SendingRowsStreaming { mut rows, .. }, _) => {
                // We have to only drop the session client _after_ we read the
                // result. Otherwise the peek will get cancelled right when we
                // drop the session client. So we wrap it up in an extra stream
                // like this, which owns the client and can return it.
                let owning_response_stream = async_stream::stream! {
                    while let Some(rows) = rows.next().await {
                        yield rows;
                    }
                    drop(session_client);
                };
                Ok(Box::pin(owning_response_stream))
            }
            r => bail!("unsupported response type: {r:?}"),
        }
    }

    /// Returns the metrics associated with the adapter layer.
    pub fn metrics(&self) -> &Metrics {
        &self.metrics
    }

    /// The current time according to the [`Client`].
    pub fn now(&self) -> DateTime<Utc> {
        to_datetime((self.now)())
    }

    /// Get a metadata and a channel that can be used to append to a webhook source.
    pub async fn get_webhook_appender(
        &self,
        database: String,
        schema: String,
        name: String,
    ) -> Result<AppendWebhookResponse, AppendWebhookError> {
        let (tx, rx) = oneshot::channel();

        // Send our request.
        self.send(Command::GetWebhook {
            database,
            schema,
            name,
            tx,
        });

        // Using our one shot channel to get the result, returning an error if the sender dropped.
        let response = rx
            .await
            .map_err(|_| anyhow::anyhow!("failed to receive webhook response"))?;

        response
    }

    /// Gets the current value of all system variables.
    pub async fn get_system_vars(&self) -> SystemVars {
        let (tx, rx) = oneshot::channel();
        self.send(Command::GetSystemVars { tx });
        rx.await.expect("coordinator unexpectedly gone")
    }

    /// Returns a snapshot of the catalog.
    ///
    /// Does a Coordinator round-trip. Session-bound callers should
    /// prefer [`SessionClient::catalog_snapshot`], which serves from the
    /// session's snapshot cache.
    pub async fn catalog_snapshot_expensive(&self) -> Arc<Catalog> {
        let (tx, rx) = oneshot::channel();
        self.send(Command::CatalogSnapshot { tx });
        let CatalogSnapshot { catalog } = rx.await.expect("coordinator unexpectedly gone");
        catalog
    }

    /// Reconciles the coordinator's scoped feature-flag working copy towards
    /// `overrides`. Used by the system-parameter sync loop from continuous
    /// LaunchDarkly evaluation.
    ///
    /// `prune_scope` bounds which objects' rows the reconcile may remove (the
    /// objects `overrides` was evaluated for). The sync loop passes the live
    /// objects from its snapshot. See
    /// [`crate::catalog::Op::UpdateScopedSystemParameters`].
    pub async fn update_scoped_system_parameters(
        &self,
        overrides: ScopedParameters,
        prune_scope: ScopedParametersScope,
    ) {
        let (tx, rx) = oneshot::channel();
        self.send(Command::UpdateScopedSystemParameters {
            overrides,
            prune_scope,
            tx,
        });
        let _ = rx.await;
    }

    /// Installs (or replaces) the shared system-parameter frontend on the
    /// coordinator, letting the create-cluster / create-replica paths resolve a
    /// new object's scoped overrides synchronously. Sent by the sync loop each
    /// time it (re)initializes the frontend. Fire-and-forget.
    pub fn install_scoped_system_parameter_frontend(&self, frontend: Arc<SystemParameterFrontend>) {
        self.send(Command::InstallScopedSystemParameterFrontend { frontend });
    }

    #[instrument(level = "debug")]
    pub(crate) fn try_send(&self, cmd: Command) -> bool {
        self.inner_cmd_tx
            .send((OpenTelemetryContext::obtain(), cmd))
            .is_ok()
    }

    #[instrument(level = "debug")]
    pub(crate) fn send(&self, cmd: Command) {
        assert!(self.try_send(cmd), "coordinator unexpectedly gone");
    }
}

/// A coordinator client that is bound to a connection.
///
/// See also [`Client`].
pub struct SessionClient {
    // Invariant: inner may only be `None` after the session has been terminated.
    // Once the session is terminated, no communication to the Coordinator
    // should be attempted.
    inner: Option<Client>,
    // Invariant: session may only be `None` during a method call. Every public
    // method must ensure that `Session` is `Some` before it returns.
    session: Option<Session>,
    timeouts: Timeout,
    segment_client: Option<mz_segment::Client>,
    environment_id: EnvironmentId,
    /// Client for frontend peek sequencing; populated at connection startup.
    peek_client: PeekClient,
}

impl SessionClient {
    /// Parses a SQL expression, reporting failures as a telemetry event if
    /// possible.
    pub fn parse<'a>(
        &self,
        sql: &'a str,
    ) -> Result<Result<Vec<StatementParseResult<'a>>, ParserStatementError>, String> {
        match mz_sql::parse::parse_with_limit(sql) {
            Ok(Err(e)) => {
                self.track_statement_parse_failure(&e);
                Ok(Err(e))
            }
            r => r,
        }
    }

    fn track_statement_parse_failure(&self, parse_error: &ParserStatementError) {
        let session = self.session.as_ref().expect("session invariant violated");
        let Some(user_id) = session.user().external_metadata.as_ref().map(|m| m.user_id) else {
            return;
        };
        let Some(segment_client) = &self.segment_client else {
            return;
        };
        let Some(statement_kind) = parse_error.statement else {
            return;
        };
        let Some((action, object_type)) = telemetry::analyze_audited_statement(statement_kind)
        else {
            return;
        };
        let event_type = StatementFailureType::ParseFailure;
        let event_name = format!(
            "{} {} {}",
            object_type.as_title_case(),
            action.as_title_case(),
            event_type.as_title_case(),
        );
        segment_client.environment_track(
            &self.environment_id,
            event_name,
            json!({
                "statement_kind": statement_kind,
                "error": &parse_error.error,
            }),
            EventDetails {
                user_id: Some(user_id),
                application_name: Some(session.application_name()),
                ..Default::default()
            },
        );
    }

    // Verify and return the named prepared statement. We need to verify each use
    // to make sure the prepared statement is still safe to use.
    pub async fn get_prepared_statement(
        &mut self,
        name: &str,
    ) -> Result<&PreparedStatement, AdapterError> {
        let catalog = self.catalog_snapshot("get_prepared_statement").await;
        Coordinator::verify_prepared_statement(&catalog, self.session(), name)?;
        Ok(self
            .session()
            .get_prepared_statement_unverified(name)
            .expect("must exist"))
    }

    /// Saves the parsed statement as a prepared statement.
    ///
    /// The prepared statement is saved in the connection's [`crate::session::Session`]
    /// under the specified name.
    pub async fn prepare(
        &mut self,
        name: String,
        stmt: Option<Statement<Raw>>,
        sql: String,
        param_types: Vec<Option<SqlScalarType>>,
    ) -> Result<(), AdapterError> {
        let catalog = self.catalog_snapshot("prepare").await;

        // Note: This failpoint is used to simulate a request outliving the external connection
        // that made it.
        let mut async_pause = false;
        (|| {
            fail::fail_point!("async_prepare", |val| {
                async_pause = val.map_or(false, |val| val.parse().unwrap_or(false))
            });
        })();
        if async_pause {
            tokio::time::sleep(Duration::from_secs(1)).await;
        };

        let desc = Coordinator::describe(&catalog, self.session(), stmt.clone(), param_types)?;
        let now = self.now();
        let state_revision = StateRevision {
            catalog_revision: catalog.transient_revision(),
            session_state_revision: self.session().state_revision(),
        };
        self.session()
            .set_prepared_statement(name, stmt, sql, desc, state_revision, now);
        Ok(())
    }

    /// Binds a statement to a portal.
    #[mz_ore::instrument(level = "debug")]
    pub async fn declare(
        &mut self,
        name: String,
        stmt: Statement<Raw>,
        sql: String,
    ) -> Result<(), AdapterError> {
        let catalog = self.catalog_snapshot("declare").await;
        let param_types = vec![];
        let desc =
            Coordinator::describe(&catalog, self.session(), Some(stmt.clone()), param_types)?;
        let params = vec![];
        let result_formats = vec![mz_pgwire_common::Format::Text; desc.arity()];
        let now = self.now();
        let logging = self.session().mint_logging(sql, Some(&stmt), now);
        let state_revision = StateRevision {
            catalog_revision: catalog.transient_revision(),
            session_state_revision: self.session().state_revision(),
        };
        self.session().set_portal(
            name,
            desc,
            Some(stmt),
            logging,
            params,
            result_formats,
            state_revision,
        )?;
        Ok(())
    }

    /// Executes a previously-bound portal.
    ///
    /// Note: the provided `cancel_future` must be cancel-safe as it's polled in a `select!` loop.
    ///
    /// `outer_ctx_extra` is Some when we are executing as part of an outer statement, e.g., a FETCH
    /// triggering the execution of the underlying query.
    #[mz_ore::instrument(level = "debug")]
    pub async fn execute(
        &mut self,
        portal_name: String,
        cancel_future: impl Future<Output = std::io::Error> + Send,
        outer_ctx_extra: Option<ExecuteContextGuard>,
    ) -> Result<(ExecuteResponse, Instant), AdapterError> {
        let execute_started = Instant::now();
        let cancel_future = cancel_future.map(|_| ()).shared();

        // Owning the end-of-execution obligation in this frame is what lets
        // cancellation report an error: the inner future can be dropped without
        // the obligation going with it. See `ExecutionLogging`.
        let mut logging = ExecutionLogging::adopt(outer_ctx_extra, &self.peek_client);

        let result = self
            .execute_attempts(portal_name, &mut logging, cancel_future)
            .await;

        logging.retire(&result);

        result.map(|response| (response, execute_started))
    }

    /// Runs the execution paths in order of preference: frontend peek
    /// sequencing, frontend read-then-write sequencing, and finally the
    /// coordinator via `Command::Execute`.
    async fn execute_attempts(
        &mut self,
        portal_name: String,
        logging: &mut ExecutionLogging,
        cancel_future: impl Future<Output = ()> + Send + Clone,
    ) -> Result<ExecuteResponse, AdapterError> {
        // Unroll SQL `EXECUTE <prepared> (...)` so the inner statement
        // flows through `try_frontend_peek` /
        // `try_frontend_read_then_write` below, rather than being
        // re-dispatched via `Command::Execute` from the coordinator's
        // `Plan::Execute` handler. Without this, a prepared statement
        // would route differently from the same statement issued
        // directly.
        //
        // On a successful unroll, `unroll_sql_execute` also takes over
        // statement logging on the outer portal, so that
        // `mz_statement_execution_history` records `EXECUTE foo (...)`
        // rather than the inner SQL.
        let portal_name = self.unroll_sql_execute(portal_name, logging).await?;

        // Attempt peek sequencing in the session task.
        // If unsupported, fall back to the Coordinator path.
        // TODO(peek-seq): wire up cancel_future
        let peek_result = self.try_frontend_peek(&portal_name, logging).await?;
        if let Some(resp) = peek_result {
            debug!("frontend peek succeeded");
            return Ok(resp);
        }
        debug!("frontend peek did not happen, trying frontend read-then-write");

        // Attempt read-then-write sequencing in the session task.
        let rtw_result = self
            .try_frontend_read_then_write_with_cancel(&portal_name, logging, cancel_future.clone())
            .await?;
        if let Some(resp) = rtw_result {
            debug!("frontend read-then-write succeeded");
            return Ok(resp);
        }
        debug!("frontend read-then-write did not happen, falling back to `Command::Execute`");

        // No frontend path took the statement over, so the coordinator retires
        // whatever entry we hold, or begins its own if we hold none.
        let outer_ctx_extra = logging.release();
        self.send_with_cancel(
            |tx, session| Command::Execute {
                portal_name,
                session,
                tx,
                outer_ctx_extra,
            },
            cancel_future,
        )
        .await
    }

    /// If the named portal binds a SQL `EXECUTE <prepared>`, resolve the
    /// prepared statement, install a fresh portal for the inner statement
    /// (carrying the EXECUTE's actual parameter values), and return that
    /// portal's name so the caller can run `try_frontend_peek` /
    /// `try_frontend_read_then_write` against it.
    ///
    /// Only ever unrolls one level: the parser rejects
    /// `PREPARE foo AS EXECUTE bar` (matching Postgres), so the inner
    /// statement is guaranteed not to be another `EXECUTE`. A failsafe below
    /// surfaces an internal error if that invariant is ever violated.
    ///
    /// When the portal does not bind an `EXECUTE`, the common case, returns the
    /// original portal name, costing only a portal lookup.
    async fn unroll_sql_execute(
        &mut self,
        portal_name: String,
        logging: &mut ExecutionLogging,
    ) -> Result<String, AdapterError> {
        let (stmt, params, outer_logging, outer_lifecycle_timestamps) = {
            let session = self.session.as_ref().expect("SessionClient invariant");
            let portal = match session.get_portal_unverified(&portal_name) {
                Some(p) => p,
                // No portal: let `try_frontend_peek` /
                // `try_frontend_read_then_write` surface the
                // standard "missing portal" error.
                None => return Ok(portal_name),
            };
            match &portal.stmt {
                Some(stmt) => (
                    Arc::clone(stmt),
                    portal.parameters.clone(),
                    Arc::clone(&portal.logging),
                    portal.lifecycle_timestamps.clone(),
                ),
                None => return Ok(portal_name),
            }
        };

        // Only EXECUTE statements need unrolling. Bail out before taking a
        // catalog snapshot in the (overwhelmingly common) non-EXECUTE case.
        if !matches!(&*stmt, Statement::Execute(_)) {
            return Ok(portal_name);
        }

        let catalog = self.catalog_snapshot("unroll_sql_execute").await;

        // Validate the outer EXECUTE portal against the (possibly newer)
        // catalog: ensures the recorded portal description still matches
        // what describing the EXECUTE would produce now.
        {
            let session = self.session.as_mut().expect("SessionClient invariant");
            Coordinator::verify_portal(&catalog, session, &portal_name)?;
        }

        // Take over the EXECUTE itself, so that a planning error below produces
        // an errored end-event in `mz_statement_execution_history` rather than
        // no entry at all. We pass the *outer* portal's `logging` and
        // pgwire-bound `params` so the recorded entry shows the user-visible
        // `EXECUTE foo (...)`, not the inner SQL. The inner statement gets its
        // own `query_total` increment when a frontend path takes it over, or,
        // on bailout, in the coordinator's `handle_execute`.
        logging.take_over(
            &self.peek_client,
            self.session.as_mut().expect("SessionClient invariant"),
            Some(&stmt),
            &params,
            &outer_logging,
            &catalog,
            outer_lifecycle_timestamps,
            TakeOver::UnrolledExecute,
        );

        self.install_inner_portal_for_execute(&catalog, &stmt, &params)
    }

    /// Helper for [`Self::unroll_sql_execute`]: plans the outer
    /// `Statement::Execute`, verifies the referenced prepared statement, and
    /// installs a fresh portal carrying the inner statement plus the
    /// EXECUTE's bound parameter values. Returns the new portal's name.
    fn install_inner_portal_for_execute(
        &mut self,
        catalog: &Arc<Catalog>,
        stmt: &Arc<Statement<Raw>>,
        params: &mz_sql::plan::Params,
    ) -> Result<String, AdapterError> {
        use mz_sql::plan::Plan;

        let execute_plan = {
            let session = self.session.as_mut().expect("SessionClient invariant");
            let conn_catalog = catalog.for_session(session);
            let (resolved_stmt, resolved_ids) =
                mz_sql::names::resolve(&conn_catalog, (**stmt).clone())?;
            let pcx = session.pcx();
            let (plan, _sql_impl_ids) = mz_sql::plan::plan(
                Some(pcx),
                &conn_catalog,
                resolved_stmt,
                params,
                &resolved_ids,
            )?;
            match plan {
                Plan::Execute(plan) => plan,
                other => {
                    // Planning a `Statement::Execute` must yield
                    // `Plan::Execute`. If it doesn't, the planner
                    // contract is broken.
                    return Err(AdapterError::Internal(format!(
                        "planning Statement::Execute yielded unexpected plan: {:?}",
                        mz_sql::plan::PlanKind::from(&other),
                    )));
                }
            }
        };

        // Verify and install the inner portal. The new portal carries the
        // inner prepared statement's `logging`, but the logging slot already
        // holds the EXECUTE-level entry, so a frontend path taking the inner
        // statement over inherits that entry instead of starting fresh from
        // this portal.
        let session = self.session.as_mut().expect("SessionClient invariant");
        Coordinator::verify_prepared_statement(catalog, session, &execute_plan.name)?;
        let ps = session
            .get_prepared_statement_unverified(&execute_plan.name)
            .expect("verified above");
        let inner_stmt = ps.stmt().cloned();
        let inner_desc = ps.desc().clone();
        let state_revision = ps.state_revision;
        let inner_logging = Arc::clone(ps.logging());

        // Failsafe: `PREPARE foo AS EXECUTE bar` is rejected by the parser,
        // so the resolved inner statement must not be another `EXECUTE`. If
        // that ever changes, we'd silently skip frontend sequencing for the
        // deeper EXECUTEs. Surface it as an internal error instead.
        if let Some(inner) = inner_stmt.as_ref() {
            if matches!(inner, Statement::Execute(_)) {
                return Err(AdapterError::Internal(format!(
                    "nested EXECUTE: prepared statement {} resolves to another EXECUTE; \
                     parser should reject `PREPARE ... AS EXECUTE ...`",
                    execute_plan.name.quoted(),
                )));
            }
        }

        session.create_new_portal(
            inner_stmt,
            inner_logging,
            inner_desc,
            execute_plan.params,
            Vec::new(),
            state_revision,
        )
    }

    fn now(&self) -> EpochMillis {
        (self.inner().now)()
    }

    fn now_datetime(&self) -> DateTime<Utc> {
        to_datetime(self.now())
    }

    /// Starts a transaction based on implicit:
    /// - `None`: InTransaction
    /// - `Some(1)`: Started
    /// - `Some(n > 1)`: InTransactionImplicit
    /// - `Some(0)`: no change
    pub fn start_transaction(&mut self, implicit: Option<usize>) -> Result<(), AdapterError> {
        let now = self.now_datetime();
        let session = self.session.as_mut().expect("session invariant violated");
        let result = match implicit {
            None => session.start_transaction(now, None, None),
            Some(stmts) => {
                session.start_transaction_implicit(now, stmts);
                Ok(())
            }
        };
        result
    }

    /// Ends a transaction. Even if an error is returned, guarantees that the transaction in the
    /// session and Coordinator has cleared its state.
    #[instrument(level = "debug")]
    pub async fn end_transaction(
        &mut self,
        action: EndTransactionAction,
    ) -> Result<ExecuteResponse, AdapterError> {
        let res = self
            .send(|tx, session| Command::Commit {
                action,
                session,
                tx,
            })
            .await;
        // Commit isn't guaranteed to set the session's state to anything specific, so clear it
        // here. It's safe to ignore the returned `TransactionStatus` because that doesn't contain
        // any data that the Coordinator must act on for correctness.
        let _ = self.session().clear_transaction();
        res
    }

    /// Fails a transaction.
    pub fn fail_transaction(&mut self) {
        let session = self.session.take().expect("session invariant violated");
        let session = session.fail_transaction();
        self.session = Some(session);
    }

    /// Fetches the catalog, served from the session-side snapshot cache when
    /// the catalog is unchanged since the cached snapshot was taken. See
    /// [`PeekClient::catalog_snapshot`].
    #[instrument(level = "debug")]
    pub async fn catalog_snapshot(&mut self, context: &str) -> Arc<Catalog> {
        self.peek_client.catalog_snapshot(context).await
    }

    /// Reports whether `enable_statement_arrival_logging` is on.
    pub async fn statement_arrival_logging_enabled(&mut self) -> bool {
        let catalog = self.catalog_snapshot("statement_arrival_logging").await;
        catalog.system_config().enable_statement_arrival_logging()
    }

    /// Reports whether `enable_extended_protocol_implicit_transaction` is on.
    pub async fn extended_protocol_implicit_transaction_enabled(&mut self) -> bool {
        let catalog = self
            .catalog_snapshot("extended_protocol_implicit_transaction")
            .await;
        catalog
            .system_config()
            .enable_extended_protocol_implicit_transaction()
    }

    /// Dumps the catalog to a JSON string.
    ///
    /// No authorization is performed, so access to this function must be limited to internal
    /// servers or superusers.
    pub async fn dump_catalog(&mut self) -> Result<CatalogDump, AdapterError> {
        let catalog = self.catalog_snapshot("dump_catalog").await;
        catalog.dump().map_err(AdapterError::from)
    }

    /// Checks the catalog for internal consistency, returning a JSON object describing the
    /// inconsistencies, if there are any.
    ///
    /// No authorization is performed, so access to this function must be limited to internal
    /// servers or superusers.
    pub async fn check_catalog(&mut self) -> Result<(), serde_json::Value> {
        let catalog = self.catalog_snapshot("check_catalog").await;
        catalog.check_consistency()
    }

    /// Checks the coordinator for internal consistency, returning a JSON object describing the
    /// inconsistencies, if there are any. This is a superset of checks that check_catalog performs,
    ///
    /// No authorization is performed, so access to this function must be limited to internal
    /// servers or superusers.
    pub async fn check_coordinator(&self) -> Result<(), serde_json::Value> {
        self.send_without_session(|tx| Command::CheckConsistency { tx })
            .await
            .map_err(|inconsistencies| {
                serde_json::to_value(inconsistencies).unwrap_or_else(|_| {
                    serde_json::Value::String("failed to serialize inconsistencies".to_string())
                })
            })
    }

    pub async fn dump_coordinator_state(&self) -> Result<serde_json::Value, anyhow::Error> {
        self.send_without_session(|tx| Command::Dump { tx }).await
    }

    /// Tells the coordinator a statement has finished execution, in the cases
    /// where we have no other reason to communicate with the coordinator.
    pub fn retire_execute(
        &self,
        guard: ExecuteContextGuard,
        reason: StatementEndedExecutionReason,
    ) {
        if !guard.is_trivial() {
            let data = guard.defuse();
            let cmd = Command::RetireExecute { data, reason };
            self.inner().send(cmd);
        }
    }

    /// Sets up a streaming COPY FROM STDIN operation.
    ///
    /// Sends a command to the coordinator to create a background batch
    /// builder task. Returns a [`CopyFromStdinWriter`] that pgwire uses
    /// to stream decoded rows.
    pub async fn start_copy_from_stdin(
        &mut self,
        target_id: CatalogItemId,
        target_name: String,
        columns: Vec<ColumnIndex>,
        row_desc: mz_repr::RelationDesc,
        params: mz_pgcopy::CopyFormatParams<'static>,
    ) -> Result<CopyFromStdinWriter, AdapterError> {
        self.send(|tx, session| Command::StartCopyFromStdin {
            target_id,
            target_name,
            columns,
            row_desc,
            params,
            session,
            tx,
        })
        .await
    }

    /// Commits staged COPY FROM STDIN batches to a table.
    ///
    /// Adds the pre-built persist batches to the session's transaction
    /// operations. The actual commit happens when the transaction ends.
    pub fn stage_copy_from_stdin_batches(
        &mut self,
        target_id: CatalogItemId,
        batches: Vec<mz_persist_client::batch::ProtoBatch>,
    ) -> Result<(), AdapterError> {
        use crate::session::{TransactionOps, WriteOp};
        use mz_storage_client::client::TableData;

        self.session()
            .add_transaction_ops(TransactionOps::Writes(vec![WriteOp {
                id: target_id,
                rows: TableData::Batches(batches.into()),
            }]))?;
        Ok(())
    }

    /// Gets the current value of all system variables.
    pub async fn get_system_vars(&self) -> SystemVars {
        self.inner().get_system_vars().await
    }

    /// Updates the specified system variables to the specified values.
    pub async fn set_system_vars(
        &mut self,
        vars: BTreeMap<String, String>,
    ) -> Result<(), AdapterError> {
        let conn_id = self.session().conn_id().clone();
        self.send_without_session(|tx| Command::SetSystemVars { vars, conn_id, tx })
            .await
    }

    /// Injects audit events into the catalog via the coordinator.
    ///
    /// No authorization is performed, so access to this function must be limited to internal
    /// servers or superusers.
    pub async fn inject_audit_events(
        &mut self,
        events: Vec<crate::catalog::InjectedAuditEvent>,
    ) -> Result<(), AdapterError> {
        let conn_id = self.session().conn_id().clone();
        self.send_without_session(|tx| Command::InjectAuditEvents {
            events,
            conn_id,
            tx,
        })
        .await
    }

    /// Terminates the client session.
    pub async fn terminate(&mut self) {
        let conn_id = self.session().conn_id().clone();
        let res = self
            .send_without_session(|tx| Command::Terminate {
                conn_id,
                tx: Some(tx),
            })
            .await;
        if let Err(e) = res {
            // Nothing we can do to handle a failed terminate so we just log and ignore it.
            error!("Unable to terminate session: {e:?}");
        }
        // Prevent any communication with Coordinator after session is terminated.
        self.inner = None;
    }

    /// Returns a mutable reference to the session bound to this client.
    pub fn session(&mut self) -> &mut Session {
        self.session.as_mut().expect("session invariant violated")
    }

    /// Returns a reference to the inner client.
    pub fn inner(&self) -> &Client {
        self.inner.as_ref().expect("inner invariant violated")
    }

    async fn send_without_session<T, F>(&self, f: F) -> T
    where
        F: FnOnce(oneshot::Sender<T>) -> Command,
    {
        let (tx, rx) = oneshot::channel();
        self.inner().send(f(tx));
        rx.await.expect("sender dropped")
    }

    #[instrument(level = "debug")]
    async fn send<T, F>(&mut self, f: F) -> Result<T, AdapterError>
    where
        F: FnOnce(oneshot::Sender<Response<T>>, Session) -> Command,
    {
        self.send_with_cancel(f, futures::future::pending()).await
    }

    /// Send a [`Command`] to the Coordinator, with the ability to cancel the command.
    ///
    /// Note: the provided `cancel_future` must be cancel-safe as it's polled in a `select!` loop.
    #[instrument(level = "debug")]
    async fn send_with_cancel<T, F>(
        &mut self,
        f: F,
        cancel_future: impl Future<Output = ()> + Send,
    ) -> Result<T, AdapterError>
    where
        F: FnOnce(oneshot::Sender<Response<T>>, Session) -> Command,
    {
        let session = self.session.take().expect("session invariant violated");
        let mut typ = None;
        let application_name = session.application_name();
        let name_hint = ApplicationNameHint::from_str(application_name);
        let conn_id = session.conn_id().clone();
        let (tx, rx) = oneshot::channel();

        // Destructure self so we can hold a mutable reference to the inner client and session at
        // the same time.
        let Self {
            inner: inner_client,
            session: client_session,
            ..
        } = self;

        // TODO(parkmycar): Leaking this invariant here doesn't feel great, but calling
        // `self.client()` doesn't work because then Rust takes a borrow on the entirity of self.
        let inner_client = inner_client.as_ref().expect("inner invariant violated");

        // ~~SPOOKY ZONE~~
        //
        // This guard prevents a race where a `Session` is returned on `rx` but never placed
        // back in `self` because the Future returned by this function is concurrently dropped
        // with the Coordinator sending a response.
        let mut guarded_rx = rx.with_guard(|response: Response<_>| {
            *client_session = Some(response.session);
        });

        inner_client.send({
            let cmd = f(tx, session);
            // Measure the success and error rate of certain commands:
            // - declare reports success of SQL statement planning
            // - execute reports success of dataflow execution
            match cmd {
                Command::Execute { .. } => typ = Some("execute"),
                Command::GetWebhook { .. } => typ = Some("webhook"),
                Command::StartCopyFromStdin { .. }
                | Command::Startup { .. }
                | Command::AuthenticatePassword { .. }
                | Command::AuthenticateGetSASLChallenge { .. }
                | Command::AuthenticateVerifySASLProof { .. }
                | Command::CheckRoleCanLogin { .. }
                | Command::CatalogSnapshot { .. }
                | Command::Commit { .. }
                | Command::CancelRequest { .. }
                | Command::PrivilegedCancelRequest { .. }
                | Command::GetSystemVars { .. }
                | Command::SetSystemVars { .. }
                | Command::UpdateScopedSystemParameters { .. }
                | Command::InstallScopedSystemParameterFrontend { .. }
                | Command::Terminate { .. }
                | 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 { .. }
                | Command::RegisterConnectionCancelWatch { .. }
                | Command::CreateInternalSubscribe { .. }
                | Command::AttemptWrite { .. }
                | Command::DropInternalSubscribe { .. } => {}
            };
            cmd
        });

        let mut cancel_future = pin::pin!(cancel_future);
        let mut cancelled = false;
        loop {
            tokio::select! {
                res = &mut guarded_rx => {
                    // We received a result, so drop our guard to drop our borrows.
                    drop(guarded_rx);

                    let res = res.expect("sender dropped");
                    let status = res.result.is_ok().then_some("success").unwrap_or("error");
                    if let Err(err) = res.result.as_ref() {
                        if name_hint.should_trace_errors() {
                            tracing::warn!(?err, ?name_hint, "adapter response error");
                        }
                    }

                    if let Some(typ) = typ {
                        inner_client
                            .metrics
                            .commands
                            .with_label_values(&[typ, status, name_hint.as_str()])
                            .inc();
                    }
                    *client_session = Some(res.session);
                    return res.result;
                },
                _ = &mut cancel_future, if !cancelled => {
                    cancelled = true;
                    inner_client.send(Command::PrivilegedCancelRequest {
                        conn_id: conn_id.clone(),
                    });
                }
            };
        }
    }

    pub fn add_idle_in_transaction_session_timeout(&mut self) {
        let session = self.session();
        let timeout_dur = session.vars().idle_in_transaction_session_timeout();
        if !timeout_dur.is_zero() {
            let timeout_dur = timeout_dur.clone();
            if let Some(txn) = session.transaction().inner() {
                let txn_id = txn.id.clone();
                let timeout = TimeoutType::IdleInTransactionSession(txn_id);
                self.timeouts.add_timeout(timeout, timeout_dur);
            }
        }
    }

    pub fn remove_idle_in_transaction_session_timeout(&mut self) {
        let session = self.session();
        if let Some(txn) = session.transaction().inner() {
            let txn_id = txn.id.clone();
            self.timeouts
                .remove_timeout(&TimeoutType::IdleInTransactionSession(txn_id));
        }
    }

    /// # Cancel safety
    ///
    /// This method is cancel safe. If `recv` is used as the event in a
    /// `tokio::select!` statement and some other branch
    /// completes first, it is guaranteed that no messages were received on this
    /// channel.
    pub async fn recv_timeout(&mut self) -> Option<TimeoutType> {
        self.timeouts.recv().await
    }

    /// Attempt to sequence a peek from the session task.
    ///
    /// Returns `Ok(Some(response))` if we handled the peek, or `Ok(None)` if the statement is
    /// not one that the frontend peek sequencing handles, in which case the Coordinator sequences
    /// it. If it returns an error, it should be returned to the user.
    pub(crate) async fn try_frontend_peek(
        &mut self,
        portal_name: &str,
        logging: &mut ExecutionLogging,
    ) -> Result<Option<ExecuteResponse>, AdapterError> {
        let session = self.session.as_mut().expect("SessionClient invariant");
        self.peek_client
            .try_frontend_peek(portal_name, session, logging)
            .await
    }

    /// Whether the frontend read-then-write path could take this portal over.
    ///
    /// The gate is deliberately cheap, a portal lookup, because
    /// every statement that reaches `execute_attempts` without being handled by
    /// the peek path is tested against it. Everything expensive, including the
    /// coordinator round-trip that registers the connection cancel watch, sits
    /// behind it.
    fn frontend_read_then_write_applies(&self, portal_name: &str) -> bool {
        let session = self.session.as_ref().expect("SessionClient invariant");
        match session.get_portal_unverified(portal_name) {
            Some(portal) => portal
                .stmt
                .as_deref()
                .is_some_and(is_read_then_write_statement),
            None => false,
        }
    }

    /// Runs frontend read-then-write while reacting to both local/session
    /// cancellation and coordinator-issued connection cancellation.
    ///
    /// Returns `Ok(None)` when the statement is not eligible for this path and
    /// the caller must fall back to the coordinator, either because
    /// `try_frontend_read_then_write` declined it or because this wrapper did.
    ///
    /// Cancellation and statement timeout are never reported for a write that
    /// may have committed. Once a write has been submitted we await its
    /// definitive result instead of returning the cancellation.
    async fn try_frontend_read_then_write_with_cancel(
        &mut self,
        portal_name: &str,
        logging: &mut ExecutionLogging,
        cancel_future: impl Future<Output = ()> + Send,
    ) -> Result<Option<ExecuteResponse>, AdapterError> {
        // Bail out before the cancel-watch registration below, which is a
        // synchronous round-trip through the coordinator's command loop. A
        // statement this path will not take over must not pay for it, and must
        // not add queueing latency for other sessions either.
        if !self.frontend_read_then_write_applies(portal_name) {
            return Ok(None);
        }

        let conn_id = self.session().conn_id().clone();
        let statement_timeout = *self.session().vars().statement_timeout();
        let inner_client = self.inner().clone();
        let attempt_state = Arc::new(FrontendWriteAttemptState::new());

        let mut cancel_future = pin::pin!(cancel_future);
        let statement_timeout = async move {
            if statement_timeout.is_zero() {
                futures::future::pending::<()>().await;
            } else {
                tokio::time::sleep(statement_timeout).await;
            }
        };
        tokio::pin!(statement_timeout);

        // Registering installs a fresh channel, so this cannot observe a
        // cancellation aimed at an earlier statement. The entry it leaves behind
        // is replaced by the next registration and removed when a statement
        // reaches the coordinator or the connection's state is cleared, so there
        // is nothing to unregister here.
        let mut connection_cancel_rx = {
            let register =
                self.peek_client
                    .call_coordinator(|tx| Command::RegisterConnectionCancelWatch {
                        conn_id: conn_id.clone(),
                        tx,
                    });
            tokio::pin!(register);
            tokio::select! {
                rx = &mut register => rx?,
                _ = &mut cancel_future => {
                    inner_client.try_send(Command::PrivilegedCancelRequest {
                        conn_id: conn_id.clone(),
                    });
                    return Err(AdapterError::Canceled);
                }
                _ = &mut statement_timeout => {
                    inner_client.try_send(Command::PrivilegedCancelRequest {
                        conn_id: conn_id.clone(),
                    });
                    return Err(AdapterError::StatementTimeout);
                }
            }
        };
        if *connection_cancel_rx.borrow() {
            return Err(AdapterError::Canceled);
        }
        let connection_cancel = async move {
            if connection_cancel_rx.wait_for(|v| *v).await.is_err() {
                futures::future::pending::<()>().await;
            }
        };
        tokio::pin!(connection_cancel);

        let frontend_read_then_write =
            self.try_frontend_read_then_write(portal_name, logging, Arc::clone(&attempt_state));
        tokio::pin!(frontend_read_then_write);

        let requested = tokio::select! {
            response = &mut frontend_read_then_write => return response,
            _ = &mut cancel_future => FrontendWriteCancellation::Canceled,
            _ = &mut connection_cancel => FrontendWriteCancellation::Canceled,
            _ = &mut statement_timeout => FrontendWriteCancellation::StatementTimeout,
        };

        attempt_state.request(requested);
        inner_client.try_send(Command::PrivilegedCancelRequest {
            conn_id: conn_id.clone(),
        });

        if !attempt_state.write_submitted() {
            return Err(requested.into());
        }

        // A submitted write can already be durable. Await its definitive result
        // rather than reporting cancellation or timeout incorrectly.
        frontend_read_then_write.await
    }

    /// Attempt to sequence a read-then-write (DELETE/UPDATE/INSERT INTO ..
    /// SELECT .. FROM) from the session task.
    ///
    /// Returns `Ok(Some(response))` if we handled the operation, or `Ok(None)`
    /// to fall back to the Coordinator's sequencing. If it returns an error, it
    /// should be returned to the user.
    async fn try_frontend_read_then_write(
        &mut self,
        portal_name: &str,
        logging: &mut ExecutionLogging,
        attempt_state: Arc<FrontendWriteAttemptState>,
    ) -> Result<Option<ExecuteResponse>, AdapterError> {
        let catalog = self.catalog_snapshot("try_frontend_read_then_write").await;

        let stmt = {
            let session = self.session.as_ref().expect("SessionClient invariant");
            let portal = match session.get_portal_unverified(portal_name) {
                Some(portal) => portal,
                None => return Ok(None), // Portal doesn't exist, fall back
            };
            portal.stmt.clone()
        };

        let stmt = match stmt {
            Some(stmt) if is_read_then_write_statement(&stmt) => stmt,
            Some(_stmt) => {
                return Ok(None);
            }
            None => {
                return Ok(None);
            }
        };

        // Verify and plan against one catalog snapshot. Pairing a stale plan
        // with a newer target generation could direct a write incorrectly.
        // A failed verification is not logged, mirroring the coordinator: the
        // portal is what statement logging draws its record from.
        Coordinator::verify_portal(
            &catalog,
            self.session.as_mut().expect("SessionClient invariant"),
            portal_name,
        )?;

        let (params, logging_info, lifecycle_timestamps) = {
            let portal = self
                .session
                .as_ref()
                .expect("SessionClient invariant")
                .get_portal_unverified(portal_name)
                .expect("verified above");
            (
                portal.parameters.clone(),
                Arc::clone(&portal.logging),
                portal.lifecycle_timestamps.clone(),
            )
        };

        // Past this point the coordinator never sees this statement, so every
        // exit has to produce an outcome for it. The remaining `Ok(None)`
        // bailouts are all above.
        logging.take_over(
            &self.peek_client,
            self.session.as_mut().expect("SessionClient invariant"),
            Some(&stmt),
            &params,
            &logging_info,
            &catalog,
            lifecycle_timestamps,
            TakeOver::StatementToRun,
        );

        // Mirror the coordinator's transaction-state gate in `handle_execute`:
        // in a multi-statement transaction (an implicit batch or an explicit
        // block), the only DML allowed is an AST-constant INSERT without
        // RETURNING, which joins the transaction's write ops and commits at
        // transaction end. All other DML is prohibited because writes on this
        // path commit immediately and cannot be rolled back at transaction
        // end. `Failed` transactions pass through, pgwire only admits
        // COMMIT/ROLLBACK in that state.
        //
        // An AST-constant source can still plan to a read, so the check on the
        // planned selection further down narrows this.
        {
            let session = self.session.as_ref().expect("SessionClient invariant");
            // `Started` does not mean "single statement" on its own. An
            // extended-protocol pipeline keeps the transaction `Started` across
            // statements until `Sync`, so once it holds write ops this
            // statement runs alongside them and belongs with the
            // multi-statement cases. Committing here would commit against a
            // snapshot that lacks those ops, reordering this statement before
            // writes that a later pipeline error would roll back.
            let contains_ops = session.transaction().contains_ops();
            match session.transaction() {
                TransactionStatus::Default | TransactionStatus::Failed(_) => {}
                TransactionStatus::Started(_) if !contains_ops => {}
                TransactionStatus::Started(_)
                | TransactionStatus::InTransactionImplicit(_)
                | TransactionStatus::InTransaction(_) => {
                    let constant_insert = matches!(
                        &*stmt,
                        Statement::Insert(InsertStatement {
                            source, returning, ..
                        }) if returning.is_empty() && ConstantVisitor::insert_source(source)
                    );
                    if !constant_insert {
                        return Err(prohibited_in_transaction(&stmt));
                    }
                }
            }
        }

        let (plan, target_cluster, resolved_ids, sql_impl_ids) = {
            let session = self.session.as_mut().expect("SessionClient invariant");
            let conn_catalog = catalog.for_session(session);
            let (stmt, resolved_ids) = mz_sql::names::resolve(&conn_catalog, (*stmt).clone())?;
            let pcx = session.pcx();
            let (plan, sql_impl_ids) =
                mz_sql::plan::plan(Some(pcx), &conn_catalog, stmt, &params, &resolved_ids)?;

            let target_cluster = match session.transaction().cluster() {
                Some(cluster_id) => crate::coord::TargetCluster::Transaction(cluster_id),
                None => crate::coord::catalog_serving::auto_run_on_catalog_server(
                    &conn_catalog,
                    session,
                    &plan,
                ),
            };

            (plan, target_cluster, resolved_ids, sql_impl_ids)
        };

        // Reject mutations in read-only mode (e.g. during 0dt upgrades). Placed
        // where the coordinator has it, in `sequence_plan`: after planning, so a
        // statement that does not plan reports the planning error, and before
        // the cluster and RBAC checks below, which the coordinator also reports
        // second. Every sub-path from here on writes (constant INSERT and the
        // OCC INSERT/UPDATE/DELETE), so one check covers them all.
        if self.peek_client.read_only {
            return Err(AdapterError::ReadOnly);
        }

        // Cluster restrictions and RBAC, mirroring the coordinator's checks
        // in sequencer.rs. Resolution may fail if the target cluster doesn't
        // exist. That gets reported later (with the correct error) by
        // `validate_read_then_write`. For the purposes of these checks we
        // treat it as "no cluster known", consistent with the coordinator.
        let (target_cluster_id, target_cluster_name) = {
            let session = self.session.as_ref().expect("SessionClient invariant");
            match catalog.resolve_target_cluster(target_cluster.clone(), session) {
                Ok(cluster) => (Some(cluster.id), Some(cluster.name.clone())),
                Err(_) => (None, None),
            }
        };

        // Record the cluster before the checks below can fail, so that their
        // error rows carry it, as the coordinator's do.
        if let (Some(logging_id), Some(cluster_id), Some(cluster_name)) =
            (logging.id(), target_cluster_id, target_cluster_name.clone())
        {
            self.peek_client
                .log_set_cluster(logging_id, cluster_id, cluster_name);
        }

        {
            let session = self.session.as_ref().expect("SessionClient invariant");
            let conn_catalog = catalog.for_session(session);
            if let Some(cluster_name) = &target_cluster_name {
                crate::coord::catalog_serving::check_cluster_restrictions(
                    cluster_name,
                    &conn_catalog,
                    &plan,
                )?;
            }
            if let Err(e) = mz_sql::rbac::check_plan(
                &conn_catalog,
                None,
                session,
                &plan,
                target_cluster_id,
                &resolved_ids,
                &sql_impl_ids,
            ) {
                return Err(e.into());
            }
        }

        // Wait for any in-flight startup builtin-table appends that this plan
        // depends on. Mirrors the frontend_peek and coordinator sequencer
        // paths, and is a no-op for plans that don't depend on builtin tables.
        {
            let session = self.session.as_mut().expect("SessionClient invariant");
            if let Some((_, wait_future)) =
                crate::coord::appends::waiting_on_startup_appends(&catalog, session, &plan)
            {
                wait_future.await;
            }
        }

        // `allows_writes` is only defined inside a transaction, which is also
        // the only place it can be false: outside one the session task opens a
        // fresh transaction with no ops. Autocommit statements therefore rely on
        // the check in `PeekClient::frontend_read_then_write` instead.
        {
            let session = self.session.as_ref().expect("SessionClient invariant");
            if session.transaction().is_in_multi_statement_transaction()
                && !session.transaction().allows_writes()
            {
                return Err(AdapterError::ReadOnlyTransaction);
            }
            if session
                .vars()
                .transaction_isolation()
                .is_bounded_staleness()
            {
                return Err(AdapterError::BoundedStalenessReadOnly);
            }
        }

        // Handle ReadThenWrite plans or Insert plans.
        let rtw_plan = match plan {
            Plan::ReadThenWrite(rtw_plan) => rtw_plan,
            Plan::Insert(insert_plan) => {
                // A constant INSERT without RETURNING is a blind write, handled
                // here through the coordinator's `insert_constant` helper, which
                // buffers the rows as session write ops.
                //
                // Deciding that needs HIR lowered to MIR, because a VALUES list
                // is planned as a `Wrap` call at the HIR level.
                //
                // Only take that path when the HIR names no persisted
                // collections (no `Get` nodes on tables or MVs). The MIR
                // optimizer can fold an MV reference into a literal when the
                // MV's plan happens to be constant, but "plan is constant" is
                // NOT the same as "content is visible at the current
                // oracle_ts". A `REFRESH AT year 30000` MV has a constant plan
                // but no durable content until the refresh fires. Folding it
                // and blind-writing the literal would skip timestamp selection
                // and linearization, producing data that was never observable.
                // Preserving the HIR-level `Get` nodes routes the INSERT through
                // the RTW path, where timestamp selection handles REFRESH and
                // other time-dependent reads correctly.
                let has_read_deps = !insert_plan.values.depends_on().is_empty();

                if !has_read_deps {
                    let optimized_mir = if insert_plan.values.as_const().is_some() {
                        // Already constant at HIR level - just lower without optimization
                        let expr = insert_plan
                            .values
                            .clone()
                            .lower(catalog.system_config(), None)?;
                        mz_expr::OptimizedMirRelationExpr(expr)
                    } else {
                        // Need to optimize to check if it becomes constant.
                        // Use one-shot expression prep so unmaterializable
                        // functions like current_user() are resolved before we
                        // decide whether this can use the blind-write path.
                        let optimizer_config =
                            optimize::OptimizerConfig::from(catalog.system_config());
                        let session = self.session.as_ref().expect("SessionClient invariant");
                        let prep = ExprPrepOneShot {
                            logical_time: EvalTime::NotAvailable,
                            session,
                            catalog_state: catalog.state(),
                        };
                        let mut optimizer =
                            optimize::view::Optimizer::new_with_prep(optimizer_config, None, prep);
                        match optimizer.optimize(insert_plan.values.clone()) {
                            Ok(expr) => expr,
                            Err(OptimizerError::UncallableFunction {
                                func: UnmaterializableFunc::MzNow,
                                ..
                            }) => {
                                // Preserve the established user-facing `mz_now()`
                                // error by falling back to the RTW validator.
                                let expr = insert_plan
                                    .values
                                    .clone()
                                    .lower(catalog.system_config(), None)?;
                                mz_expr::OptimizedMirRelationExpr(expr)
                            }
                            Err(e) => return Err(e.into()),
                        }
                    };

                    let inner_mir = optimized_mir.into_inner();
                    if inner_mir.as_const().is_some() && insert_plan.returning.is_empty() {
                        let session = self.session.as_mut().expect("SessionClient invariant");
                        let result = Coordinator::insert_constant(
                            &catalog,
                            session,
                            insert_plan.id,
                            inner_mir,
                        );

                        return Ok(Some(result?));
                    }
                }

                let desc_arity = match catalog.try_get_entry(&insert_plan.id) {
                    Some(table) => {
                        let desc = table
                            .relation_desc_latest()
                            .ok_or_else(|| AdapterError::Internal("table has no desc".into()))?;
                        desc.arity()
                    }
                    None => {
                        return Err(AdapterError::Catalog(mz_catalog::memory::error::Error {
                            kind: mz_catalog::memory::error::ErrorKind::Sql(
                                mz_sql::catalog::CatalogError::UnknownItem(
                                    insert_plan.id.to_string(),
                                ),
                            ),
                        }));
                    }
                };

                let finishing = RowSetFinishing {
                    order_by: vec![],
                    limit: None,
                    offset: 0,
                    project: (0..desc_arity).collect(),
                };

                ReadThenWritePlan {
                    id: insert_plan.id,
                    selection: insert_plan.values,
                    finishing,
                    assignments: BTreeMap::new(),
                    kind: MutationKind::Insert,
                    returning: insert_plan.returning,
                }
            }
            _ => {
                return Err(AdapterError::Internal(
                    "unexpected plan type for mutation".into(),
                ));
            }
        };

        // The syntactic predicate for "reads persisted state", see the module
        // docs on `frontend_read_then_write`. Inside a transaction, only a write
        // that reads nothing can run on this path.
        //
        // The AST gate above is not enough to establish this. It admits
        // INSERTs whose source is constant in the AST, and such a statement can
        // still plan to a selection with `Get` nodes, because SQL-implemented
        // builtins (`pg_get_viewdef`, `text` to `reg*` casts, ...) read system
        // relations. So decide on the planned selection, and do it before we
        // execute a dataflow for a statement we would then refuse.
        {
            let session = self.session.as_ref().expect("SessionClient invariant");
            let in_transaction = session
                .transaction()
                .may_share_transaction_with_other_statements();
            let depends_on = rtw_plan.selection.depends_on();
            if in_transaction && !depends_on.is_empty() {
                // Report the reasons that hold wherever the statement runs
                // before the one that holds only here. A statement carrying
                // `mz_now`, or reading a system table, never works anywhere.
                // Answering with the transaction state names the one condition
                // the caller could remove, which tells them to retry outside a
                // transaction and get the same refusal again.
                if contains_mz_now(&rtw_plan) {
                    return Err(AdapterError::Unsupported(
                        "calls to mz_now in write statements",
                    ));
                }
                validate_selection_dependencies(&catalog, &depends_on, DependencyPolicy::UserDml)?;
                return Err(prohibited_in_transaction(&stmt));
            }
        }

        let session = self.session.as_mut().expect("SessionClient invariant");
        self.peek_client
            .frontend_read_then_write(
                session,
                rtw_plan,
                target_cluster,
                &catalog,
                logging.id(),
                attempt_state,
            )
            .await
            .map(Some)
    }
}

/// Whether a statement is one the frontend read-then-write path sequences.
///
/// These are the statement kinds that plan to a `ReadThenWrite` or an `Insert`.
/// Note that not every one of them ends up on the OCC path: an `INSERT` whose
/// source folds to a constant is dispatched as a blind write instead.
fn is_read_then_write_statement(stmt: &Statement<Raw>) -> bool {
    matches!(
        stmt,
        Statement::Delete(_) | Statement::Update(_) | Statement::Insert(_)
    )
}

/// Builds the error for DML that cannot run in a transaction block, mirroring
/// the coordinator's redaction in `handle_execute`: statements that can carry
/// sensitive literals are redacted because the error message is persisted in
/// `mz_statement_execution_history`.
fn prohibited_in_transaction(stmt: &Statement<Raw>) -> AdapterError {
    let op = if StatementKind::from(stmt).is_sensitive() {
        stmt.to_ast_string_redacted()
    } else {
        stmt.to_string()
    };
    AdapterError::OperationProhibitsTransaction(op)
}

impl Drop for SessionClient {
    fn drop(&mut self) {
        // We may not have a session if this client was dropped while awaiting
        // a response. In this case, it is the coordinator's responsibility to
        // terminate the session.
        if let Some(session) = self.session.take() {
            // We may not have a connection to the Coordinator if the session was
            // prematurely terminated, for example due to a timeout.
            if let Some(inner) = &self.inner {
                inner.send(Command::Terminate {
                    conn_id: session.conn_id().clone(),
                    tx: None,
                })
            }
        }
    }
}

/// Renders SQL for statement arrival logging: parsed and displayed with its
/// literals redacted, which is the same redaction the statement log applies,
/// and capped by [`truncate_sql_for_logging`].
/// When the text does not parse or exceeds the statement batch size limit, a
/// placeholder with the byte length is returned. Raw text is never returned,
/// so a statement that crashes the parser is not captured, an accepted
/// limitation.
pub fn redact_sql_for_logging(sql: &str) -> String {
    match mz_sql_parser::parser::parse_statements_with_limit(sql) {
        Ok(Ok(stmts)) => truncate_sql_for_logging(
            stmts
                .into_iter()
                .map(|stmt| stmt.ast.to_ast_string_redacted())
                .join("; "),
        ),
        Ok(Err(_)) => format!("<unparseable ({} bytes)>", sql.len()),
        Err(_) => format!("<too large ({} bytes)>", sql.len()),
    }
}

/// The maximum length in bytes of SQL text included in a log line.
///
/// Loki rejects a log entry larger than 256 KiB and drops it whole. A span's
/// fields are repeated in every event logged inside the span, twice when it is
/// the innermost span (the JSON format writes both `span` and `spans`), so one
/// statement can appear several times in an entry.
const MAX_LOGGED_SQL_LEN: usize = 8 * 1024;

/// Truncates SQL text to `MAX_LOGGED_SQL_LEN` bytes for a log line, noting
/// the original length.
pub fn truncate_sql_for_logging(sql: String) -> String {
    if sql.len() <= MAX_LOGGED_SQL_LEN {
        return sql;
    }
    let end = sql.floor_char_boundary(MAX_LOGGED_SQL_LEN);
    format!("{}... <truncated from {} bytes>", &sql[..end], sql.len())
}

#[derive(Hash, PartialEq, Eq, PartialOrd, Ord, Clone, Debug)]
pub enum TimeoutType {
    IdleInTransactionSession(TransactionId),
}

impl Display for TimeoutType {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        match self {
            TimeoutType::IdleInTransactionSession(txn_id) => {
                writeln!(f, "Idle in transaction session for transaction '{txn_id}'")
            }
        }
    }
}

impl From<TimeoutType> for AdapterError {
    fn from(timeout: TimeoutType) -> Self {
        match timeout {
            TimeoutType::IdleInTransactionSession(_) => {
                AdapterError::IdleInTransactionSessionTimeout
            }
        }
    }
}

struct Timeout {
    tx: mpsc::UnboundedSender<TimeoutType>,
    rx: mpsc::UnboundedReceiver<TimeoutType>,
    active_timeouts: BTreeMap<TimeoutType, AbortOnDropHandle<()>>,
}

impl Timeout {
    fn new() -> Self {
        let (tx, rx) = mpsc::unbounded_channel();
        Timeout {
            tx,
            rx,
            active_timeouts: BTreeMap::new(),
        }
    }

    /// # Cancel safety
    ///
    /// This method is cancel safe. If `recv` is used as the event in a
    /// `tokio::select!` statement and some other branch
    /// completes first, it is guaranteed that no messages were received on this
    /// channel.
    ///
    /// <https://docs.rs/tokio/latest/tokio/sync/mpsc/struct.UnboundedReceiver.html#cancel-safety>
    async fn recv(&mut self) -> Option<TimeoutType> {
        self.rx.recv().await
    }

    fn add_timeout(&mut self, timeout: TimeoutType, duration: Duration) {
        let tx = self.tx.clone();
        let timeout_key = timeout.clone();
        let handle = mz_ore::task::spawn(|| format!("{timeout_key}"), async move {
            tokio::time::sleep(duration).await;
            let _ = tx.send(timeout);
        })
        .abort_on_drop();
        self.active_timeouts.insert(timeout_key, handle);
    }

    fn remove_timeout(&mut self, timeout: &TimeoutType) {
        self.active_timeouts.remove(timeout);

        // Remove the timeout from the rx queue if it exists.
        let mut timeouts = Vec::new();
        while let Ok(pending_timeout) = self.rx.try_recv() {
            if timeout != &pending_timeout {
                timeouts.push(pending_timeout);
            }
        }
        for pending_timeout in timeouts {
            self.tx.send(pending_timeout).expect("rx is in this struct");
        }
    }
}

/// A wrapper around a Stream of PeekResponseUnary that records when it sees the
/// first row data in the given histogram. It also keeps track of whether we have already observed
/// the end of the underlying stream.
#[derive(Derivative)]
#[derivative(Debug)]
pub struct RecordFirstRowStream {
    /// The underlying stream of rows.
    #[derivative(Debug = "ignore")]
    pub rows: Box<dyn Stream<Item = PeekResponseUnary> + Unpin + Send + Sync>,
    /// The Instant when execution started.
    pub execute_started: Instant,
    /// The histogram where the time since `execute_started` will be recorded when we see the first
    /// row.
    pub time_to_first_row_seconds: Histogram,
    /// Whether we've seen any rows.
    pub saw_rows: bool,
    /// The Instant when we saw the first row.
    pub recorded_first_row_instant: Option<Instant>,
    /// Whether we have already observed the end of the underlying stream.
    pub no_more_rows: bool,
    /// Whether the first-to-last-byte metric has already been recorded for this stream.
    pub metric_recorded: bool,
}

impl RecordFirstRowStream {
    /// Create a new [`RecordFirstRowStream`]
    pub fn new(
        rows: Box<dyn Stream<Item = PeekResponseUnary> + Unpin + Send + Sync>,
        execute_started: Instant,
        client: &SessionClient,
        instance_id: Option<ComputeInstanceId>,
        strategy: Option<StatementExecutionStrategy>,
    ) -> Self {
        let histogram = Self::histogram(client, instance_id, strategy);
        Self {
            rows,
            execute_started,
            time_to_first_row_seconds: histogram,
            saw_rows: false,
            recorded_first_row_instant: None,
            no_more_rows: false,
            metric_recorded: false,
        }
    }

    fn histogram(
        client: &SessionClient,
        instance_id: Option<ComputeInstanceId>,
        strategy: Option<StatementExecutionStrategy>,
    ) -> Histogram {
        let session = client.session.as_ref().expect("session invariant");
        let isolation_level = *session.vars().transaction_isolation();
        let name_hint = ApplicationNameHint::from_str(session.application_name());

        client
            .inner()
            .metrics()
            .by_cluster
            .time_to_first_row_seconds(instance_id, isolation_level, strategy, name_hint)
    }

    /// If you want to match [`RecordFirstRowStream`]'s logic but don't need
    /// a UnboundedReceiver, you can tell it when to record an observation.
    pub fn record(
        execute_started: Instant,
        client: &SessionClient,
        instance_id: Option<ComputeInstanceId>,
        strategy: Option<StatementExecutionStrategy>,
    ) {
        Self::histogram(client, instance_id, strategy)
            .observe(execute_started.elapsed().as_secs_f64());
    }

    pub async fn recv(&mut self) -> Option<PeekResponseUnary> {
        let msg = self.rows.next().await;
        if !self.saw_rows && matches!(msg, Some(PeekResponseUnary::Rows(_))) {
            self.saw_rows = true;
            self.time_to_first_row_seconds
                .observe(self.execute_started.elapsed().as_secs_f64());
            self.recorded_first_row_instant = Some(Instant::now());
        }
        if msg.is_none() {
            self.no_more_rows = true;
        }
        msg
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[mz_ore::test]
    fn truncate_sql_for_logging_caps_length() {
        let short = "SELECT 1".to_string();
        assert_eq!(truncate_sql_for_logging(short.clone()), short);

        let exact = "x".repeat(MAX_LOGGED_SQL_LEN);
        assert_eq!(truncate_sql_for_logging(exact.clone()), exact);

        // A multi-byte character straddles the cap, so the cut must move back
        // to the previous character boundary.
        let long = format!(
            "{}é{}",
            "x".repeat(MAX_LOGGED_SQL_LEN - 1),
            "y".repeat(1000)
        );
        let truncated = truncate_sql_for_logging(long.clone());
        assert_eq!(
            truncated,
            format!(
                "{}... <truncated from {} bytes>",
                "x".repeat(MAX_LOGGED_SQL_LEN - 1),
                long.len()
            )
        );
    }
}