drizzle 0.2.0

//! AWS Aurora Serverless Data API transaction.
//!
//! The Data API has first-class server-side transactions:
//!
//! * [`Drizzle::transaction`] issues a `BeginTransaction` call; the returned
//!   `transactionId` is threaded into every subsequent `ExecuteStatement` via
//!   the `transactionId` field.
//! * Commit / rollback go through `CommitTransaction` / `RollbackTransaction`
//!   (not raw SQL).
//! * Savepoints use regular `SAVEPOINT` / `RELEASE SAVEPOINT` /
//!   `ROLLBACK TO SAVEPOINT` SQL that runs inside the transaction context.

use std::marker::PhantomData;
use std::sync::{Arc, Mutex};

use crate::transaction::savepoint::{AsyncSavepointState, async_savepoint};

use aws_sdk_rdsdata::Client;
use drizzle_core::dialect::ParamStyle;
use drizzle_core::error::DrizzleError;
use drizzle_core::traits::ToSQL;
use drizzle_postgres::aws_data_api::Row;
use drizzle_postgres::builder::{
    self, DeleteInitial, InsertInitial, QueryBuilder, SelectInitial, UpdateInitial,
    delete::DeleteBuilder, insert::InsertBuilder, select::SelectBuilder, update::UpdateBuilder,
};
use drizzle_postgres::common::PostgresTransactionType;
use drizzle_postgres::traits::PostgresTable;
use drizzle_postgres::transaction::{IsolationLevel, TransactionConfig};
use drizzle_postgres::values::PostgresValue;

use crate::builder::postgres::aws_data_api::{
    Rows, aws_error, decode_rows, encode_params, execute_statement_raw,
};

/// Returns an error indicating the transaction has already been consumed.
fn tx_consumed_error() -> DrizzleError {
    DrizzleError::TransactionError("Transaction already consumed".into())
}

/// AWS Data API transaction builder wrapper. See
/// `TransactionBuilder` for the
/// typestate-advancing methods; executor methods live below.
pub type TransactionBuilder<'tx, Schema, Builder, State> =
    crate::transaction::postgres::typestate::TransactionBuilder<
        'tx,
        &'tx Transaction<Schema>,
        Schema,
        Builder,
        State,
    >;

/// Active AWS Aurora Data API transaction.
///
/// Owns the `transactionId` returned by `BeginTransaction` and threads it into
/// every `ExecuteStatement`. You do not commit or roll back by hand:
/// [`Drizzle::transaction`](crate::postgres::aws::Drizzle::transaction) ends
/// the transaction with `CommitTransaction` when the callback returns `Ok` and
/// with `RollbackTransaction` when it returns `Err`. If the transaction is
/// dropped while still open, for example because the future returned by
/// `transaction` was dropped before it finished, `Drop` spawns a best-effort
/// `RollbackTransaction` on the current Tokio runtime.
///
/// Cloning a `Client` is cheap (internal `Arc`), so a transaction can freely
/// reuse the ambient client.
pub struct Transaction<Schema = ()> {
    client: Client,
    resource_arn: Arc<str>,
    secret_arn: Arc<str>,
    database: Option<Arc<str>>,
    tx_id: Mutex<Option<String>>,
    config: TransactionConfig,
    savepoints: AsyncSavepointState,
    schema: Schema,
}

impl<Schema> Drop for Transaction<Schema> {
    fn drop(&mut self) {
        let Some(transaction_id) = self.tx_id.get_mut().ok().and_then(Option::take) else {
            return;
        };
        let Ok(runtime) = tokio::runtime::Handle::try_current() else {
            return;
        };

        let client = self.client.clone();
        let resource_arn = Arc::clone(&self.resource_arn);
        let secret_arn = Arc::clone(&self.secret_arn);
        let _rollback = runtime.spawn(async move {
            let _ = client
                .rollback_transaction()
                .resource_arn(resource_arn.as_ref())
                .secret_arn(secret_arn.as_ref())
                .transaction_id(transaction_id)
                .send()
                .await;
        });
    }
}

