drizzle 0.1.14

//! Transaction wrapper for the Durable Objects SQL driver.
//!
//! Obtained via
//! [`Drizzle::transaction`](crate::builder::sqlite::durable::Drizzle::transaction).
//! Supports the same query-builder surface as `Drizzle` plus nested
//! savepoints through [`Transaction::savepoint`].

use std::marker::PhantomData;
use std::sync::atomic::AtomicU32;

use ::worker::{SqlStorage, SqlStorageValue};
use drizzle_core::error::DrizzleError;
use drizzle_core::traits::ToSQL;

use crate::transaction::savepoint::sync_savepoint;

#[cfg(feature = "sqlite")]
use drizzle_sqlite::{
    builder::{
        self, QueryBuilder, delete::DeleteBuilder, insert::InsertBuilder, select::SelectBuilder,
        update::UpdateBuilder,
    },
    builder::{DeleteInitial, InsertInitial, SelectInitial, UpdateInitial},
    traits::SQLiteTable,
    values::SQLiteValue,
};

use crate::builder::sqlite::durable::sqlite_value_to_storage;

/// Query builder scoped to a [`Transaction`]. See
/// [`crate::transaction::sqlite::typestate::TransactionBuilder`] for the
/// typestate-advancing methods; executor methods live below in this module.
pub type TransactionBuilder<'tx, Schema, Builder, State> =
    crate::transaction::sqlite::typestate::TransactionBuilder<
        'tx,
        Transaction<Schema>,
        Schema,
        Builder,
        State,
    >;

/// Transaction handle for a Durable Object's SQL storage.
///
/// Provides the same query-building surface as
/// [`Drizzle`](crate::builder::sqlite::durable::Drizzle) plus
/// [`Transaction::savepoint`] for nested savepoints.
pub struct Transaction<Schema = ()> {
    conn: SqlStorage,
    savepoint_depth: AtomicU32,
    schema: Schema,
}

impl<Schema> std::fmt::Debug for Transaction<Schema> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Transaction").finish()
    }
}

impl<Schema> Transaction<Schema> {
    pub(crate) fn new(conn: SqlStorage, schema: Schema) -> Self {
        Self {
            conn,
            savepoint_depth: AtomicU32::new(0),
            schema,
        }
    }

    /// Gets a reference to the schema.
    #[inline]
    pub fn schema(&self) -> &Schema {
        &self.schema
    }

    /// Gets a reference to the underlying [`SqlStorage`] handle.
    #[inline]
    pub fn inner(&self) -> &SqlStorage {
        &self.conn
    }

    /// Executes a nested savepoint within this transaction.
    ///
    /// The callback receives a reference to this transaction for executing
    /// queries. If the callback returns `Ok`, the savepoint is released. If
    /// it returns `Err` or panics, the savepoint is rolled back. The outer
    /// transaction is unaffected either way. Savepoints can be nested — each
    /// level gets its own savepoint name.
    pub fn savepoint<F, R>(&self, f: F) -> drizzle_core::error::Result<R>
    where
        F: FnOnce(&Self) -> drizzle_core::error::Result<R>,
    {
        sync_savepoint(
            &self.savepoint_depth,
            |sql| {
                self.conn
                    .exec(sql, None)
                    .map(|_| ())
                    .map_err(|e| DrizzleError::Other(e.to_string().into()))
            },
            || f(self),
        )
    }

    sqlite_transaction_constructors!();

    /// Executes a query within the transaction and returns the number of rows
    /// written.
    pub fn execute<'q, T>(&self, query: T) -> drizzle_core::error::Result<u64>
    where
        T: ToSQL<'q, SQLiteValue<'q>>,
    {
        let cursor = exec_in_tx(&self.conn, &query)?;
        // Drain so `rows_written` is populated.
        let _ = cursor
            .to_array::<serde::de::IgnoredAny>()
            .map_err(|e| DrizzleError::Other(e.to_string().into()))?;
        Ok(cursor.rows_written() as u64)
    }

    /// Runs a query and returns all matching rows within the transaction.
    pub fn all<'q, T, R, C>(&self, query: T) -> drizzle_core::error::Result<C>
    where
        R: for<'de> serde::Deserialize<'de>,
        T: ToSQL<'q, SQLiteValue<'q>>,
        C: Default + Extend<R>,
    {
        let cursor = exec_in_tx(&self.conn, &query)?;
        let rows: Vec<R> = cursor
            .to_array::<R>()
            .map_err(|e| DrizzleError::Other(e.to_string().into()))?;
        let mut out = C::default();
        out.extend(rows);
        Ok(out)
    }

    /// Runs a query and returns a single row within the transaction.
    pub fn get<'q, T, R>(&self, query: T) -> drizzle_core::error::Result<R>
    where
        R: for<'de> serde::Deserialize<'de>,
        T: ToSQL<'q, SQLiteValue<'q>>,
    {
        let cursor = exec_in_tx(&self.conn, &query)?;
        cursor
            .to_array::<R>()
            .map_err(|e| DrizzleError::Other(e.to_string().into()))?
            .into_iter()
            .next()
            .ok_or(DrizzleError::NotFound)
    }
}

fn exec_in_tx<'q, T>(
    conn: &SqlStorage,
    query: &T,
) -> drizzle_core::error::Result<::worker::SqlCursor>
where
    T: ToSQL<'q, SQLiteValue<'q>>,
{
    let sql = query.to_sql();
    let (sql_str, params) = sql.build();
    let values: Vec<SqlStorageValue> = params.into_iter().map(sqlite_value_to_storage).collect();
    conn.exec(&sql_str, Some(values))
        .map_err(|e| DrizzleError::Other(e.to_string().into()))
}

// =============================================================================
// Terminal methods on TransactionBuilder (execute / all / get)
// =============================================================================

#[cfg(feature = "durable")]
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,
{
    /// Runs the query and returns the number of rows written.
    pub fn execute(self) -> drizzle_core::error::Result<u64> {
        let cursor = exec_in_tx(&self.runner.conn, &self.builder.sql)?;
        let _ = cursor
            .to_array::<serde::de::IgnoredAny>()
            .map_err(|e| DrizzleError::Other(e.to_string().into()))?;
        Ok(cursor.rows_written() as u64)
    }

    /// Runs the query and returns all matching rows deserialized into `R`.
    pub fn all<R>(self) -> drizzle_core::error::Result<Vec<R>>
    where
        R: for<'de> serde::Deserialize<'de>,
    {
        let cursor = exec_in_tx(&self.runner.conn, &self.builder.sql)?;
        cursor
            .to_array::<R>()
            .map_err(|e| DrizzleError::Other(e.to_string().into()))
    }

    /// Runs the query and returns the first matching row.
    pub fn get<R>(self) -> drizzle_core::error::Result<R>
    where
        R: for<'de> serde::Deserialize<'de>,
    {
        let cursor = exec_in_tx(&self.runner.conn, &self.builder.sql)?;
        cursor
            .to_array::<R>()
            .map_err(|e| DrizzleError::Other(e.to_string().into()))?
            .into_iter()
            .next()
            .ok_or(DrizzleError::NotFound)
    }
}