use crate::error::{EFError, EFResult};
use crate::provider::{IAsyncConnection, IsolationLevel};
use crate::session::{Session, SessionConnection, SharedTxn};
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum TxOutcome<T> {
Commit(T),
Rollback(T),
}
pub trait ITransaction: Send + Sync {
fn commit(self: Box<Self>) -> Pin<Box<dyn Future<Output = EFResult<()>> + Send + 'static>>;
fn rollback(self: Box<Self>) -> Pin<Box<dyn Future<Output = EFResult<()>> + Send + 'static>>;
fn create_point<'a>(
&'a mut self,
name: &'a str,
) -> Pin<Box<dyn Future<Output = EFResult<()>> + Send + 'a>>;
fn release_point<'a>(
&'a mut self,
name: &'a str,
) -> Pin<Box<dyn Future<Output = EFResult<()>> + Send + 'a>>;
fn rollback_point<'a>(
&'a mut self,
name: &'a str,
) -> Pin<Box<dyn Future<Output = EFResult<()>> + Send + 'a>>;
fn set_isolation<'a>(
&'a mut self,
level: IsolationLevel,
) -> Pin<Box<dyn Future<Output = EFResult<()>> + Send + 'a>>;
fn connection(&mut self) -> &mut (dyn IAsyncConnection + Send);
fn is_scoped(&self) -> bool {
false
}
}
enum TxnState {
Owned(Option<Box<dyn IAsyncConnection>>),
Shared {
shared: Arc<SharedTxn>,
facade: SessionConnection,
},
}
pub struct DbTransaction {
state: Option<TxnState>,
session: Option<Arc<Session>>,
}
impl DbTransaction {
pub fn new(conn: Box<dyn IAsyncConnection>) -> Self {
Self {
state: Some(TxnState::Owned(Some(conn))),
session: None,
}
}
pub(crate) fn new_shared(
shared: Arc<SharedTxn>,
facade: SessionConnection,
session: Option<Arc<Session>>,
) -> Self {
Self {
state: Some(TxnState::Shared { shared, facade }),
session,
}
}
}
impl ITransaction for DbTransaction {
fn commit(self: Box<Self>) -> Pin<Box<dyn Future<Output = EFResult<()>> + Send + 'static>> {
Box::pin(async move {
let mut this = self;
let Some(state) = this.state.take() else {
return Err(EFError::transaction("transaction already consumed"));
};
let result = match state {
TxnState::Owned(conn) => match conn {
Some(mut conn) => conn.commit_transaction().await,
None => Err(EFError::transaction("transaction already consumed")),
},
TxnState::Shared { shared, .. } => match shared.take().await {
Ok(mut conn) => conn.commit_transaction().await,
Err(e) => Err(e),
},
};
if let Some(session) = this.session.as_ref() {
session.clear().await;
}
result
})
}
fn rollback(self: Box<Self>) -> Pin<Box<dyn Future<Output = EFResult<()>> + Send + 'static>> {
Box::pin(async move {
let mut this = self;
let Some(state) = this.state.take() else {
return Err(EFError::transaction("transaction already consumed"));
};
let result = match state {
TxnState::Owned(conn) => match conn {
Some(mut conn) => conn.rollback_transaction().await,
None => Err(EFError::transaction("transaction already consumed")),
},
TxnState::Shared { shared, .. } => match shared.take().await {
Ok(mut conn) => conn.rollback_transaction().await,
Err(e) => Err(e),
},
};
if let Some(session) = this.session.as_ref() {
session.clear().await;
}
result
})
}
fn create_point<'a>(
&'a mut self,
name: &'a str,
) -> Pin<Box<dyn Future<Output = EFResult<()>> + Send + 'a>> {
Box::pin(async move {
let state = self
.state
.as_mut()
.ok_or_else(|| EFError::transaction("transaction already consumed"))?;
match state {
TxnState::Owned(conn) => {
conn.as_mut()
.ok_or_else(|| EFError::transaction("transaction consumed"))?
.create_savepoint(name)
.await
}
TxnState::Shared { facade, .. } => facade.create_savepoint(name).await,
}
})
}
fn release_point<'a>(
&'a mut self,
name: &'a str,
) -> Pin<Box<dyn Future<Output = EFResult<()>> + Send + 'a>> {
Box::pin(async move {
let state = self
.state
.as_mut()
.ok_or_else(|| EFError::transaction("transaction already consumed"))?;
match state {
TxnState::Owned(conn) => {
conn.as_mut()
.ok_or_else(|| EFError::transaction("transaction consumed"))?
.release_savepoint(name)
.await
}
TxnState::Shared { facade, .. } => facade.release_savepoint(name).await,
}
})
}
fn rollback_point<'a>(
&'a mut self,
name: &'a str,
) -> Pin<Box<dyn Future<Output = EFResult<()>> + Send + 'a>> {
Box::pin(async move {
let state = self
.state
.as_mut()
.ok_or_else(|| EFError::transaction("transaction already consumed"))?;
match state {
TxnState::Owned(conn) => {
conn.as_mut()
.ok_or_else(|| EFError::transaction("transaction consumed"))?
.rollback_to_savepoint(name)
.await
}
TxnState::Shared { facade, .. } => facade.rollback_to_savepoint(name).await,
}
})
}
fn set_isolation<'a>(
&'a mut self,
level: IsolationLevel,
) -> Pin<Box<dyn Future<Output = EFResult<()>> + Send + 'a>> {
Box::pin(async move {
let state = self
.state
.as_mut()
.ok_or_else(|| EFError::transaction("transaction already consumed"))?;
match state {
TxnState::Owned(conn) => {
conn.as_mut()
.ok_or_else(|| EFError::transaction("transaction consumed"))?
.set_transaction_isolation(level)
.await
}
TxnState::Shared { facade, .. } => facade.set_transaction_isolation(level).await,
}
})
}
fn connection(&mut self) -> &mut (dyn IAsyncConnection + Send) {
match self.state.as_mut().expect("transaction already consumed") {
TxnState::Owned(conn) => conn
.as_mut()
.expect("transaction already consumed")
.as_mut(),
TxnState::Shared { facade, .. } => facade,
}
}
fn is_scoped(&self) -> bool {
matches!(&self.state, Some(TxnState::Shared { .. }))
}
}
impl Drop for DbTransaction {
fn drop(&mut self) {
let state = self.state.take();
let session = self.session.take();
if let Some(session) = &session {
session.clear_now();
}
let Some(state) = state else {
return;
};
let Ok(handle) = tokio::runtime::Handle::try_current() else {
return;
};
handle.spawn(async move {
match state {
TxnState::Owned(Some(mut conn)) => {
let _ = conn.rollback_transaction().await;
}
TxnState::Owned(None) => {}
TxnState::Shared { shared, .. } => {
if let Ok(mut conn) = shared.take().await {
let _ = conn.rollback_transaction().await;
}
}
}
if let Some(session) = session {
session.clear().await;
}
});
}
}