sqlite-rwc 0.4.0

use crate::drivers::Driver;
use crate::pool::{
    ConnectionAdapter, ConnectionPool, ConnectionPoolConfig, ConnectionPoolError, PooledConnection,
};
use crate::{ReadTransaction, Transaction};
use flume::Sender;
use std::marker::PhantomData;
use std::sync::Arc;
use tracing::error;

#[derive(Debug, thiserror::Error)]
pub enum AsyncConnectionError<E> {
    #[error(transparent)]
    Connection(#[from] E),
    #[error("Communication with worker failed")]
    Worker,
}

impl<E> AsyncConnectionError<E> {
    fn into<T: From<E>>(self) -> AsyncConnectionError<T> {
        match self {
            AsyncConnectionError::Connection(e) => AsyncConnectionError::Connection(T::from(e)),
            AsyncConnectionError::Worker => AsyncConnectionError::Worker,
        }
    }
}

pub trait AsyncRuntime: Send + 'static + Sync {
    type JoinError: std::error::Error + Send + Sync + 'static;
    type JoinHandle<T: Send + 'static>: Future<Output = Result<T, Self::JoinError>> + Send + 'static;
    fn spawn_blocking<F, T>(closure: F) -> Self::JoinHandle<T>
    where
        F: FnOnce() -> T + Send + 'static,
        T: Send + 'static;
}

pub type AsyncConnectionPool<T> = ConnectionPool<T, AsyncConnectionAdapter<T>>;
pub type AsyncPooledConnection<T> = PooledConnection<T, AsyncConnectionAdapter<T>>;

impl<T, A> ConnectionPool<T, A>
where
    T: Driver,
    A: ConnectionAdapter<T> + Send + 'static,
    <T as Driver>::Connection: Send + 'static,
    <T as Driver>::Error: Send + 'static,
    <T as Driver>::ConnectionError: Send + 'static,
{
    /// Create new connection pool in an async safe way,
    ///
    /// # Errors
    ///
    /// Returns error if connection pool could not be created or we failed to join the task.
    pub async fn new_async<R: AsyncRuntime>(
        config: ConnectionPoolConfig,
    ) -> Result<Arc<Self>, ConnectionPoolError<T::ConnectionError>> {
        R::spawn_blocking(move || Self::new(config))
            .await
            .map_err(|e| ConnectionPoolError::Other(Box::new(e)))?
    }

    /// Similar to [`ConnectionPool::connection`] but async friendly.
    ///
    /// # Errors
    ///
    /// Returns error if the connection could not be acquired or failed to join the task.
    #[allow(clippy::type_complexity)]
    pub async fn connection_async<R: AsyncRuntime>(
        self: &Arc<Self>,
    ) -> Result<PooledConnection<T, A>, ConnectionPoolError<T::Error>> {
        let pool = self.clone();
        R::spawn_blocking(move || pool.connection())
            .await
            .map_err(|e| ConnectionPoolError::Other(Box::new(e)))?
    }
}

/// Defines behaviors suitable for use in an async environment.
///
/// Since the [`Driver`]s are synchronous, each connection is placed on a dedicated thread
/// and closures are transferred over a channel to these threads where they will be executed.
pub struct AsyncConnectionAdapter<T: Driver> {
    sender: Sender<AsyncCommand<T::Connection>>,
}

