bodega 0.7.1

use snafu::{ResultExt, Snafu};
use sqlx::{postgres::PgPoolOptions, Executor, Pool, Postgres};

use crate::{Error, Result};

pub type Db = Pool<Postgres>;

pub async fn new_db_pool(db_connect_url: &str, max_connections: u32) -> Result<Db> {
    PgPoolOptions::new()
        .max_connections(max_connections)
        .connect(db_connect_url)
        .await
        .map_err(|e| Error::FailedToCreateDBPool {
            message: e.to_string(),
        })
}

#[derive(Debug, Snafu)]
pub enum DbModelManagerError {
    #[snafu(display("Error checking DB connectivity: "))]
    Connectivity { source: sqlx::Error },

    #[snafu(display("Error starting DB transaction: "))]
    TransactionInit { source: sqlx::Error },

    #[snafu(display("Error committing DB transaction: "))]
    TransactionCommit { source: sqlx::Error },

    #[snafu(display("Error rolling back DB transaction: "))]
    TransactionRollback { source: sqlx::Error },
}

impl DbModelManagerError {
    /// Convenience method for getting a reference to the underlying [`sqlx::Error`].
    pub fn source(&self) -> &sqlx::Error {
        match self {
            DbModelManagerError::Connectivity { source } => source,
            DbModelManagerError::TransactionInit { source } => source,
            DbModelManagerError::TransactionCommit { source } => source,
            DbModelManagerError::TransactionRollback { source } => source,
        }
    }
}

/// Acts as an interface to a db connection pool that is Clone + Send + Sync.
///
/// This type can be cloned freely, as the underlying pool is already a smart
/// pointer.
#[derive(Debug, Clone)]
pub struct DbModelManager {
    db: Db,
}

impl DbModelManager {
    pub async fn new(db_connect_url: &str, max_connections: u32) -> Result<Self> {
        let db = new_db_pool(db_connect_url, max_connections).await?;

        Ok(DbModelManager { db })
    }

    pub fn new_from_pool(pool: Db) -> Self {
        pool.into()
    }

    pub async fn check_db_connectivity(&self) -> Result<()> {
        sqlx::query("SELECT 1")
            .execute(self.db())
            .await
            .context(ConnectivitySnafu)?;
        Ok(())
    }

    /// Begin a new transaction.
    pub async fn begin(&self) -> Result<Transaction<'_>> {
        let mut raw = self.db().begin().await.context(TransactionInitSnafu)?;
        raw.execute("SET TRANSACTION ISOLATION LEVEL SERIALIZABLE;")
            .await
            .context(TransactionInitSnafu)?;

        Ok(Transaction(raw))
    }

    /// Get a reference to the db pool.
    ///
    /// We would prefer not to expose this but because we're providing this type
    /// in this crate, we have to expose it. Care should be used to not use this
    /// outside of running migrations, perhaps.
    pub fn db(&self) -> &Db {
        &self.db
    }
}

impl From<Db> for DbModelManager {
    fn from(db: Db) -> Self {
        Self { db }
    }
}

/// We wrap a transaction in this type to prevent a caller from outside of this
/// crate having direct access to the transaction, and therefore access to an
/// executor that can manipulate the database without going through the exposed
/// interfaces of this crate.
pub struct Transaction<'a>(sqlx::Transaction<'a, Postgres>);

impl Transaction<'_> {
    pub(crate) fn executor(&mut self) -> impl Executor<'_, Database = Postgres> {
        &mut *self.0
    }

    /// Commit the underlying transaction.
    pub async fn commit(self) -> Result<()> {
        Ok(self.0.commit().await.context(TransactionCommitSnafu)?)
    }

    /// Roll the underlying transaction back.
    pub async fn rollback(self) -> Result<()> {
        Ok(self.0.rollback().await.context(TransactionRollbackSnafu)?)
    }
}

/// Indicates that this type can provide an executor.
pub trait AsExecutor: private::ActualExecutor {
    fn is_transaction(&self) -> bool {
        false
    }
}

// by splitting the trait, we can expose functions generic over implementors
// of the public trait, but have the actual function for getting the executor
// in the private interface. This prevents people from getting at the raw
// executor outside of the context of this crate. If they _could_ get at this,
// then it would be possible to execute a query against the db outside of this
// crate.
pub(crate) mod private {
    use sqlx::{Executor, Postgres};

    pub trait ActualExecutor {
        fn as_executor(&mut self) -> impl Executor<'_, Database = Postgres>;
    }
}

impl private::ActualExecutor for Transaction<'_> {
    fn as_executor(&mut self) -> impl Executor<'_, Database = Postgres> {
        self.executor()
    }
}

impl AsExecutor for Transaction<'_> {
    fn is_transaction(&self) -> bool {
        true
    }
}

impl private::ActualExecutor for DbModelManager {
    // an unfortunate side-effect of supporting transactions is that we need
    // as_executor to operate on mutable references. This means that, when using
    // the db model manager as the executor, we need a mutable reference,
    // despite nothing actually requiring it to be mutable other than this
    // interface.
    fn as_executor(&mut self) -> impl Executor<'_, Database = Postgres> {
        self.db()
    }
}

// the model manager by itself is not a transaction
impl AsExecutor for DbModelManager {}