use chrono::{DateTime, NaiveDate, Utc};
use rainier_orm::sea_query::Value;
use rainier_orm::{Dialect, Error, ExecOutcome, Executor, PoolConfig, Result, Row};
use sea_orm::{
ConnectOptions, ConnectionTrait, Database, DatabaseConnection, DbBackend, QueryResult,
Statement, Value as SeaValue,
};
#[derive(Clone)]
pub struct SeaOrmExecutor {
db: DatabaseConnection,
dialect: Dialect,
}
impl SeaOrmExecutor {
pub fn new(db: DatabaseConnection) -> Self {
let dialect = match db.get_database_backend() {
DbBackend::MySql => Dialect::MySql,
DbBackend::Postgres => Dialect::Postgres,
DbBackend::Sqlite => Dialect::Sqlite,
};
Self { db, dialect }
}
pub async fn connect(url: &str, pool: &PoolConfig) -> Result<Self> {
let mut opts = ConnectOptions::new(url.to_owned());
opts.max_connections(pool.max_connections)
.min_connections(pool.min_connections)
.acquire_timeout(pool.acquire_timeout)
.test_before_acquire(pool.test_before_acquire);
if let Some(idle) = pool.idle_timeout {
opts.idle_timeout(idle);
}
if let Some(life) = pool.max_lifetime {
opts.max_lifetime(life);
}
let db = Database::connect(opts).await.map_err(Error::from)?;
Ok(Self::new(db))
}
pub fn connection(&self) -> &DatabaseConnection {
&self.db
}
fn backend(&self) -> DbBackend {
match self.dialect {
Dialect::MySql => DbBackend::MySql,
Dialect::Postgres => DbBackend::Postgres,
Dialect::Sqlite => DbBackend::Sqlite,
}
}
fn statement(&self, sql: &str, params: Vec<Value>) -> Statement {
let values: Vec<SeaValue> = params.into_iter().collect();
Statement::from_sql_and_values(self.backend(), sql, values)
}
}
impl Executor for SeaOrmExecutor {
fn dialect(&self) -> Dialect {
self.dialect
}
async fn fetch_all(&self, sql: &str, params: Vec<Value>) -> Result<Vec<Box<dyn Row>>> {
let stmt = self.statement(sql, params);
let rows = self.db.query_all(stmt).await.map_err(Error::from)?;
let dialect = self.dialect;
Ok(rows
.into_iter()
.map(|r| Box::new(SeaOrmRow { row: r, dialect }) as Box<dyn Row>)
.collect())
}
async fn execute(&self, sql: &str, params: Vec<Value>) -> Result<ExecOutcome> {
let stmt = self.statement(sql, params);
let res = self.db.execute(stmt).await.map_err(Error::from)?;
Ok(ExecOutcome {
rows_affected: res.rows_affected(),
last_insert_id: res.last_insert_id() as i64,
})
}
}
struct SeaOrmRow {
row: QueryResult,
dialect: Dialect,
}
macro_rules! getter {
($name:ident, $t:ty) => {
fn $name(&self, col: &str) -> Result<Option<$t>> {
self.row.try_get::<Option<$t>>("", col).map_err(Error::from)
}
};
}
impl Row for SeaOrmRow {
getter!(get_bool, bool);
getter!(get_i32, i32);
getter!(get_i64, i64);
getter!(get_string, String);
getter!(get_bytes, Vec<u8>);
getter!(get_datetime, DateTime<Utc>);
getter!(get_naive_date, NaiveDate);
fn get_u32(&self, col: &str) -> Result<Option<u32>> {
match self.dialect {
Dialect::MySql => self.row.try_get::<Option<u32>>("", col).map_err(Error::from),
_ => {
Ok(self.row.try_get::<Option<i64>>("", col).map_err(Error::from)?.map(|v| v as u32))
}
}
}
fn get_u64(&self, col: &str) -> Result<Option<u64>> {
match self.dialect {
Dialect::MySql => self.row.try_get::<Option<u64>>("", col).map_err(Error::from),
_ => {
Ok(self.row.try_get::<Option<i64>>("", col).map_err(Error::from)?.map(|v| v as u64))
}
}
}
fn get_f64(&self, col: &str) -> Result<Option<f64>> {
if let Ok(v) = self.row.try_get::<Option<f64>>("", col) {
return Ok(v);
}
use std::str::FromStr;
let dec =
self.row.try_get::<Option<rust_decimal::Decimal>>("", col).map_err(Error::from)?;
Ok(dec.map(|d| f64::from_str(&d.to_string()).unwrap_or(0.0)))
}
}