149bec6
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 is_write_sql(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])
}
}
fn strip_sql_lead(sql: &str) -> &str {
let mut s = sql.trim_start();
loop {
if let Some(rest) = s.strip_prefix("--") {
let end = rest.find('\n').unwrap_or(rest.len());
s = rest[end..].trim_start();
continue;
}
if let Some(rest) = s.strip_prefix("/*") {
let end = rest.find("*/").map(|i| i + 2).unwrap_or(rest.len());
s = rest[end..].trim_start();
continue;
}
break;
}
s
}
fn first_keyword(sql: &str) -> Option<String> {
let lead = strip_sql_lead(sql);
let head: String = lead
.chars()
.take_while(|c| c.is_ascii_alphanumeric() || *c == '_')
.collect();
if head.is_empty() {
None
} else {
Some(head.to_uppercase())
}
}
fn is_write_sql(sql: &str) -> bool {
let Some(kw) = first_keyword(sql) else {
return true;
};
if kw == "WITH" {
return with_clause_is_write(sql);
}
matches!(
kw.as_str(),
"INSERT"
| "UPDATE"
| "DELETE"
| "MERGE"
| "UPSERT"
| "REPLACE"
| "TRUNCATE"
| "CREATE"
| "DROP"
| "ALTER"
| "RENAME"
| "GRANT"
| "REVOKE"
| "COMMIT"
| "ROLLBACK"
| "SAVEPOINT"
| "LOCK"
| "UNLOCK"
| "COPY"
| "CALL"
| "SET"
| "RESET"
| "DO"
)
}
fn with_clause_is_write(sql: &str) -> bool {
let upper = sql.to_uppercase();
contains_top_level(&upper, "INSERT")
|| contains_top_level(&upper, "UPDATE")
|| contains_top_level(&upper, "DELETE")
|| contains_top_level(&upper, "MERGE")
}
fn contains_top_level(haystack: &str, needle: &str) -> bool {
let bytes = haystack.as_bytes();
let n = needle.as_bytes();
let mut depth: i32 = 0;
let mut in_string = false;
let mut i = 0;
while i < bytes.len() {
let c = bytes[i] as char;
if in_string {
if c == '\'' {
if i + 1 < bytes.len() && bytes[i + 1] == b'\'' {
i += 2;
continue;
}
in_string = false;
}
i += 1;
continue;
}
match c {
'\'' => in_string = true,
'(' => depth += 1,
')' => depth -= 1,
_ if depth == 0 && i + n.len() <= bytes.len() && &bytes[i..i + n.len()] == n => {
let before_ok = i == 0
|| !(bytes[i - 1] as char).is_ascii_alphanumeric() && bytes[i - 1] != b'_';
let after_ok = i + n.len() == bytes.len()
|| !(bytes[i + n.len()] as char).is_ascii_alphanumeric()
&& bytes[i + n.len()] != b'_';
if before_ok && after_ok {
return true;
}
}
_ => {}
}
i += 1;
}
false
}
#[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
}
}