7cf7ea3
use std::fmt::{Debug, Display};
use std::future::Future;
use std::pin::Pin;
use std::sync::atomic::Ordering;
use sea_orm::prelude::async_trait::async_trait;
use sea_orm::{
AccessMode, ConnectionTrait, DatabaseConnection, DatabaseTransaction, DbBackend, DbErr,
ExecResult, IsolationLevel, QueryResult, Statement, TransactionError, TransactionOptions,
TransactionTrait,
};
use crate::conn::DbConn;
use crate::routing;
impl DbConn {
async fn route<'a, T, F, Fut>(&'a self, sql: &str, operation: F) -> Result<T, DbErr>
where
F: Fn(&'a DatabaseConnection) -> Fut + Send,
Fut: Future<Output = Result<T, DbErr>> + Send,
T: Send,
{
if !routing::requires_primary(sql)
&& let Some(replica) = self.replica()
{
match operation(replica).await {
Ok(result) => return Ok(result),
// 副本失败时静默回落会掩盖副本故障，把首次失败记进日志。
Err(error) => {
if self
.replica_failure_logged
.compare_exchange(false, true, Ordering::Relaxed, Ordering::Relaxed)
.is_ok()
{
tracing::warn!(
%error,
"replica query failed; falling back to the primary"
);
}
}
}
}
operation(self.write()).await
}
}
#[async_trait]
impl ConnectionTrait for DbConn {
fn get_database_backend(&self) -> DbBackend {
self.write().get_database_backend()
}
async fn execute_raw(&self, stmt: Statement) -> Result<ExecResult, DbErr> {
self.route(&stmt.sql, |connection| connection.execute_raw(stmt.clone()))
.await
}
async fn execute_unprepared(&self, sql: &str) -> Result<ExecResult, DbErr> {
self.route(sql, |connection| connection.execute_unprepared(sql))
.await
}
async fn query_one_raw(&self, stmt: Statement) -> Result<Option<QueryResult>, DbErr> {
self.route(&stmt.sql, |connection| {
connection.query_one_raw(stmt.clone())
})
.await
}
async fn query_all_raw(&self, stmt: Statement) -> Result<Vec<QueryResult>, DbErr> {
self.route(&stmt.sql, |connection| {
connection.query_all_raw(stmt.clone())
})
.await
}
fn support_returning(&self) -> bool {
self.write().support_returning()
}
fn is_mock_connection(&self) -> bool {
self.write().is_mock_connection()
}
}
#[async_trait]
impl TransactionTrait for DbConn {
type Transaction = DatabaseTransaction;
async fn begin(&self) -> Result<Self::Transaction, DbErr> {
self.write().begin().await
}
async fn begin_with_config(
&self,
isolation_level: Option<IsolationLevel>,
access_mode: Option<AccessMode>,
) -> Result<Self::Transaction, DbErr> {
self.write()
.begin_with_config(isolation_level, access_mode)
.await
}
async fn begin_with_options(
&self,
options: TransactionOptions,
) -> Result<Self::Transaction, DbErr> {
self.write().begin_with_options(options).await
}
async fn transaction<F, T, E>(&self, callback: F) -> Result<T, TransactionError<E>>
where
F: for<'c> FnOnce(
&'c Self::Transaction,
) -> Pin<Box<dyn Future<Output = Result<T, E>> + Send + 'c>>
+ Send,
T: Send,
E: Display + Debug + Send,
{
self.write().transaction(callback).await
}
async fn transaction_with_config<F, T, E>(
&self,
callback: F,
isolation_level: Option<IsolationLevel>,
access_mode: Option<AccessMode>,
) -> Result<T, TransactionError<E>>
where
F: for<'c> FnOnce(
&'c Self::Transaction,
) -> Pin<Box<dyn Future<Output = Result<T, E>> + Send + 'c>>
+ Send,
T: Send,
E: Display + Debug + Send,
{
self.write()
.transaction_with_config(callback, isolation_level, access_mode)
.await
}
}