sqlx_plus_core/

sqlx_base.rs

use std::sync::Once;

use async_trait::async_trait;
use futures::executor::block_on;
use sqlx::{postgres::PgRow, query, query_as, Error, FromRow, Pool, Row, Transaction};

#[cfg(feature = "mysql")]
use sqlx::MySql;
#[cfg(feature = "postgres")]
use sqlx::Postgres;
#[cfg(feature = "sqlite")]
use sqlx::Sqlite;
use tracing::debug;

// use crate::logger;

#[cfg(feature = "sqlite")]
pub type DbPool = Pool<Sqlite>;
#[cfg(feature = "postgres")]
pub type DbPool = Pool<Postgres>;
#[cfg(feature = "mysql")]
pub type DbPool = Pool<MySql>;

#[cfg(feature = "sqlite")]
pub type DbRow = sqlx::sqlite::SqliteRow;
#[cfg(feature = "postgres")]
pub type DbRow = sqlx::postgres::PgRow;
#[cfg(feature = "mysql")]
pub type DbRow = sqlx::mysql::MySqlRow;

#[cfg(feature = "sqlite")]
pub type DbTransaction<'a> = Transaction<'a, Sqlite>;
#[cfg(feature = "postgres")]
pub type DbTransaction<'a> = Transaction<'a, Postgres>;
#[cfg(feature = "mysql")]
pub type DbTransaction<'a> = Transaction<'a, MySql>;

#[derive(Debug)]
pub enum DbType {
    Sqlite,
    Postgres,
    Mysql,
}

pub fn get_db_type() -> DbType {
    #[cfg(feature = "sqlite")]
    {
        return DbType::Sqlite;
    }

    #[cfg(feature = "postgres")]
    {
        return DbType::Postgres;
    }

    #[cfg(feature = "mysql")]
    {
        return DbType::Mysql;
    }
}

static INIT: Once = Once::new();

// fn initialize_logger() {
//     INIT.call_once(|| {
//         block_on(logger::run());
//     });
// }

/// 定义 XPlusCore trait
#[async_trait]
pub trait SqlxBase {
    fn pool(&self) -> &DbPool;

    async fn create_table_if_not_exists(
        &self,
        table_name: &str,
        create_sql: &str,
    ) -> Result<(), Error> {
        let db_type = get_db_type();
        let table_exists_sql = match db_type {
            DbType::Sqlite => format!(
                "SELECT COUNT(*) FROM sqlite_master WHERE type='table' AND name='{}'",
                table_name
            ),
            DbType::Postgres => format!(
                "SELECT EXISTS (SELECT 1 FROM information_schema.tables WHERE table_schema = 'public' AND table_name = '{}')",
                table_name
            ),
            DbType::Mysql => format!(
                "SELECT COUNT(*) FROM information_schema.tables WHERE table_schema = DATABASE() AND table_name = '{}'",
                table_name
            ),
        };

        let table_exists: (bool,) = query_as(&table_exists_sql).fetch_one(self.pool()).await?;
        if !table_exists.0 {
            query(create_sql).execute(self.pool()).await?;
        }
        Ok(())
    }

    async fn insert(
        &self,
        sql: &str,
        transaction: Option<&mut DbTransaction<'_>>,
    ) -> Result<u64, Error> {
        let sql_lower = sql.to_lowercase();
        if sql_lower.contains("returning") {
            return Err(Error::Protocol("The 'insert' method does not support 'RETURNING' clause. Use 'insert_and_return_key' instead.".into()));
        }
        debug!("The 「Sqlx pPlus」「Insert)」 sql:{}", sql_lower);
        let result = match transaction {
            Some(tx) => query(sql).execute(&mut **tx).await?,
            None => query(sql).execute(self.pool()).await?,
        };
        let rows_affected = result.rows_affected();
        Ok(rows_affected)
    }

    async fn insert_and_return_key<T>(
        &self,
        sql: &str,
        transaction: Option<&mut DbTransaction<'_>>,
    ) -> Result<T, Error>
    where
        T: for<'r> FromRow<'r, DbRow> + Unpin + Send,
    {
        let sql_lower = sql.to_lowercase();
        if !sql_lower.contains("returning") {
            return Err(Error::Protocol(
                "The 'insert_and_return_key' method requires a 'RETURNING' clause.".into(),
            ));
        }
        let result = match transaction {
            Some(tx) => query_as::<_, T>(&sql).fetch_one(&mut **tx).await?,
            None => query_as::<_, T>(&sql).fetch_one(self.pool()).await?,
        };
        Ok(result)
    }

    async fn update(
        &self,
        sql: &str,
        transaction: Option<&mut DbTransaction<'_>>,
    ) -> Result<u64, Error> {
        let result = match transaction {
            Some(tx) => query(sql).execute(&mut **tx).await?,
            None => query(sql).execute(self.pool()).await?,
        };
        Ok(result.rows_affected())
    }

    async fn delete(
        &self,
        sql: &str,
        transaction: Option<&mut DbTransaction<'_>>,
    ) -> Result<u64, Error> {
        let result = match transaction {
            Some(tx) => query(sql).execute(&mut **tx).await?,
            None => query(sql).execute(self.pool()).await?,
        };
        Ok(result.rows_affected())
    }

    async fn select_one<T>(&self, sql: &str) -> Result<Option<T>, Error>
    where
        T: for<'r> FromRow<'r, DbRow> + Unpin + Send,
    {
        let result = query_as::<_, T>(sql).fetch_optional(self.pool()).await?;
        Ok(result)
    }

    async fn query_one(&self, sql: &str) -> Result<PgRow, Error> {
        let row = query(sql).fetch_one(self.pool()).await?;
        Ok(row)
    }

    async fn select_all<T>(&self, sql: &str) -> Result<Vec<T>, Error>
    where
        T: for<'r> FromRow<'r, DbRow> + Unpin + Send,
    {
        let result = query_as::<_, T>(sql).fetch_all(self.pool()).await?;
        Ok(result)
    }

    async fn select_page<T>(&self, sql: &str) -> Result<Vec<T>, Error>
    where
        T: for<'r> FromRow<'r, DbRow> + Unpin + Send,
    {
        let sql_lower = sql.to_lowercase();
        if !sql_lower.contains("limit") || !sql_lower.contains("offset") {
            return Err(Error::Protocol("SQL must contain LIMIT and OFFSET".into()));
        }

        let result = query_as::<_, T>(sql).fetch_all(self.pool()).await?;
        Ok(result)
    }

    async fn count(&self, sql: &str) -> Result<i64, Error> {
        let row = query(sql).fetch_one(self.pool()).await?;
        let count: i64 = row.try_get(0)?;
        Ok(count)
    }

    async fn health_check(&self) -> Result<bool, Error> {
        let row: (i64,) = sqlx::query_as("SELECT $1")
            .bind(150_i64)
            .fetch_one(self.pool())
            .await
            .unwrap_or_else(|_| panic!("Sqlx Plus Failed to to execute query"));
        Ok(row.0 == 150)
    }

    async fn begin(&self) -> Result<DbTransaction<'_>, Error> {
        let transaction = self.pool().begin().await?;
        Ok(transaction)
    }

    async fn commit(transaction: DbTransaction<'_>) -> Result<(), Error> {
        transaction.commit().await?;
        Ok(())
    }

    async fn rollback(transaction: DbTransaction<'_>) -> Result<(), Error> {
        transaction.rollback().await?;
        Ok(())
    }
}

#[async_trait]
impl SqlxBase for DbPool {
    fn pool(&self) -> &DbPool {
        // initialize_logger();
        self
    }
}