drizzle 0.1.14

use drizzle_core::error::DrizzleError;
use drizzle_core::traits::ToSQL;
use drizzle_postgres::builder::{DeleteInitial, InsertInitial, SelectInitial, UpdateInitial};
use drizzle_postgres::traits::PostgresTable;
use postgres::fallible_iterator::FallibleIterator;
use postgres::{Row, Transaction as PgTransaction};
use std::cell::RefCell;
use std::marker::PhantomData;
use std::sync::atomic::AtomicU32;

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

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

use drizzle_postgres::builder::{
    self, QueryBuilder, delete::DeleteBuilder, insert::InsertBuilder, select::SelectBuilder,
    update::UpdateBuilder,
};
use drizzle_postgres::common::PostgresTransactionType;
use drizzle_postgres::values::PostgresValue;
use smallvec::SmallVec;

use crate::builder::postgres::postgres_sync::Rows;

/// `postgres_sync`-specific transaction builder.
///
/// This is a thin type alias over the dialect-shared
/// [`crate::transaction::postgres::typestate::TransactionBuilder`]; every
/// typestate-advancing method (`.value`/`.values`/`.r#where`/`.set`/
/// `.on_conflict`/`.returning`/`.from`/`.join`/etc.) lives on the generic
/// struct over there. Executor methods (`.execute`/`.all`/`.rows`/`.get`)
/// — the only parts that need `postgres::Transaction`-specific access —
/// stay below in this module.
pub type TransactionBuilder<'tx, 'conn, Schema, Builder, State> =
    crate::transaction::postgres::typestate::TransactionBuilder<
        'tx,
        &'tx Transaction<'conn, Schema>,
        Schema,
        Builder,
        State,
    >;

/// Transaction wrapper that provides the same query building capabilities as Drizzle
pub struct Transaction<'conn, Schema = ()> {
    tx: RefCell<Option<PgTransaction<'conn>>>,
    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 {
        f.debug_struct("Transaction")
            .field("tx_type", &self.tx_type)
            .field("is_active", &self.tx.borrow().is_some())
            .finish()
    }
}

impl<'conn, Schema> Transaction<'conn, Schema> {
    /// Creates a new transaction wrapper
    pub(crate) const fn new(
        tx: PgTransaction<'conn>,
        tx_type: PostgresTransactionType,
        schema: Schema,
    ) -> Self {
        Self {
            tx: RefCell::new(Some(tx)),
            tx_type,
            savepoint_depth: AtomicU32::new(0),
            schema,
        }
    }

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

    /// Gets the transaction type
    #[inline]
    pub const fn tx_type(&self) -> PostgresTransactionType {
        self.tx_type
    }

    /// Executes a raw SQL string with no parameters.
    fn execute_raw(&self, sql: &str) -> drizzle_core::error::Result<()> {
        let mut tx_ref = self.tx.borrow_mut();
        let tx = tx_ref.as_mut().ok_or_else(tx_consumed_error)?;
        tx.execute(sql, &[]).map_err(DrizzleError::from)?;
        Ok(())
    }