impl<T> AsyncConnectionAdapter<T>
where
    T: Driver,
    <T as Driver>::Connection: Send + 'static,
    <T as Driver>::Error: Send + 'static,
{
    fn new(mut connection: T::Connection) -> Result<Self, std::io::Error> {
        let (sender, receiver) = flume::bounded(1);

        std::thread::Builder::new()
            .name("sqlite-wrc-async".to_owned())
            .spawn(move || {
                while let Ok(cmd) = receiver.recv() {
                    match cmd {
                        AsyncCommand::Execute(execute) => {
                            execute(&mut connection);
                        }
                    }
                }
            })?;

        Ok(AsyncConnectionAdapter { sender })
    }

    async fn read<F, R, E>(&self, closure: F) -> Result<R, AsyncConnectionError<E>>
    where
        F: FnOnce(&mut T::Connection) -> Result<R, E> + Send + 'static,
        R: Send + 'static,
        E: From<T::Error> + Send + 'static,
    {
        let (sx, rx) = oneshot::channel();
        self.sender
            .send_async(AsyncCommand::Execute(Box::new(
                move |conn: &mut T::Connection| {
                    let r = closure(conn);
                    let _ = sx.send(r);
                },
            )))
            .await
            .map_err(|_| AsyncConnectionError::Worker)?;

        Ok(rx.await.map_err(|_| AsyncConnectionError::Worker)??)
    }
    async fn read_transaction_closure<F, R, E>(
        &self,
        closure: F,
    ) -> Result<R, AsyncConnectionError<E>>
    where
        F: FnOnce(&mut ReadTransaction<'_, T>) -> Result<R, E> + Send + 'static,
        R: Send + 'static,
        E: From<T::Error> + Send + 'static,
    {
        let (sx, rx) = oneshot::channel();
        self.sender
            .send_async(AsyncCommand::Execute(Box::new(
                move |conn: &mut T::Connection| {
                    let r = ReadTransaction::scoped(conn, closure);
                    let _ = sx.send(r);
                },
            )))
            .await
            .map_err(|_| AsyncConnectionError::Worker)?;

        Ok(rx.await.map_err(|_| AsyncConnectionError::Worker)??)
    }

    async fn transaction_closure<F, R, E>(
        &mut self,
        pool: Arc<ConnectionPool<T, Self>>,
        closure: F,
    ) -> Result<R, AsyncConnectionError<E>>
    where
        F: FnOnce(&mut Transaction<'_, T>) -> Result<R, E> + Send + 'static,
        R: Send + 'static,
        E: From<T::Error> + Send + 'static,
    {
        let (sx, rx) = oneshot::channel();
        self.sender
            .send_async(AsyncCommand::Execute(Box::new(
                move |_: &mut T::Connection| {
                    let r = pool.transaction_closure(closure);
                    let _ = sx.send(r);
                },
            )))
            .await
            .map_err(|_| AsyncConnectionError::Worker)?;

        Ok(rx.await.map_err(|_| AsyncConnectionError::Worker)??)
    }

    async fn transaction_closure_async<F, R, E>(
        &mut self,
        pool: Arc<ConnectionPool<T, Self>>,
        closure: F,
    ) -> Result<R, AsyncConnectionError<E>>
    where
        F: AsyncFnOnce(&mut AsyncTransaction<T>) -> Result<R, AsyncConnectionError<E>>,
        R: Send + 'static,
        E: From<T::Error> + Send + 'static,
    {
        let mut tx = self
            .transaction(pool)
            .await
            .map_err(AsyncConnectionError::into)?;
        let r = closure(&mut tx).await;
        let tx_result = if r.is_ok() {
            tx.commit().await
        } else {
            tx.rollback().await
        };
        tx_result.map_err(AsyncConnectionError::into)?;
        r
    }

    async fn read_transaction_closure_async<F, R, E>(
        &mut self,
        closure: F,
    ) -> Result<R, AsyncConnectionError<E>>
    where
        F: AsyncFnOnce(&mut AsyncReadTransaction<T>) -> Result<R, AsyncConnectionError<E>>,
        R: Send + 'static,
        E: From<T::Error> + Send + 'static,
    {
        let mut tx = self
            .read_transaction()
            .await
            .map_err(AsyncConnectionError::into)?;
        let r = closure(&mut tx).await;
        tx.commit().await.map_err(AsyncConnectionError::into)?;
        r
    }

    async fn transaction(
        &mut self,
        pool: Arc<ConnectionPool<T, Self>>,
    ) -> Result<AsyncTransaction<'_, T>, AsyncConnectionError<T::Error>> {
        let (sx, rx) = oneshot::channel();
        self.sender
            .send_async(AsyncCommand::Execute(Box::new(
                move |_: &mut T::Connection| {
                    let (async_rx, mut tx) = match pool.transaction() {
                        Ok(tx) => {
                            let (async_tx, async_rx) = flume::bounded(1);
                            let _ = sx.send(Ok(async_tx));
                            (async_rx, tx)
                        }
                        Err(e) => {
                            let _ = sx.send(Err(e));
                            return;
                        }
                    };
                    while let Ok(cmd) = async_rx.recv() {
                        match cmd {
                            AsyncTxCommand::Op(op) => op(&mut tx),
                            AsyncTxCommand::Commit(sx) => {
                                let _ = sx.send(tx.commit());
                                return;
                            }
                            AsyncTxCommand::Rollback(sx) => {
                                let _ = sx.send(tx.rollback());
                                return;
                            }
                            AsyncTxCommand::ReadOp(_) => {
                                tracing::warn!("Received invalid command");
                            }
                        }
                    }
                },
            )))
            .await
            .map_err(|_| AsyncConnectionError::Worker)?;

        Ok(rx
            .await
            .map_err(|_| AsyncConnectionError::Worker)?
            .map(AsyncTransaction::new)?)
    }

    async fn read_transaction(
        &mut self,
    ) -> Result<AsyncReadTransaction<'_, T>, AsyncConnectionError<T::Error>> {
        let (sx, rx) = oneshot::channel();
        self.sender
            .send_async(AsyncCommand::Execute(Box::new(
                move |conn: &mut T::Connection| {
                    let (async_rx, mut tx): (_, ReadTransaction<T>) =
                        match ReadTransaction::new(conn) {
                            Ok(tx) => {
                                let (async_tx, async_rx) = flume::bounded(1);
                                let _ = sx.send(Ok(async_tx));
                                (async_rx, tx)
                            }
                            Err(e) => {
                                let _ = sx.send(Err(e));
                                return;
                            }
                        };
                    while let Ok(cmd) = async_rx.recv() {
                        match cmd {
                            AsyncTxCommand::ReadOp(op) => op(&mut tx),
                            AsyncTxCommand::Commit(sx) => {
                                let _ = sx.send(tx.commit());
                                return;
                            }
                            AsyncTxCommand::Rollback(sx) => {
                                let _ = sx.send(tx.rollback());
                                return;
                            }
                            AsyncTxCommand::Op(_) => {
                                tracing::warn!("Received invalid command");
                            }
                        }
                    }
                },
            )))
            .await
            .map_err(|_| AsyncConnectionError::Worker)?;

        Ok(rx
            .await
            .map_err(|_| AsyncConnectionError::Worker)?
            .map(AsyncReadTransaction::new)?)
    }
}

