drizzle 0.1.12

//! 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::atomic::AtomicU32;
use std::sync::{Arc, Mutex};

use crate::transaction::savepoint::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::values::PostgresValue;

use crate::builder::postgres::aws_data_api::{
    Rows, aws_error, decode_rows, encode_params, execute_statement_raw,
};
use crate::builder::postgres::rows::DecodeRows as _;

/// 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
/// [`crate::transaction::postgres::typestate::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` until `commit()` or `rollback()` consumes it.
/// 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>>,
    tx_type: PostgresTransactionType,
    savepoint_depth: AtomicU32,
    schema: Schema,
}

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("tx_type", &self.tx_type)
            .field("is_active", &is_active)
            .field("savepoint_depth", &self.savepoint_depth)
            .field("resource_arn", &self.resource_arn)
            .field("database", &self.database)
            .finish_non_exhaustive()
    }
}

impl<Schema> Transaction<Schema> {
    /// Construct a new transaction handle.
    pub(crate) const fn new(
        client: Client,
        resource_arn: Arc<str>,
        secret_arn: Arc<str>,
        database: Option<Arc<str>>,
        transaction_id: String,
        tx_type: PostgresTransactionType,
        schema: Schema,
    ) -> Self {
        Self {
            client,
            resource_arn,
            secret_arn,
            database,
            tx_id: Mutex::new(Some(transaction_id)),
            tx_type,
            savepoint_depth: AtomicU32::new(0),
            schema,
        }
    }

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

    /// Isolation / transaction type configured on begin.
    #[inline]
    pub const fn tx_type(&self) -> PostgresTransactionType {
        self.tx_type
    }

    /// 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.savepoint_depth,
            |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 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<()> {
        let tx_id = self
            .tx_id
            .lock()
            .map_err(|_| tx_consumed_error())?
            .take()
            .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(|_| ())
            .map_err(|e| aws_error("commit_transaction", &e))
    }

    /// 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())?
            .take()
            .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(|_| ())
            .map_err(|e| aws_error("rollback_transaction", &e))
    }

    /// 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,
    > {
        // 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)?;
        execute_statement_raw(
            &self.client,
            &self.resource_arn,
            &self.secret_arn,
            self.database.as_deref(),
            sql,
            params,
            Some(&tx_id),
        )
        .await
    }
}

// =============================================================================
// 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 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`.