35ef7e8
use std::sync::atomic::{AtomicUsize, Ordering};
use crate::error::ExportError;
use config::database::{DatabaseConfig, DatabaseRole as ConfigRole};
use sea_orm::prelude::async_trait::async_trait;
use sea_orm::{Database, 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,
primary: DatabaseConnection,
replicas: Vec<DatabaseConnection>,
round_robin: AtomicUsize,
}
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 = Database::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 = Database::connect(opts).await?;
replicas.push(conn);
}
Ok(Self {
config,
primary,
replicas,
round_robin: AtomicUsize::new(0),
})
}
pub fn config(&self) -> &DatabaseConfig {
&self.config
}
pub fn primary(&self) -> &DatabaseConnection {
&self.primary
}
pub fn replicas(&self) -> &[DatabaseConnection] {
&self.replicas
}
pub fn replica(&self, index: usize) -> Option<&DatabaseConnection> {
self.replicas.get(index)
}
pub fn replica_count(&self) -> usize {
self.replicas.len()
}
pub fn has_replicas(&self) -> bool {
!self.replicas.is_empty()
}
pub fn write(&self) -> &DatabaseConnection {
&self.primary
}
pub fn read(&self) -> &DatabaseConnection {
self.pick_replica().unwrap_or(&self.primary)
}
pub fn for_role(&self, role: DatabaseRole) -> &DatabaseConnection {
match role {
DatabaseRole::Primary => self.write(),
DatabaseRole::Replica => self.read(),
}
}
pub fn route_for(&self, sql: &str) -> &DatabaseConnection {
if db::requires_primary(sql) {
self.write()
} else {
self.read()
}
}
fn pick_replica(&self) -> Option<&DatabaseConnection> {
if self.replicas.is_empty() {
return None;
}
let idx = self.round_robin.fetch_add(1, Ordering::Relaxed) % self.replicas.len();
Some(&self.replicas[idx])
}
}
#[async_trait]
impl sea_orm::ConnectionTrait for AppDatabase {
fn get_database_backend(&self) -> DbBackend {
self.primary.get_database_backend()
}
async fn execute_raw(&self, stmt: Statement) -> Result<ExecResult, DbErr> {
self.route_for(&stmt.sql).execute_raw(stmt).await
}
async fn execute_unprepared(&self, sql: &str) -> Result<ExecResult, DbErr> {
self.route_for(sql).execute_unprepared(sql).await
}
async fn query_one_raw(&self, stmt: Statement) -> Result<Option<QueryResult>, DbErr> {
self.route_for(&stmt.sql).query_one_raw(stmt).await
}
async fn query_all_raw(&self, stmt: Statement) -> Result<Vec<QueryResult>, DbErr> {
self.route_for(&stmt.sql).query_all_raw(stmt).await
}
}