impl<Schema> std::fmt::Debug for Transaction<Schema> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let is_active = self.tx_id.lock().is_ok_and(|g| g.is_some());
        f.debug_struct("Transaction")
            .field("config", &self.config)
            .field("is_active", &is_active)
            .field("savepoints", &self.savepoints)
            .field("resource_arn", &self.resource_arn)
            .field("database", &self.database)
            .finish_non_exhaustive()
    }
}

impl<Schema> Transaction<Schema> {
    /// Construct a new transaction handle.
    pub(crate) fn new(
        client: Client,
        resource_arn: Arc<str>,
        secret_arn: Arc<str>,
        database: Option<Arc<str>>,
        transaction_id: String,
        config: TransactionConfig,
        schema: Schema,
    ) -> Self {
        Self {
            client,
            resource_arn,
            secret_arn,
            database,
            tx_id: Mutex::new(Some(transaction_id)),
            config,
            savepoints: AsyncSavepointState::new(),
            schema,
        }
    }

    /// Schema handle.
    #[inline]
    pub const fn schema(&self) -> &Schema {
        &self.schema
    }

    /// Legacy isolation view.
    ///
    /// This cannot distinguish server-default isolation from explicit
    /// `READ COMMITTED`. Use [`Self::config`] when that distinction matters.
    #[deprecated(since = "0.2.0", note = "use config()")]
    #[inline]
    pub const fn tx_type(&self) -> PostgresTransactionType {
        match self.config.isolation() {
            None | Some(IsolationLevel::ReadCommitted) => PostgresTransactionType::ReadCommitted,
            Some(IsolationLevel::ReadUncommitted) => PostgresTransactionType::ReadUncommitted,
            Some(IsolationLevel::RepeatableRead) => PostgresTransactionType::RepeatableRead,
            Some(IsolationLevel::Serializable) => PostgresTransactionType::Serializable,
        }
    }

    /// Configuration used to begin this transaction.
    #[inline]
    pub const fn config(&self) -> TransactionConfig {
        self.config
    }

    /// Current transaction id, if the transaction is still open.
    pub fn transaction_id(&self) -> Option<String> {
        self.tx_id.lock().ok().and_then(|g| g.clone())
    }

    /// Run a nested savepoint block.
    ///
    /// On `Ok`: `RELEASE SAVEPOINT`.
    /// On `Err`: `ROLLBACK TO SAVEPOINT` + `RELEASE SAVEPOINT`.
    /// The outer transaction stays live either way.
    ///
    /// # Errors
    ///
    /// Returns [`DrizzleError`] if the savepoint cannot be created/released, or the inner closure returns an error.
    pub async fn savepoint<F, R>(&self, f: F) -> drizzle_core::error::Result<R>
    where
        F: AsyncFnOnce(&Self) -> drizzle_core::error::Result<R>,
    {
        async_savepoint(
            &self.savepoints,
            |sql| async move { self.execute(sql.as_str()).await.map(|_| ()) },
            f(self),
        )
        .await
    }

    postgres_transaction_constructors!();

    // Inline execution methods.

    /// Run a raw SQL / built query and return affected row count.
    ///
    /// # Errors
    ///
    /// Returns [`DrizzleError`] if the Data API call fails or the SQL is invalid.
    pub async fn execute<'q, T>(&self, query: T) -> drizzle_core::error::Result<u64>
    where
        T: ToSQL<'q, PostgresValue<'q>>,
    {
        let sql = query.to_sql();
        let (sql_str, params) = {
            #[cfg(feature = "profiling")]
            drizzle_core::drizzle_profile_scope!("postgres.aws_data_api", "tx.execute");
            let (sql_str, params) = sql.build_with(ParamStyle::ColonNumbered);
            drizzle_core::drizzle_trace_query!(&sql_str, params.len());
            (sql_str, params)
        };

        let sql_params = encode_params(params.as_slice());
        let out = self.run_statement(&sql_str, sql_params).await?;
        Ok(out.number_of_records_updated.max(0).cast_unsigned())
    }