impl<T> ConnectionAdapter<T> for AsyncConnectionAdapter<T>
where
    T: Driver,
    <T as Driver>::Connection: Send + 'static,
    <T as Driver>::Error: Send + 'static,
{
    fn from_driver_connection(connection: T::Connection) -> Self {
        AsyncConnectionAdapter::new(connection).expect("Failed to create named thread")
    }
}

impl<T> PooledConnection<T, AsyncConnectionAdapter<T>>
where
    T: Driver,
    <T as Driver>::Connection: Send + 'static,
    <T as Driver>::Error: Send + 'static,
{
    /// Perform a read operation without a stable scope.
    ///
    /// Unlike [`Self::read_transaction_closure`], this method uses the default Sqlite transaction behavior
    /// where each statement will create its own transaction.
    ///
    /// # Errors
    ///
    /// Returns error if the closure failed.
    pub async fn read<F, R, E>(&mut self, closure: F) -> Result<R, AsyncConnectionError<E>>
    where
        F: FnOnce(&mut T::Connection) -> Result<R, E> + Send + 'static,
        R: Send + 'static,
        E: From<T::Error> + Send + 'static,
    {
        self.connection_mut().read(closure).await
    }

    /// Perform a read operation with a stable scope.
    ///
    /// This method uses an explicit transaction to ensure the reader does not take into
    /// account ongoing modifications of writers during the execution of the statements in
    /// the scope of the closure.
    ///
    /// Use [`Self::read`] if you want to have the default sqlite behavior, where each statement
    /// gets their own transaction.
    ///
    /// # Errors
    ///
    /// Returns error if the closure failed or if the read transaction could not be initiated.
    pub async fn read_transaction_closure<F, R, E>(
        &mut self,
        closure: F,
    ) -> Result<R, AsyncConnectionError<E>>
    where
        F: FnOnce(&mut ReadTransaction<'_, T>) -> Result<R, E> + Send + 'static,
        R: Send + 'static,
        E: From<T::Error> + Send + 'static,
    {
        self.connection_mut()
            .read_transaction_closure(closure)
            .await
    }

    /// Create a transaction and run the given `closure` in its scope.
    ///
    /// If the closure returns an error the transaction is aborted, otherwise it is commited.
    ///
    /// # Errors
    ///
    /// Returns error if the closure failed or the transaction failed to commit.
    pub async fn transaction_closure<F, R, E>(
        &mut self,
        closure: F,
    ) -> Result<R, AsyncConnectionError<E>>
    where
        F: FnOnce(&mut Transaction<'_, T>) -> Result<R, E> + Send + 'static,
        R: Send + 'static,
        E: From<T::Error> + Send + 'static,
    {
        let pool = self.pool.clone();
        self.connection_mut()
            .transaction_closure(pool, closure)
            .await
    }

    /// Run a transaction but in an async closure.
    ///
    /// This allows the transaction code to call other async code. Note that it is less efficient
    /// then [`Self::transaction_closure`] as each statement is sent over a channel rather than the entire
    /// transaction closure.
    ///
    /// The behavior is similar to [`Self::transaction_closure`], if the closure returns an error, the
    /// transaction is rolled back, otherwise it is commited.
    ///
    /// # Errors
    ///
    /// Returns errors if the transaction failed or the closure failed.
    pub async fn transaction_closure_async<F, R, E>(
        &mut self,
        closure: F,
    ) -> Result<R, AsyncConnectionError<E>>
    where
        F: AsyncFnOnce(&mut AsyncTransaction<T>) -> Result<R, AsyncConnectionError<E>>,
        R: Send + 'static,
        E: From<T::Error> + Send + 'static,
    {
        let pool = self.pool.clone();
        self.connection_mut()
            .transaction_closure_async(pool, closure)
            .await
    }

    /// Run a read transaction but in an async closure.
    ///
    /// This allows the transaction code to call other async code. Note that it is less efficient
    /// then [`Self::read_transaction_closure`] as each statement is sent over a channel rather than the entire
    /// transaction closure.
    ///
    /// The behavior is similar to [`Self::read_transaction_closure`], if the closure returns an error, the
    /// transaction is rolled back, otherwise it is commited.
    ///
    /// # Errors
    ///
    /// Returns errors if the transaction failed or the closure failed.
    pub async fn read_transaction_closure_async<F, R, E>(
        &mut self,
        closure: F,
    ) -> Result<R, AsyncConnectionError<E>>
    where
        F: AsyncFnOnce(&mut AsyncReadTransaction<T>) -> Result<R, AsyncConnectionError<E>>,
        R: Send + 'static,
        E: From<T::Error> + Send + 'static,
    {
        self.connection_mut()
            .read_transaction_closure_async(closure)
            .await
    }

    /// Create a new transaction type that needs to be manually commited.
    ///
    /// If the transaction is dropped without an explicit call to [`AsyncTransaction::rollback`] or
    /// [`AsyncTransaction::commit`], the transaction will be rolled back.
    ///
    /// # Errors
    ///
    /// Returns an error if the transaction could not be created.
    pub async fn transaction(
        &mut self,
    ) -> Result<AsyncTransaction<'_, T>, AsyncConnectionError<T::Error>> {
        let pool = self.pool.clone();
        self.connection_mut().transaction(pool).await
    }

    /// Create a new read transaction type that needs to be manually commited.
    ///
    /// If the transaction is dropped without an explicit call to [`AsyncReadTransaction::commit`], the transaction will be rolled back.
    ///
    ///
    /// # Errors
    ///
    /// Returns an error if the transaction could not be created.
    pub async fn read_transaction(
        &mut self,
    ) -> Result<AsyncReadTransaction<'_, T>, AsyncConnectionError<T::Error>> {
        self.connection_mut().read_transaction().await
    }
}

