axum-postgres-tx 0.3.0

use std::{ops::Deref, sync::Arc};

#[cfg(feature = "bb8")]
use bb8_postgres::tokio_postgres::{self, Transaction};
#[cfg(feature = "deadpool")]
use deadpool_postgres::tokio_postgres::{self, Transaction};
use parking_lot::{ArcMutexGuard, Mutex, RawMutex};

use super::{OwnedTx, Pool};

#[derive(Clone)]
pub(super) struct Extension(Arc<Mutex<LazyTx>>);

impl Extension {
    pub(super) async fn get(
        &self,
        read_only: bool,
    ) -> Result<ArcMutexGuard<RawMutex, LazyTx>, super::Error> {
        let mut tx = self
            .0
            .try_lock_arc()
            .ok_or(super::Error::OverlappingExtractors)?;
        tx.get(read_only).await?;
        Ok(tx)
    }

    /// Commit the transaction if one has been started and the lock is available.
    ///
    /// This is called by the [`Layer`](super::layer::Layer) after the response
    /// is sent. If the lock cannot be acquired (e.g., the [`Tx`](super::tx::Tx)
    /// guard is still held by a spawned task), the commit is silently skipped
    /// and the transaction will be rolled back when the connection is returned
    /// to the pool.
    pub(super) async fn commit(&self) -> Result<(), tokio_postgres::Error> {
        if let Some(mut tx) = self.0.try_lock_arc() {
            tx.commit().await?;
        }
        Ok(())
    }
}

impl From<Pool> for Extension {
    fn from(value: Pool) -> Self {
        Self(Arc::new(Mutex::new(value.into())))
    }
}

pub(super) struct LazyTx(LazyTxState);

impl LazyTx {
    async fn get(&mut self, read_only: bool) -> Result<(), super::Error> {
        match &self.0 {
            LazyTxState::Pool(pool) => {
                #[cfg(feature = "bb8")]
                let conn = pool.get_owned().await?;
                #[cfg(feature = "deadpool")]
                let conn = pool.get().await?;
                self.0 = LazyTxState::Tx(OwnedTx::new(conn, read_only).await?);
                Ok(())
            }
            LazyTxState::Tx(_) => Ok(()),
            LazyTxState::Resolved => Err(super::Error::OverlappingExtractors),
        }
    }

    async fn commit(&mut self) -> Result<(), tokio_postgres::Error> {
        match std::mem::replace(&mut self.0, LazyTxState::Resolved) {
            LazyTxState::Pool(_) | LazyTxState::Resolved => Ok(()),
            LazyTxState::Tx(tx) => tx.tx.commit().await,
        }
    }
}

impl From<Pool> for LazyTx {
    fn from(value: Pool) -> Self {
        Self(LazyTxState::Pool(value))
    }
}

impl AsRef<Transaction<'static>> for LazyTx {
    fn as_ref(&self) -> &Transaction<'static> {
        match &self.0 {
            LazyTxState::Pool(_) => {
                panic!("BUG: exposed LazyTx pool")
            }
            LazyTxState::Tx(tx) => &tx.tx,
            LazyTxState::Resolved => {
                panic!("BUG: exposed resolved LazyTx")
            }
        }
    }
}

impl Deref for LazyTx {
    type Target = Transaction<'static>;
    fn deref(&self) -> &Self::Target {
        self.as_ref()
    }
}

enum LazyTxState {
    Pool(Pool),
    Tx(OwnedTx),
    Resolved,
}