    /// Run a query and collect all rows into `C`.
    ///
    /// # Errors
    ///
    /// Returns [`DrizzleError`] if the Data API call fails or row decoding fails.
    pub async fn all<'q, T, R, C>(&self, query: T) -> drizzle_core::error::Result<C>
    where
        R: for<'r> TryFrom<&'r Row>,
        for<'r> <R as TryFrom<&'r Row>>::Error: Into<drizzle_core::error::DrizzleError>,
        T: ToSQL<'q, PostgresValue<'q>>,
        C: core::iter::FromIterator<R>,
    {
        let sql = query.to_sql();
        let (sql_str, params) = {
            #[cfg(feature = "profiling")]
            drizzle_core::drizzle_profile_scope!("postgres.aws_data_api", "tx.all");
            let (sql_str, params) = sql.build_with(ParamStyle::ColonNumbered);
            drizzle_core::drizzle_trace_query!(&sql_str, params.len());
            (sql_str, params)
        };

        let sql_params = encode_params(params.as_slice());
        let out = self.run_statement(&sql_str, sql_params).await?;
        let rows = decode_rows(out);
        rows.into_iter()
            .map(|row| R::try_from(&row).map_err(Into::into))
            .collect()
    }

    /// Run a query and return a lazy decoded-row cursor.
    ///
    /// # Errors
    ///
    /// Returns [`DrizzleError`] if the Data API call fails.
    pub async fn rows<'q, T, R>(&self, query: T) -> drizzle_core::error::Result<Rows<R>>
    where
        R: for<'r> TryFrom<&'r Row>,
        for<'r> <R as TryFrom<&'r Row>>::Error: Into<drizzle_core::error::DrizzleError>,
        T: ToSQL<'q, PostgresValue<'q>>,
    {
        let sql = query.to_sql();
        let (sql_str, params) = {
            #[cfg(feature = "profiling")]
            drizzle_core::drizzle_profile_scope!("postgres.aws_data_api", "tx.rows");
            let (sql_str, params) = sql.build_with(ParamStyle::ColonNumbered);
            drizzle_core::drizzle_trace_query!(&sql_str, params.len());
            (sql_str, params)
        };

        let sql_params = encode_params(params.as_slice());
        let out = self.run_statement(&sql_str, sql_params).await?;
        Ok(Rows::new(decode_rows(out)))
    }

    /// Run a query and return a single row (errors if empty).
    ///
    /// # Errors
    ///
    /// Returns [`DrizzleError`] if the Data API call fails, no rows match (returns `DrizzleError::NotFound`), or decoding fails.
    pub async fn get<'q, T, R>(&self, query: T) -> drizzle_core::error::Result<R>
    where
        R: for<'r> TryFrom<&'r Row>,
        for<'r> <R as TryFrom<&'r Row>>::Error: Into<drizzle_core::error::DrizzleError>,
        T: ToSQL<'q, PostgresValue<'q>>,
    {
        let sql = query.to_sql();
        let (sql_str, params) = {
            #[cfg(feature = "profiling")]
            drizzle_core::drizzle_profile_scope!("postgres.aws_data_api", "tx.get");
            let (sql_str, params) = sql.build_with(ParamStyle::ColonNumbered);
            drizzle_core::drizzle_trace_query!(&sql_str, params.len());
            (sql_str, params)
        };

        let sql_params = encode_params(params.as_slice());
        let out = self.run_statement(&sql_str, sql_params).await?;
        let row = decode_rows(out)
            .into_iter()
            .next()
            .ok_or(DrizzleError::NotFound)?;
        R::try_from(&row).map_err(Into::into)
    }

