clt_database/turso/src/
connection.rs

1use crate::turso::assert_send_sync;
2use crate::turso::transaction::DropBehavior;
3use crate::turso::transaction::TransactionBehavior;
4use crate::turso::Error;
5use crate::turso::IntoParams;
6use crate::turso::Row;
7use crate::turso::Rows;
8use crate::turso::Statement;
9use std::fmt::Debug;
10use std::sync::atomic::AtomicU8;
11use std::sync::atomic::Ordering;
12use std::sync::Arc;
13use std::sync::Mutex;
14use std::task::Waker;
15pub type Result<T> = std::result::Result<T, Error>;
16
17/// Atomic wrapper for [DropBehavior]
18pub(crate) struct AtomicDropBehavior {
19    inner: AtomicU8,
20}
21
22impl AtomicDropBehavior {
23    fn new(behavior: DropBehavior) -> Self {
24        Self {
25            inner: AtomicU8::new(behavior.into()),
26        }
27    }
28
29    fn load(&self, ordering: Ordering) -> DropBehavior {
30        self.inner.load(ordering).into()
31    }
32
33    pub(crate) fn store(&self, behavior: DropBehavior, ordering: Ordering) {
34        self.inner.store(behavior.into(), ordering);
35    }
36}
37
38// A database connection.
39pub struct Connection {
40    /// Inner is an Option so that when a Connection is dropped we can take the inner
41    /// (Actual connection) out of it and put it back into the ConnectionPool
42    /// the only time inner will be None is just before the Connection is freed after the
43    /// inner connection has been recyled into the connection pool
44    inner: Option<Arc<crate::turso_sdk_kit::rsapi::TursoConnection>>,
45    pub(crate) transaction_behavior: TransactionBehavior,
46    /// If there is a dangling transaction after it was dropped without being finished,
47    /// [Connection::dangling_tx] will be set to the [DropBehavior] of the dangling transaction,
48    /// and the corresponding action will be taken when a new transaction is requested
49    /// or the connection queries/executes.
50    /// We cannot do this eagerly on Drop because drop is not async.
51    ///
52    /// By default, the value is [DropBehavior::Ignore] which effectively does nothing.
53    pub(crate) dangling_tx: AtomicDropBehavior,
54    pub(crate) extra_io: Option<Arc<dyn Fn(Waker) -> Result<()> + Send + Sync>>,
55}
56
57assert_send_sync!(Connection);
58
59impl Clone for Connection {
60    fn clone(&self) -> Self {
61        Self {
62            inner: self.inner.clone(),
63            transaction_behavior: self.transaction_behavior,
64            dangling_tx: AtomicDropBehavior::new(self.dangling_tx.load(Ordering::SeqCst)),
65            extra_io: self.extra_io.clone(),
66        }
67    }
68}
69
70impl Connection {
71    pub(crate) fn create(
72        conn: Arc<crate::turso_sdk_kit::rsapi::TursoConnection>,
73        extra_io: Option<Arc<dyn Fn(Waker) -> Result<()> + Send + Sync>>,
74    ) -> Self {
75        #[allow(clippy::arc_with_non_send_sync)]
76        let connection = Connection {
77            inner: Some(conn),
78            transaction_behavior: TransactionBehavior::Deferred,
79            dangling_tx: AtomicDropBehavior::new(DropBehavior::Ignore),
80            extra_io,
81        };
82        connection
83    }
84
85    pub(crate) async fn maybe_handle_dangling_tx(&self) -> Result<()> {
86        match self.dangling_tx.load(Ordering::SeqCst) {
87            DropBehavior::Rollback => {
88                let mut stmt = self.prepare("ROLLBACK").await?;
89                stmt.execute(()).await?;
90                self.dangling_tx
91                    .store(DropBehavior::Ignore, Ordering::SeqCst);
92            }
93            DropBehavior::Commit => {
94                let mut stmt = self.prepare("COMMIT").await?;
95                stmt.execute(()).await?;
96                self.dangling_tx
97                    .store(DropBehavior::Ignore, Ordering::SeqCst);
98            }
99            DropBehavior::Ignore => {}
100            DropBehavior::Panic => {
101                panic!("Transaction dropped unexpectedly.");
102            }
103        }
104        Ok(())
105    }
106
107    /// Query the database with SQL.
108    pub async fn query(&self, sql: impl AsRef<str>, params: impl IntoParams) -> Result<Rows> {
109        self.maybe_handle_dangling_tx().await?;
110        let mut stmt = self.prepare(sql).await?;
111        stmt.query(params).await
112    }
113
114    /// Execute SQL statement on the database.
115    pub async fn execute(&self, sql: impl AsRef<str>, params: impl IntoParams) -> Result<u64> {
116        self.maybe_handle_dangling_tx().await?;
117        let mut stmt = self.prepare(sql).await?;
118        stmt.execute(params).await
119    }
120
121    /// get the inner connection
122    fn get_inner_connection(&self) -> Result<Arc<crate::turso_sdk_kit::rsapi::TursoConnection>> {
123        match &self.inner {
124            Some(inner) => Ok(inner.clone()),
125            None => Err(Error::Misuse("inner connection must be set".to_string())),
126        }
127    }
128
129    /// Execute a batch of SQL statements on the database.
130    pub async fn execute_batch(&self, sql: impl AsRef<str>) -> Result<()> {
131        self.maybe_handle_dangling_tx().await?;
132        self.prepare_execute_batch(sql).await?;
133        Ok(())
134    }
135
136    /// Prepare a SQL statement for later execution.
137    pub async fn prepare(&self, sql: impl AsRef<str>) -> Result<Statement> {
138        let conn = self.get_inner_connection()?;
139        let stmt = conn.prepare_single(sql)?;
140
141        #[allow(clippy::arc_with_non_send_sync)]
142        let statement = Statement {
143            conn: self.clone(),
144            inner: Arc::new(Mutex::new(stmt)),
145        };
146        Ok(statement)
147    }
148
149    /// Prepare a SQL statement for later execution, caching it in the connection.
150    pub async fn prepare_cached(&self, sql: impl AsRef<str>) -> Result<Statement> {
151        let conn = self.get_inner_connection()?;
152        let stmt = conn.prepare_cached(sql)?;
153
154        #[allow(clippy::arc_with_non_send_sync)]
155        let statement = Statement {
156            conn: self.clone(),
157            inner: Arc::new(Mutex::new(stmt)),
158        };
159        Ok(statement)
160    }
161
162    async fn prepare_execute_batch(&self, sql: impl AsRef<str>) -> Result<()> {
163        self.maybe_handle_dangling_tx().await?;
164        let conn = self.get_inner_connection()?;
165        let mut sql = sql.as_ref();
166        while let Some((stmt, offset)) = conn.prepare_first(sql)? {
167            let mut stmt = Statement {
168                conn: self.clone(),
169                inner: Arc::new(Mutex::new(stmt)),
170            };
171            let _ = stmt.execute(()).await?;
172            sql = &sql[offset..];
173        }
174        Ok(())
175    }
176
177    /// Query a pragma.
178    pub async fn pragma_query<F>(&self, pragma_name: &str, mut f: F) -> Result<()>
179    where
180        F: FnMut(&Row) -> std::result::Result<(), crate::turso_sdk_kit::rsapi::TursoError>,
181    {
182        let sql = format!("PRAGMA {pragma_name}");
183        let mut stmt = self.prepare(&sql).await?;
184        let mut rows = stmt.query(()).await?;
185        while let Some(row) = rows.next().await? {
186            f(&row)?;
187        }
188        Ok(())
189    }
190
191    /// Set a pragma value.
192    pub async fn pragma_update<V: std::fmt::Display>(
193        &self,
194        pragma_name: &str,
195        pragma_value: V,
196    ) -> Result<Vec<Row>> {
197        let sql = format!("PRAGMA {pragma_name} = {pragma_value}");
198        let mut stmt = self.prepare(&sql).await?;
199        let mut rows = stmt.query(()).await?;
200        let mut collected = Vec::new();
201        while let Some(row) = rows.next().await? {
202            collected.push(row);
203        }
204        Ok(collected)
205    }
206
207    /// Returns the rowid of the last row inserted.
208    pub fn last_insert_rowid(&self) -> i64 {
209        let conn = self.get_inner_connection().unwrap();
210        conn.last_insert_rowid()
211    }
212
213    /// Flush dirty pages to disk.
214    /// This will write the dirty pages to the WAL.
215    pub fn cacheflush(&self) -> Result<()> {
216        let conn = self.get_inner_connection()?;
217        conn.cacheflush()?;
218        Ok(())
219    }
220
221    pub fn is_autocommit(&self) -> Result<bool> {
222        let conn = self.get_inner_connection()?;
223        Ok(conn.get_auto_commit())
224    }
225
226    /// Sets maximum total accumuated timeout. If the duration is None or Zero, we unset the busy handler for this Connection
227    ///
228    /// This api defers slighty from: https://www.sqlite.org/c3ref/busy_timeout.html
229    ///
230    /// Instead of sleeping for linear amount of time specified by the user,
231    /// we will sleep in phases, until the the total amount of time is reached.
232    /// This means we first sleep of 1ms, then if we still return busy, we sleep for 2 ms, and repeat until a maximum of 100 ms per phase.
233    ///
234    /// Example:
235    /// 1. Set duration to 5ms
236    /// 2. Step through query -> returns Busy -> sleep/yield for 1 ms
237    /// 3. Step through query -> returns Busy -> sleep/yield for 2 ms
238    /// 4. Step through query -> returns Busy -> sleep/yield for 2 ms (totaling 5 ms of sleep)
239    /// 5. Step through query -> returns Busy -> return Busy to user
240    pub fn busy_timeout(&self, duration: std::time::Duration) -> Result<()> {
241        let conn = self.get_inner_connection()?;
242        conn.set_busy_timeout(duration);
243        Ok(())
244    }
245}
246
247impl Debug for Connection {
248    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
249        f.debug_struct("Connection").finish()
250    }
251}
clt_database/turso/src/connection.rs

clt_database/turso/src/
connection.rs
