ec0bdb8
use crate::error::ExportError;
use config::database::{DatabaseConfig, DatabaseRole as ConfigRole};
use sea_orm::prelude::async_trait::async_trait;
use sea_orm::{
Database as SeaDatabase, DatabaseConnection, DbBackend, DbErr, ExecResult, QueryResult,
Statement,
};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum DatabaseRole {
Primary,
Replica,
}
impl From<ConfigRole> for DatabaseRole {
fn from(role: ConfigRole) -> Self {
match role {
ConfigRole::Primary => DatabaseRole::Primary,
ConfigRole::Replica => DatabaseRole::Replica,
}
}
}
impl From<DatabaseRole> for ConfigRole {
fn from(role: DatabaseRole) -> Self {
match role {
DatabaseRole::Primary => ConfigRole::Primary,
DatabaseRole::Replica => ConfigRole::Replica,
}
}
}
pub struct AppDatabase {
config: DatabaseConfig,
connection: db::DbConn,
}
impl AppDatabase {
pub async fn connect(config: DatabaseConfig) -> Result<Self, ExportError> {
config.validate().map_err(ExportError::InvalidConfig)?;
let primary_opts = config
.primary_connect_options()
.map_err(ExportError::InvalidConfig)?;
let primary = SeaDatabase::connect(primary_opts).await?;
let mut replicas = Vec::with_capacity(config.replica_count());
for i in 0..config.replica_count() {
let opts = config
.replica_connect_options(i)
.map_err(ExportError::InvalidConfig)?;
let conn = SeaDatabase::connect(opts).await?;
replicas.push(conn);
}
Ok(Self {
config,
connection: db::DbConn::new(primary, replicas),
})
}
pub fn config(&self) -> &DatabaseConfig {
&self.config
}
pub fn primary(&self) -> &DatabaseConnection {
self.connection.write()
}
pub fn replicas(&self) -> &[DatabaseConnection] {
self.connection.replicas()
}
pub fn replica(&self, index: usize) -> Option<&DatabaseConnection> {
self.connection.replicas().get(index)
}
pub fn replica_count(&self) -> usize {
self.connection.replicas().len()
}
pub fn has_replicas(&self) -> bool {
!self.connection.replicas().is_empty()
}
pub fn write(&self) -> &DatabaseConnection {
self.connection.write()
}
pub fn read(&self) -> &db::DbConn {
self.connection.read()
}
pub fn connection(&self) -> &db::DbConn {
&self.connection
}
}
#[async_trait]
impl sea_orm::ConnectionTrait for AppDatabase {
fn get_database_backend(&self) -> DbBackend {
self.connection.get_database_backend()
}
async fn execute_raw(&self, stmt: Statement) -> Result<ExecResult, DbErr> {
self.connection.read().execute_raw(stmt).await
}
async fn execute_unprepared(&self, sql: &str) -> Result<ExecResult, DbErr> {
self.connection.read().execute_unprepared(sql).await
}
async fn query_one_raw(&self, stmt: Statement) -> Result<Option<QueryResult>, DbErr> {
self.connection.read().query_one_raw(stmt).await
}
async fn query_all_raw(&self, stmt: Statement) -> Result<Vec<QueryResult>, DbErr> {
self.connection.read().query_all_raw(stmt).await
}
}