enum AsyncTxCommand<T: Driver> {
    #[allow(clippy::type_complexity)]
    Op(Box<dyn FnOnce(&mut Transaction<'_, T>) + Send + 'static>),
    #[allow(clippy::type_complexity)]
    ReadOp(Box<dyn FnOnce(&mut ReadTransaction<'_, T>) + Send + 'static>),
    Commit(oneshot::Sender<Result<(), T::Error>>),
    Rollback(oneshot::Sender<Result<(), T::Error>>),
}

/// A means to interface with an ongoing transaction initiated via
/// [`PooledConnection::transaction`].
pub struct AsyncTransaction<'c, T: Driver>(
    flume::Sender<AsyncTxCommand<T>>,
    PhantomData<&'c mut T>,
);

impl<T: Driver> AsyncTransaction<'_, T> {
    fn new(sender: flume::Sender<AsyncTxCommand<T>>) -> Self {
        Self(sender, PhantomData)
    }

    /// Execute a single statement on the active transaction.
    ///
    /// # Errors
    ///
    /// Returns an error if the communication with the worker thread failed or the closure
    /// returned an error.
    pub async fn run<F, R, E>(&mut self, closure: F) -> Result<R, AsyncConnectionError<E>>
    where
        F: FnOnce(&mut Transaction<'_, T>) -> Result<R, E> + Send + 'static,
        R: Send + 'static,
        E: From<T::Error> + Send + 'static,
    {
        let (sx, rx) = oneshot::channel();
        self.0
            .send_async(AsyncTxCommand::Op(Box::new(
                move |tx: &mut Transaction<'_, T>| {
                    let r = closure(tx);
                    let _ = sx.send(r);
                },
            )))
            .await
            .map_err(|_| AsyncConnectionError::Worker)?;

        Ok(rx.await.map_err(|_| AsyncConnectionError::Worker)??)
    }