    /// Commit via the service-level `CommitTransaction` call.
    pub(crate) async fn commit(&self) -> drizzle_core::error::Result<()> {
        if let Err(error) = self.savepoints.ensure_usable() {
            self.rollback().await?;
            return Err(error);
        }
        if self.savepoints.aborted().is_aborted() {
            self.rollback().await?;
            return Err(crate::transaction::savepoint::aborted_transaction_error());
        }
        let tx_id = self
            .tx_id
            .lock()
            .map_err(|_| tx_consumed_error())?
            .clone()
            .ok_or_else(tx_consumed_error)?;
        // CommitTransaction doesn't take a database — transaction id is enough.
        self.client
            .commit_transaction()
            .resource_arn(self.resource_arn.as_ref())
            .secret_arn(self.secret_arn.as_ref())
            .transaction_id(tx_id)
            .send()
            .await
            .map_err(|e| aws_error("commit_transaction", &e))?;
        self.tx_id.lock().map_err(|_| tx_consumed_error())?.take();
        Ok(())
    }

    /// Roll back via the service-level `RollbackTransaction` call.
    pub(crate) async fn rollback(&self) -> drizzle_core::error::Result<()> {
        let tx_id = self
            .tx_id
            .lock()
            .map_err(|_| tx_consumed_error())?
            .clone()
            .ok_or_else(tx_consumed_error)?;
        // RollbackTransaction doesn't take a database — transaction id is enough.
        self.client
            .rollback_transaction()
            .resource_arn(self.resource_arn.as_ref())
            .secret_arn(self.secret_arn.as_ref())
            .transaction_id(tx_id)
            .send()
            .await
            .map_err(|e| aws_error("rollback_transaction", &e))?;
        self.tx_id.lock().map_err(|_| tx_consumed_error())?.take();
        Ok(())
    }

    /// Internal helper — runs a statement with this transaction's id threaded in.
    pub(crate) async fn run_statement(
        &self,
        sql: &str,
        params: Vec<aws_sdk_rdsdata::types::SqlParameter>,
    ) -> drizzle_core::error::Result<
        aws_sdk_rdsdata::operation::execute_statement::ExecuteStatementOutput,
    > {
        self.savepoints.ensure_usable()?;
        // Clone the id out of the `Mutex` so no guard is held across the `.await` below.
        // Holding a `MutexGuard` over an await would make this future `!Send` (on older
        // compilers) and risks lock contention stalls.
        let tx_id = self
            .tx_id
            .lock()
            .map_err(|_| tx_consumed_error())?
            .clone()
            .ok_or_else(tx_consumed_error)?;
        let result = execute_statement_raw(
            &self.client,
            &self.resource_arn,
            &self.secret_arn,
            self.database.as_deref(),
            sql,
            params,
            Some(&tx_id),
        )
        .await;
        // A failed statement aborts the PostgreSQL transaction behind the
        // Data API: every later statement fails and COMMIT rolls back.
        if result.is_err() {
            self.savepoints.aborted().mark();
        }
        result
    }
}

// =============================================================================
// TransactionBuilder trailing-impls (execute / all / get)
// =============================================================================

impl<'tx, 'q, Schema, State, Table, Mk, Rw, Grouped>
    TransactionBuilder<'tx, Schema, QueryBuilder<'q, Schema, State, Table, Mk, Rw, Grouped>, State>