    /// 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.
    ///
    /// ```no_run
    /// # use drizzle::postgres::prelude::*;
    /// # use drizzle::postgres::sync::Drizzle;
    /// # use drizzle::postgres::common::PostgresTransactionType;
    /// # #[PostgresTable] struct User { #[column(serial, primary)] id: i32, name: String }
    /// # #[derive(PostgresSchema)] struct S { user: User }
    /// # fn main() -> drizzle::Result<()> {
    /// # let client = ::postgres::Client::connect("host=localhost user=postgres", ::postgres::NoTls)?;
    /// # let (mut db, S { user }) = Drizzle::new(client, S::new());
    /// db.transaction(PostgresTransactionType::ReadCommitted, |tx| {
    ///     tx.insert(user).values([InsertUser::new("Alice")]).execute()?;
    ///
    ///     // This savepoint fails — only its changes roll back
    ///     let _: Result<(), _> = tx.savepoint(|stx| {
    ///         stx.insert(user).values([InsertUser::new("Bad")]).execute()?;
    ///         Err(drizzle::error::DrizzleError::Other("oops".into()))
    ///     });
    ///
    ///     let users: Vec<SelectUser> = tx.select(()).from(user).all()?;
    ///     assert_eq!(users.len(), 1); // only Alice
    ///     Ok(())
    /// })?;
    /// # Ok(()) }
    /// ```
    ///
    /// # Errors
    ///
    /// Returns [`DrizzleError`] if the savepoint cannot be created/released, or the inner closure returns an error.
    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.execute_raw(sql),
            || f(self),
        )
    }

    postgres_transaction_constructors!('conn);

    /// Execute a statement within the transaction and return the number of affected rows.
    ///
    /// # Errors
    ///
    /// Returns [`DrizzleError`] if the database call fails or the SQL is invalid.
    pub fn execute<'q, T>(&self, query: T) -> drizzle_core::error::Result<u64>
    where
        T: ToSQL<'q, PostgresValue<'q>>,
    {
        #[cfg(feature = "profiling")]
        drizzle_core::drizzle_profile_scope!("postgres.sync", "tx.execute");
        let query_sql = query.to_sql();
        let (sql, params) = query_sql.build();
        drizzle_core::drizzle_trace_query!(&sql, params.len());

        let mut tx_ref = self.tx.borrow_mut();
        let tx = tx_ref.as_mut().ok_or_else(tx_consumed_error)?;

        let param_refs = {
            #[cfg(feature = "profiling")]
            drizzle_core::drizzle_profile_scope!("postgres.sync", "tx.execute.param_refs");
            let mut param_refs: SmallVec<[&(dyn postgres::types::ToSql + Sync); 8]> =
                SmallVec::with_capacity(params.len());
            param_refs.extend(
                params
                    .iter()
                    .map(|&p| p as &(dyn postgres::types::ToSql + Sync)),
            );
            param_refs
        };

        let mut typed_params: SmallVec<
            [(&(dyn postgres::types::ToSql + Sync), postgres::types::Type); 8],
        > = SmallVec::with_capacity(params.len());
        let mut all_typed = true;
        for p in &params {
            if let Some(ty) = crate::builder::postgres::prepared_common::postgres_sync_param_type(p)
            {
                typed_params.push((*p as &(dyn postgres::types::ToSql + Sync), ty));
            } else {
                all_typed = false;
                break;
            }
        }

        if all_typed {
            #[cfg(feature = "profiling")]
            drizzle_core::drizzle_profile_scope!("postgres.sync", "tx.execute.db_typed");
            let mut rows = tx
                .query_typed_raw(&sql, typed_params)
                .map_err(DrizzleError::from)?;
            while rows.next().map_err(DrizzleError::from)?.is_some() {}
            return Ok(rows.rows_affected().unwrap_or(0));
        }

        #[cfg(feature = "profiling")]
        drizzle_core::drizzle_profile_scope!("postgres.sync", "tx.execute.db");
        Ok(tx
            .execute(&sql, &param_refs[..])
            .map_err(DrizzleError::from)?)
    }

    /// Runs the query and returns all matching rows (for SELECT queries)
    ///
    /// # Errors
    ///
    /// Returns [`DrizzleError`] if the query fails or row decoding fails.
    pub 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: std::iter::FromIterator<R>,
    {
        self.rows(query)?
            .collect::<drizzle_core::error::Result<C>>()
    }

    /// Runs the query and returns a lazy row cursor.
    ///
    /// # Errors
    ///
    /// Returns [`DrizzleError`] if the query fails.
    pub 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>>,
    {
        #[cfg(feature = "profiling")]
        drizzle_core::drizzle_profile_scope!("postgres.sync", "tx.all");
        let sql = query.to_sql();
        let (sql_str, params) = sql.build();
        drizzle_core::drizzle_trace_query!(&sql_str, params.len());

        #[cfg(feature = "profiling")]
        drizzle_core::drizzle_profile_scope!("postgres.sync", "tx.all.param_refs");
        let mut param_refs: SmallVec<[&(dyn postgres::types::ToSql + Sync); 8]> =
            SmallVec::with_capacity(params.len());
        param_refs.extend(
            params
                .iter()
                .map(|&p| p as &(dyn postgres::types::ToSql + Sync)),
        );

        let mut tx_ref = self.tx.borrow_mut();
        let tx = tx_ref.as_mut().ok_or_else(tx_consumed_error)?;

        let rows = tx
            .query(&sql_str, &param_refs[..])
            .map_err(DrizzleError::from)?;

        Ok(Rows::new(rows))
    }

    /// Runs the query and returns a single row (for SELECT queries)
    ///
    /// # Errors
    ///
    /// Returns [`DrizzleError`] if the query fails, no rows match, or decoding fails.
    pub 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>>,
    {
        #[cfg(feature = "profiling")]
        drizzle_core::drizzle_profile_scope!("postgres.sync", "tx.get");
        let sql = query.to_sql();
        let (sql_str, params) = sql.build();
        drizzle_core::drizzle_trace_query!(&sql_str, params.len());

        #[cfg(feature = "profiling")]
        drizzle_core::drizzle_profile_scope!("postgres.sync", "tx.get.param_refs");
        let mut param_refs: SmallVec<[&(dyn postgres::types::ToSql + Sync); 8]> =
            SmallVec::with_capacity(params.len());
        param_refs.extend(
            params
                .iter()
                .map(|&p| p as &(dyn postgres::types::ToSql + Sync)),
        );

        let mut tx_ref = self.tx.borrow_mut();
        let tx = tx_ref.as_mut().ok_or_else(tx_consumed_error)?;

        let row = tx
            .query_one(&sql_str, &param_refs[..])
            .map_err(DrizzleError::from)?;

        R::try_from(&row).map_err(Into::into)
    }

    /// Commits the transaction
    pub(crate) fn commit(&self) -> drizzle_core::error::Result<()> {
        let tx = self.tx.borrow_mut().take().ok_or_else(tx_consumed_error)?;
        tx.commit().map_err(DrizzleError::from)
    }

    /// Rolls back the transaction
    pub(crate) fn rollback(&self) -> drizzle_core::error::Result<()> {
        let tx = self.tx.borrow_mut().take().ok_or_else(tx_consumed_error)?;
        tx.rollback().map_err(DrizzleError::from)
    }
}

