rust_ef_sqlite/
connection.rs

1use async_trait::async_trait;
2use r2d2::PooledConnection;
3use r2d2_sqlite::SqliteConnectionManager;
4use rust_ef::error::{EFError, EFResult};
5use rust_ef::provider::{DbValue, IAsyncConnection, IsolationLevel};
6use std::sync::Arc;
7use std::sync::Mutex as StdMutex;
8use tokio::sync::Mutex as TokioMutex;
9
10use crate::sync_ops::{execute_sync, query_sync};
11
12/// The underlying SQLite connection handle.
13///
14/// `Pooled` wraps an `r2d2::PooledConnection` (used for file-based databases
15/// with a connection pool). The `Option` allows temporarily taking the
16/// connection out so it can be moved into `spawn_blocking` — the connection is
17/// returned from the blocking task and put back for subsequent operations.
18///
19/// `Shared` wraps an `Arc<tokio::sync::Mutex<Connection>>` (used for `:memory:`
20/// databases, which must share a single connection). Operations on `Shared`
21/// call `rusqlite` directly (in-memory operations are sub-millisecond, so
22/// blocking the tokio worker is acceptable).
23enum SqliteConnectionInner {
24    Pooled(StdMutex<Option<PooledConnection<SqliteConnectionManager>>>),
25    Shared(Arc<TokioMutex<rusqlite::Connection>>),
26}
27
28pub struct SqliteConnection {
29    inner: SqliteConnectionInner,
30    #[cfg(feature = "tracing")]
31    slow_query_threshold: Option<std::time::Duration>,
32}
33
34impl SqliteConnection {
35    pub(crate) fn new_pooled(conn: PooledConnection<SqliteConnectionManager>) -> Self {
36        Self {
37            inner: SqliteConnectionInner::Pooled(StdMutex::new(Some(conn))),
38            #[cfg(feature = "tracing")]
39            slow_query_threshold: None,
40        }
41    }
42
43    pub(crate) fn new_shared(conn: Arc<TokioMutex<rusqlite::Connection>>) -> Self {
44        Self {
45            inner: SqliteConnectionInner::Shared(conn),
46            #[cfg(feature = "tracing")]
47            slow_query_threshold: None,
48        }
49    }
50
51    fn threshold(&self) -> Option<std::time::Duration> {
52        #[cfg(feature = "tracing")]
53        {
54            self.slow_query_threshold
55        }
56        #[cfg(not(feature = "tracing"))]
57        {
58            None
59        }
60    }
61}
62
63#[async_trait]
64impl IAsyncConnection for SqliteConnection {
65    async fn execute(&mut self, sql: &str, params: &[DbValue]) -> EFResult<u64> {
66        let _guard = rust_ef::observability::QueryGuard::new(sql, self.threshold());
67        match &self.inner {
68            SqliteConnectionInner::Pooled(m) => {
69                let conn = {
70                    let mut guard = m.lock().unwrap_or_else(|p| p.into_inner());
71                    guard.take().ok_or_else(|| {
72                        EFError::query(
73                            "SQLite connection unavailable (previous spawn_blocking panicked)",
74                        )
75                    })?
76                };
77                let sql_owned = sql.to_string();
78                let params_owned = params.to_vec();
79                let (result, conn) = tokio::task::spawn_blocking(move || {
80                    (execute_sync(&conn, &sql_owned, &params_owned), conn)
81                })
82                .await
83                .map_err(|e| EFError::query(format!("spawn_blocking panicked: {}", e)))?;
84                let mut guard = m.lock().unwrap_or_else(|p| p.into_inner());
85                *guard = Some(conn);
86                result
87            }
88            SqliteConnectionInner::Shared(m) => {
89                let conn = m.lock().await;
90                execute_sync(&conn, sql, params)
91            }
92        }
93    }
94
95    async fn query(&mut self, sql: &str, params: &[DbValue]) -> EFResult<Vec<Vec<DbValue>>> {
96        let _guard = rust_ef::observability::QueryGuard::new(sql, self.threshold());
97        match &self.inner {
98            SqliteConnectionInner::Pooled(m) => {
99                let conn = {
100                    let mut guard = m.lock().unwrap_or_else(|p| p.into_inner());
101                    guard.take().ok_or_else(|| {
102                        EFError::query(
103                            "SQLite connection unavailable (previous spawn_blocking panicked)",
104                        )
105                    })?
106                };
107                let sql_owned = sql.to_string();
108                let params_owned = params.to_vec();
109                let (result, conn) = tokio::task::spawn_blocking(move || {
110                    (query_sync(&conn, &sql_owned, &params_owned), conn)
111                })
112                .await
113                .map_err(|e| EFError::query(format!("spawn_blocking panicked: {}", e)))?;
114                let mut guard = m.lock().unwrap_or_else(|p| p.into_inner());
115                *guard = Some(conn);
116                result
117            }
118            SqliteConnectionInner::Shared(m) => {
119                let conn = m.lock().await;
120                query_sync(&conn, sql, params)
121            }
122        }
123    }
124
125    async fn begin_transaction(&mut self) -> EFResult<()> {
126        self.execute("BEGIN TRANSACTION", &[]).await.map(|_| ())
127    }
128
129    async fn commit_transaction(&mut self) -> EFResult<()> {
130        self.execute("COMMIT", &[]).await.map(|_| ())
131    }
132
133    async fn rollback_transaction(&mut self) -> EFResult<()> {
134        self.execute("ROLLBACK", &[]).await.map(|_| ())
135    }
136
137    async fn create_savepoint(&mut self, name: &str) -> EFResult<()> {
138        self.execute(&format!("SAVEPOINT {}", name), &[])
139            .await
140            .map(|_| ())
141    }
142
143    async fn release_savepoint(&mut self, name: &str) -> EFResult<()> {
144        self.execute(&format!("RELEASE {}", name), &[])
145            .await
146            .map(|_| ())
147    }
148
149    async fn rollback_to_savepoint(&mut self, name: &str) -> EFResult<()> {
150        self.execute(&format!("ROLLBACK TO {}", name), &[])
151            .await
152            .map(|_| ())
153    }
154
155    async fn set_transaction_isolation(&mut self, level: IsolationLevel) -> EFResult<()> {
156        let sql = match level {
157            IsolationLevel::ReadUncommitted => "PRAGMA read_uncommitted = ON",
158            _ => "PRAGMA read_uncommitted = OFF",
159        };
160        self.execute(sql, &[]).await.map(|_| ())
161    }
162
163    #[cfg(feature = "tracing")]
164    fn set_slow_query_threshold(&mut self, threshold: std::time::Duration) {
165        self.slow_query_threshold = Some(threshold);
166    }
167}
rust_ef_sqlite/connection.rs

rust_ef_sqlite/
connection.rs