where
    State: builder::ExecutableState,
{
    /// Run the builder and return affected row count.
    ///
    /// # Errors
    ///
    /// Returns [`DrizzleError`] if the Data API call fails or the SQL is invalid.
    pub async fn execute(self) -> drizzle_core::error::Result<u64> {
        let (sql_str, params) = {
            #[cfg(feature = "profiling")]
            drizzle_core::drizzle_profile_scope!("postgres.aws_data_api", "tx_builder.execute");
            let (sql_str, params) = self.builder.sql.build_with(ParamStyle::ColonNumbered);
            drizzle_core::drizzle_trace_query!(&sql_str, params.len());
            (sql_str, params)
        };

        let sql_params = encode_params(params.as_slice());
        let out = self.runner.run_statement(&sql_str, sql_params).await?;
        Ok(out.number_of_records_updated.max(0).cast_unsigned())
    }

    /// Run the builder and collect all rows using the builder's row type.
    ///
    /// # Errors
    ///
    /// Returns [`DrizzleError`] if the Data API call fails or row decoding fails.
    pub async fn all<R>(self) -> drizzle_core::error::Result<Vec<R>>
    where
        R: for<'r> TryFrom<&'r Row>,
        for<'r> <R as TryFrom<&'r Row>>::Error: Into<drizzle_core::error::DrizzleError>,
    {
        let (sql_str, params) = {
            #[cfg(feature = "profiling")]
            drizzle_core::drizzle_profile_scope!("postgres.aws_data_api", "tx_builder.all");
            let (sql_str, params) = self.builder.sql.build_with(ParamStyle::ColonNumbered);
            drizzle_core::drizzle_trace_query!(&sql_str, params.len());
            (sql_str, params)
        };

        let sql_params = encode_params(params.as_slice());
        let out = self.runner.run_statement(&sql_str, sql_params).await?;
        let rows = decode_rows(out);
        let mut decoded = Vec::with_capacity(rows.len());
        for row in &rows {
            decoded.push(R::try_from(row).map_err(Into::into)?);
        }
        Ok(decoded)
    }

    /// Run the builder and return a lazy decoded-row cursor.
    ///
    /// # Errors
    ///
    /// Returns [`DrizzleError`] if the Data API call fails.
    pub async fn rows(self) -> drizzle_core::error::Result<Rows<Rw>>
    where
        Rw: for<'r> TryFrom<&'r Row>,
        for<'r> <Rw as TryFrom<&'r Row>>::Error: Into<drizzle_core::error::DrizzleError>,
    {
        let (sql_str, params) = {
            #[cfg(feature = "profiling")]
            drizzle_core::drizzle_profile_scope!("postgres.aws_data_api", "tx_builder.rows");
            let (sql_str, params) = self.builder.sql.build_with(ParamStyle::ColonNumbered);
            drizzle_core::drizzle_trace_query!(&sql_str, params.len());
            (sql_str, params)
        };

        let sql_params = encode_params(params.as_slice());
        let out = self.runner.run_statement(&sql_str, sql_params).await?;
        Ok(Rows::new(decode_rows(out)))
    }

    /// Run the builder and return a single row.
    ///
    /// # Errors
    ///
    /// Returns [`DrizzleError`] if the Data API call fails, no rows match (returns `DrizzleError::NotFound`), or decoding fails.
    pub async fn get<R>(self) -> drizzle_core::error::Result<R>
    where
        R: for<'r> TryFrom<&'r Row>,
        for<'r> <R as TryFrom<&'r Row>>::Error: Into<drizzle_core::error::DrizzleError>,
    {
        let (sql_str, params) = {
            #[cfg(feature = "profiling")]
            drizzle_core::drizzle_profile_scope!("postgres.aws_data_api", "tx_builder.get");
            let (sql_str, params) = self.builder.sql.build_with(ParamStyle::ColonNumbered);
            drizzle_core::drizzle_trace_query!(&sql_str, params.len());
            (sql_str, params)
        };

        let sql_params = encode_params(params.as_slice());
        let out = self.runner.run_statement(&sql_str, sql_params).await?;
        let row = decode_rows(out)
            .into_iter()
            .next()
            .ok_or(DrizzleError::NotFound)?;
        R::try_from(&row).map_err(Into::into)
    }
}

// `ToSQL for TransactionBuilder` is now provided by the shared `DrizzleBuilder`
// impl in `crate::builder::postgres::common`.