use sqlx::Acquire;
use std::sync::Arc;
use tokio::sync::{
RwLock,
mpsc::{Receiver, Sender},
oneshot,
};
use wasmtime::component::{Accessor, AccessorTask, JoinHandle};
use crate::{
core::bindings::{
SqlHostState,
executor::{ErasedExecutor, QueryOrRaw},
generated::wasm_sql::core::{
connection::Connection, transaction::Transaction, util_types::Error,
},
},
execute_with,
sqldb::SqlDatabase,
};
#[allow(dead_code)]
pub enum TransactionCommand {
FetchAll {
query: QueryOrRaw,
cb: oneshot::Sender<Result<Vec<<SqlDatabase as sqlx::Database>::Row>, Error>>,
},
Execute {
query: QueryOrRaw,
cb: oneshot::Sender<Result<<SqlDatabase as sqlx::Database>::QueryResult, Error>>,
},
Commit {
cb: oneshot::Sender<Result<(), Error>>,
},
Rollback {
cb: oneshot::Sender<Result<(), Error>>,
},
}
#[allow(dead_code)]
pub struct ConnectionBoundTask {
pub(crate) resource: wasmtime::component::Resource<Connection>,
pub(crate) receiver: Receiver<TransactionCommand>,
}
#[allow(dead_code)]
#[derive(Clone)]
pub enum TransactionImpl {
Tx(Arc<RwLock<Option<sqlx::Transaction<'static, SqlDatabase>>>>),
ConnectionBound {
handle: Arc<JoinHandle>,
sender: Sender<TransactionCommand>,
},
}
#[allow(dead_code)]
trait TransactionErasedOps {
async fn commit(self) -> Result<(), Error>;
async fn rollback(self) -> Result<(), Error>;
}
impl ErasedExecutor<SqlHostState> for TransactionImpl {
async fn fetch_all<T>(
&self,
query: QueryOrRaw,
accessor: &Accessor<T, SqlHostState>,
) -> Result<Vec<<SqlDatabase as sqlx::Database>::Row>, Error> {
match self {
TransactionImpl::Tx(rw_lock) => {
let mut guard = rw_lock.write().await;
if let Some(tx) = guard.as_mut() {
execute_with!(&mut **tx, accessor, query, fetch_all)
} else {
Err(Error::TransactionClosed)
}
}
TransactionImpl::ConnectionBound { handle: _, sender } => {
let (s, r) = oneshot::channel();
sender
.send(TransactionCommand::FetchAll { query, cb: s })
.await?;
r.await?
}
}
}
async fn execute<T>(
&self,
query: super::executor::QueryOrRaw,
accessor: &Accessor<T, SqlHostState>,
) -> Result<<SqlDatabase as sqlx::Database>::QueryResult, Error> {
match self {
TransactionImpl::Tx(rw_lock) => {
let mut guard = rw_lock.write().await;
if let Some(tx) = guard.as_mut() {
execute_with!(&mut **tx, accessor, query, execute)
} else {
Err(Error::TransactionClosed)
}
}
TransactionImpl::ConnectionBound { handle: _, sender } => {
let (s, r) = oneshot::channel();
sender
.send(TransactionCommand::Execute { query, cb: s })
.await?;
r.await?
}
}
}
}
impl TransactionErasedOps for TransactionImpl {
async fn commit(self) -> Result<(), Error> {
match self {
TransactionImpl::Tx(rw_lock) => {
let mut guard = rw_lock.write().await;
if let Some(tx) = guard.take() {
Ok(tx.commit().await?)
} else {
Err(Error::TransactionClosed)
}
}
TransactionImpl::ConnectionBound { handle: _, sender } => {
let (s, r) = oneshot::channel();
sender.send(TransactionCommand::Commit { cb: s }).await?;
r.await?
}
}
}
async fn rollback(self) -> Result<(), Error> {
match self {
TransactionImpl::Tx(rw_lock) => {
let mut guard = rw_lock.write().await;
if let Some(tx) = guard.take() {
Ok(tx.rollback().await?)
} else {
Err(Error::TransactionClosed)
}
}
TransactionImpl::ConnectionBound { handle: _, sender } => {
let (s, r) = oneshot::channel();
sender.send(TransactionCommand::Rollback { cb: s }).await?;
r.await?
}
}
}
}
impl<T> AccessorTask<T, SqlHostState> for ConnectionBoundTask {
async fn run(
mut self,
accessor: &wasmtime::component::Accessor<T, SqlHostState>,
) -> Result<(), wasmtime::Error> {
let conn = accessor.with(|mut access| {
let state = access.get();
state
.table
.get(&self.resource)
.map(|x| x.connection.clone())
})?;
let mut guard = conn.write().await;
let mut tx = Some(guard.begin().await?);
while let Some(cmd) = self.receiver.recv().await {
match cmd {
TransactionCommand::FetchAll { query, cb } => {
let res = if let Some(ref mut tx) = tx {
execute_with!(tx, accessor, query, fetch_all)
} else {
Err(Error::TransactionClosed)
};
let _ = cb.send(res);
}
TransactionCommand::Execute { query, cb } => {
let res = if let Some(ref mut tx) = tx {
execute_with!(tx, accessor, query, execute)
} else {
Err(Error::TransactionClosed)
};
let _ = cb.send(res);
}
TransactionCommand::Commit { cb } => {
let res = if let Some(tx) = tx.take() {
tx.commit().await.map_err(|e| e.into())
} else {
Err(Error::TransactionClosed)
};
let _ = cb.send(res);
}
TransactionCommand::Rollback { cb } => {
let res = if let Some(tx) = tx.take() {
tx.rollback().await.map_err(|e| e.into())
} else {
Err(Error::TransactionClosed)
};
let _ = cb.send(res);
}
}
}
Ok(())
}
}
impl crate::core::bindings::generated::wasm_sql::core::transaction::HostTransaction
for SqlHostState
{
async fn drop(
&mut self,
rep: wasmtime::component::Resource<Transaction>,
) -> wasmtime::Result<()> {
self.table.delete(rep)?;
Ok(())
}
}
impl crate::core::bindings::generated::wasm_sql::core::transaction::Host for SqlHostState {}
impl crate::core::bindings::generated::wasm_sql::core::transaction::HostTransactionWithStore
for SqlHostState
{
async fn commit<T>(
accessor: &wasmtime::component::Accessor<T, Self>,
this: wasmtime::component::Resource<Transaction>,
) -> Result<(), Error> {
let tx_impl = accessor.with(|mut access| access.get().table.delete(this))?;
tx_impl.commit().await
}
async fn rollback<T>(
accessor: &wasmtime::component::Accessor<T, Self>,
this: wasmtime::component::Resource<Transaction>,
) -> Result<(), Error> {
let tx_impl = accessor.with(|mut access| access.get().table.delete(this))?;
tx_impl.rollback().await
}
}