// `TransactionBuilder<CTEInit>::select` and `.with` are now provided by
// the shared `DrizzleBuilder` typestate impls (see
// `crate::builder::postgres::common`).

impl<'tx, 'q, S, Schema, State, Table, Mk, Rw, Grouped>
    TransactionBuilder<'tx, '_, S, QueryBuilder<'q, Schema, State, Table, Mk, Rw, Grouped>, State>
where
    State: builder::ExecutableState,
{
    /// Runs the query and returns the number of affected rows
    pub fn execute(self) -> drizzle_core::error::Result<u64> {
        #[cfg(feature = "profiling")]
        drizzle_core::drizzle_profile_scope!("postgres.sync", "tx_builder.execute");
        let (sql_str, params) = self.builder.sql.build();
        drizzle_core::drizzle_trace_query!(&sql_str, params.len());

        let mut tx_ref = self.runner.tx.borrow_mut();
        let tx = tx_ref.as_mut().ok_or_else(tx_consumed_error)?;

        let param_refs = {
            #[cfg(feature = "profiling")]
            drizzle_core::drizzle_profile_scope!("postgres.sync", "tx_builder.execute.param_refs");
            let mut param_refs: SmallVec<[&(dyn postgres::types::ToSql + Sync); 8]> =
                SmallVec::with_capacity(params.len());
            param_refs.extend(
                params
                    .iter()
                    .map(|&p| p as &(dyn postgres::types::ToSql + Sync)),
            );
            param_refs
        };

        let mut typed_params: SmallVec<
            [(&(dyn postgres::types::ToSql + Sync), postgres::types::Type); 8],
        > = SmallVec::with_capacity(params.len());
        let mut all_typed = true;
        for p in &params {
            if let Some(ty) = crate::builder::postgres::prepared_common::postgres_sync_param_type(p)
            {
                typed_params.push((*p as &(dyn postgres::types::ToSql + Sync), ty));
            } else {
                all_typed = false;
                break;
            }
        }

        if all_typed {
            #[cfg(feature = "profiling")]
            drizzle_core::drizzle_profile_scope!("postgres.sync", "tx_builder.execute.db_typed");
            let mut rows = tx
                .query_typed_raw(&sql_str, typed_params)
                .map_err(DrizzleError::from)?;
            while rows.next().map_err(DrizzleError::from)?.is_some() {}
            return Ok(rows.rows_affected().unwrap_or(0));
        }

        #[cfg(feature = "profiling")]
        drizzle_core::drizzle_profile_scope!("postgres.sync", "tx_builder.execute.db");
        Ok(tx
            .execute(&sql_str, &param_refs[..])
            .map_err(DrizzleError::from)?)
    }

    /// Runs the query and returns all matching rows using the builder's row type.
    pub fn all<R, Proof, AggProof>(self) -> drizzle_core::error::Result<Vec<R>>
    where
        for<'r> Mk: drizzle_core::row::DecodeSelectedRef<&'r ::postgres::Row, R>
            + drizzle_core::row::MarkerScopeValidFor<Proof>
            + drizzle_core::row::StrictDecodeMarker
            + drizzle_core::row::MarkerColumnCountValid<::postgres::Row, Rw, R>,
        Mk: drizzle_core::row::MarkerAggValidFor<Grouped, AggProof>,
    {
        #[cfg(feature = "profiling")]
        drizzle_core::drizzle_profile_scope!("postgres.sync", "tx_builder.all");
        let (sql_str, params) = self.builder.sql.build();
        drizzle_core::drizzle_trace_query!(&sql_str, params.len());

        #[cfg(feature = "profiling")]
        drizzle_core::drizzle_profile_scope!("postgres.sync", "tx_builder.all.param_refs");
        let mut param_refs: SmallVec<[&(dyn postgres::types::ToSql + Sync); 8]> =
            SmallVec::with_capacity(params.len());
        param_refs.extend(
            params
                .iter()
                .map(|&p| p as &(dyn postgres::types::ToSql + Sync)),
        );

        let mut tx_ref = self.runner.tx.borrow_mut();
        let tx = tx_ref.as_mut().ok_or_else(tx_consumed_error)?;
        let rows = tx
            .query(&sql_str, &param_refs[..])
            .map_err(DrizzleError::from)?;

        let mut decoded = Vec::with_capacity(rows.len());
        for row in &rows {
            decoded.push(<Mk as drizzle_core::row::DecodeSelectedRef<
                &::postgres::Row,
                R,
            >>::decode(row)?);
        }
        Ok(decoded)
    }

    /// Runs the query and returns a lazy row cursor using the builder's row type.
    pub 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>,
    {
        #[cfg(feature = "profiling")]
        drizzle_core::drizzle_profile_scope!("postgres.sync", "tx_builder.rows");
        let (sql_str, params) = self.builder.sql.build();
        drizzle_core::drizzle_trace_query!(&sql_str, params.len());

        #[cfg(feature = "profiling")]
        drizzle_core::drizzle_profile_scope!("postgres.sync", "tx_builder.rows.param_refs");
        let mut param_refs: SmallVec<[&(dyn postgres::types::ToSql + Sync); 8]> =
            SmallVec::with_capacity(params.len());
        param_refs.extend(
            params
                .iter()
                .map(|&p| p as &(dyn postgres::types::ToSql + Sync)),
        );

        let mut tx_ref = self.runner.tx.borrow_mut();
        let tx = tx_ref.as_mut().ok_or_else(tx_consumed_error)?;

        let rows = tx
            .query(&sql_str, &param_refs[..])
            .map_err(DrizzleError::from)?;

        Ok(Rows::new(rows))
    }

    /// Runs the query and returns a single row using the builder's row type.
    pub fn get<R, Proof, AggProof>(self) -> drizzle_core::error::Result<R>
    where
        for<'r> Mk: drizzle_core::row::DecodeSelectedRef<&'r ::postgres::Row, R>
            + drizzle_core::row::MarkerScopeValidFor<Proof>
            + drizzle_core::row::StrictDecodeMarker
            + drizzle_core::row::MarkerColumnCountValid<::postgres::Row, Rw, R>,
        Mk: drizzle_core::row::MarkerAggValidFor<Grouped, AggProof>,
    {
        #[cfg(feature = "profiling")]
        drizzle_core::drizzle_profile_scope!("postgres.sync", "tx_builder.get");
        let (sql_str, params) = self.builder.sql.build();
        drizzle_core::drizzle_trace_query!(&sql_str, params.len());

        #[cfg(feature = "profiling")]
        drizzle_core::drizzle_profile_scope!("postgres.sync", "tx_builder.get.param_refs");
        let mut param_refs: SmallVec<[&(dyn postgres::types::ToSql + Sync); 8]> =
            SmallVec::with_capacity(params.len());
        param_refs.extend(
            params
                .iter()
                .map(|&p| p as &(dyn postgres::types::ToSql + Sync)),
        );

        let mut tx_ref = self.runner.tx.borrow_mut();
        let tx = tx_ref.as_mut().ok_or_else(tx_consumed_error)?;
        let row = tx
            .query_one(&sql_str, &param_refs[..])
            .map_err(DrizzleError::from)?;

        <Mk as drizzle_core::row::DecodeSelectedRef<&::postgres::Row, R>>::decode(&row)
    }
}

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