    /// Commit the transaction.
    ///
    /// # Errors
    ///
    /// Returns an error if the transaction could not be commited.
    pub async fn commit(self) -> Result<(), AsyncConnectionError<T::Error>> {
        let (sx, rx) = oneshot::channel();
        self.0
            .send_async(AsyncTxCommand::Commit(sx))
            .await
            .map_err(|_| AsyncConnectionError::Worker)?;
        Ok(rx.await.map_err(|_| AsyncConnectionError::Worker)??)
    }

    /// Rollback the transaction.
    ///
    /// # Errors
    ///
    /// Returns an error if the transaction could not be rolled back.
    pub async fn rollback(self) -> Result<(), AsyncConnectionError<T::Error>> {
        let (sx, rx) = oneshot::channel();
        self.0
            .send_async(AsyncTxCommand::Rollback(sx))
            .await
            .map_err(|_| AsyncConnectionError::Worker)?;
        Ok(rx.await.map_err(|_| AsyncConnectionError::Worker)??)
    }
}

/// A means to interface with an ongoing read transaction initiated via
/// [`PooledConnection::read_transaction`].
pub struct AsyncReadTransaction<'c, T: Driver>(
    flume::Sender<AsyncTxCommand<T>>,
    PhantomData<&'c mut T>,
);

impl<T: Driver> AsyncReadTransaction<'_, T> {
    fn new(sender: flume::Sender<AsyncTxCommand<T>>) -> Self {
        Self(sender, PhantomData)
    }

    /// Execute a single statement on the active transaction.
    ///
    /// # Errors
    ///
    /// Returns an error if the communication with the worker thread failed or the closure
    /// returned an error.
    pub async fn run<F, R, E>(&mut self, closure: F) -> Result<R, AsyncConnectionError<E>>
    where
        F: FnOnce(&mut ReadTransaction<'_, T>) -> Result<R, E> + Send + 'static,
        R: Send + 'static,
        E: From<T::Error> + Send + 'static,
    {
        let (sx, rx) = oneshot::channel();
        self.0
            .send_async(AsyncTxCommand::ReadOp(Box::new(
                move |tx: &mut ReadTransaction<'_, T>| {
                    let r = closure(tx);
                    let _ = sx.send(r);
                },
            )))
            .await
            .map_err(|_| AsyncConnectionError::Worker)?;

        Ok(rx.await.map_err(|_| AsyncConnectionError::Worker)??)
    }

    /// Commit the transaction.
    ///
    /// # Errors
    ///
    /// Returns an error if the transaction could not be commited.
    pub async fn commit(self) -> Result<(), AsyncConnectionError<T::Error>> {
        let (sx, rx) = oneshot::channel();
        self.0
            .send_async(AsyncTxCommand::Commit(sx))
            .await
            .map_err(|_| AsyncConnectionError::Worker)?;
        Ok(rx.await.map_err(|_| AsyncConnectionError::Worker)??)
    }
}

enum AsyncCommand<T> {
    Execute(Box<dyn FnOnce(&mut T) + Send + 'static>),
}