a3s_orm/drivers/sqlite/
transaction.rs

1use std::future::Future;
2use std::pin::Pin;
3use std::sync::atomic::{AtomicU64, Ordering};
4use std::sync::Arc;
5
6use async_trait::async_trait;
7use tokio::sync::OwnedMutexGuard;
8
9use crate::{ExecuteResult, Executor, QueryResult, Transaction};
10
11use super::{SqliteError, SqliteExecutor, SqliteRow, SqliteSavepoint, SqliteSavepointError};
12
13static NEXT_SAVEPOINT_ID: AtomicU64 = AtomicU64::new(1);
14
15/// An exclusive SQLite transaction.
16///
17/// The transaction owns the connection gate, so queries issued through other
18/// clones of the executor wait until this transaction commits or rolls back.
19pub struct SqliteTransaction {
20    executor: SqliteExecutor,
21    guard: Option<OwnedMutexGuard<()>>,
22    operation_lock: Arc<tokio::sync::Mutex<()>>,
23    completed: bool,
24}
25
26impl SqliteTransaction {
27    pub(crate) fn new(executor: SqliteExecutor, guard: OwnedMutexGuard<()>) -> Self {
28        Self {
29            executor,
30            guard: Some(guard),
31            operation_lock: Arc::new(tokio::sync::Mutex::new(())),
32            completed: false,
33        }
34    }
35
36    /// Run an operation in a nested savepoint.
37    ///
38    /// Savepoint cleanup owns the transaction operation gate, so cancellation
39    /// cannot race with subsequent statements in the outer transaction.
40    pub async fn savepoint<T, E, F>(&self, operation: F) -> Result<T, SqliteSavepointError<E>>
41    where
42        T: Send,
43        E: std::error::Error + Send + Sync + 'static,
44        F: for<'a> FnOnce(
45            &'a SqliteSavepoint,
46        ) -> Pin<Box<dyn Future<Output = Result<T, E>> + Send + 'a>>,
47    {
48        let guard = self.operation_lock.clone().lock_owned().await;
49        let id = NEXT_SAVEPOINT_ID.fetch_add(1, Ordering::Relaxed);
50        let savepoint = SqliteSavepoint::begin(self.executor.clone(), guard, id)
51            .await
52            .map_err(SqliteSavepointError::Begin)?;
53        match operation(&savepoint).await {
54            Ok(value) => {
55                savepoint
56                    .release()
57                    .await
58                    .map_err(SqliteSavepointError::Release)?;
59                Ok(value)
60            }
61            Err(operation) => match savepoint.rollback().await {
62                Ok(()) => Err(SqliteSavepointError::Operation(operation)),
63                Err(cleanup) => {
64                    Err(SqliteSavepointError::OperationAndCleanup { operation, cleanup })
65                }
66            },
67        }
68    }
69}
70
71#[async_trait]
72impl Executor for SqliteTransaction {
73    type Row = SqliteRow;
74    type Error = SqliteError;
75
76    async fn execute(&self, query: &crate::CompiledQuery) -> Result<ExecuteResult, Self::Error> {
77        let _operation = self.operation_lock.clone().lock_owned().await;
78        self.executor.execute_unlocked(query).await
79    }
80
81    async fn fetch_all(
82        &self,
83        query: &crate::CompiledQuery,
84    ) -> Result<QueryResult<Self::Row>, Self::Error> {
85        let _operation = self.operation_lock.clone().lock_owned().await;
86        self.executor.fetch_all_unlocked(query).await
87    }
88}
89
90#[async_trait]
91impl Transaction for SqliteTransaction {
92    async fn commit(mut self) -> Result<(), Self::Error> {
93        let _operation = self.operation_lock.lock().await;
94        self.executor.execute_control("COMMIT").await?;
95        self.completed = true;
96        Ok(())
97    }
98
99    async fn rollback(mut self) -> Result<(), Self::Error> {
100        let _operation = self.operation_lock.lock().await;
101        self.executor.execute_control("ROLLBACK").await?;
102        self.completed = true;
103        Ok(())
104    }
105}
106
107impl Drop for SqliteTransaction {
108    fn drop(&mut self) {
109        if self.completed {
110            return;
111        }
112        let Some(guard) = self.guard.take() else {
113            return;
114        };
115        let executor = self.executor.clone();
116        if let Ok(runtime) = tokio::runtime::Handle::try_current() {
117            runtime.spawn(async move {
118                let _guard = guard;
119                let _ = executor.execute_control("ROLLBACK").await;
120            });
121        }
122    }
123}
a3s_orm/drivers/sqlite/transaction.rs

a3s_orm/drivers/sqlite/
transaction.rs
