Skip to main content

sql_peas/
connection.rs

1use core::ffi::{c_char, c_int, c_void};
2use std::collections::HashMap;
3use std::marker::PhantomData;
4use std::ops::{Deref, DerefMut};
5use std::path::Path;
6
7use crate::error::Result;
8use crate::statement::{Handle as StatementHandle, Statement};
9
10/// A connection.
11pub struct Connection {
12    raw: Raw,
13    busy_callback: Option<Box<dyn FnMut(usize) -> bool + Send>>,
14    phantom: PhantomData<ffi::sqlite3>,
15    statements: HashMap<StatementHandle, crate::Statement>,
16}
17
18/// A thread-safe connection.
19pub struct ConnectionThreadSafe(Connection);
20
21/// Flags for opening a connection.
22#[derive(Clone, Copy, Debug)]
23pub struct OpenFlags(c_int);
24
25struct Raw(*mut ffi::sqlite3);
26
27impl Connection {
28    /// Open a read-write connection to a new or existing database.
29    pub fn open<T: AsRef<Path>>(path: T) -> Result<Connection> {
30        Connection::open_with_flags(path, OpenFlags::new().with_create().with_read_write())
31    }
32
33    /// Open a connection with specific flags.
34    pub fn open_with_flags<T: AsRef<Path>>(path: T, flags: OpenFlags) -> Result<Connection> {
35        let mut raw = std::ptr::null_mut();
36        unsafe {
37            let code = ffi::sqlite3_open_v2(
38                path_to_cstr!(path.as_ref()).as_ptr(),
39                &mut raw,
40                flags.0,
41                std::ptr::null(),
42            );
43            match code {
44                ffi::SQLITE_OK => {}
45                code => match crate::error::last(raw) {
46                    Some(error) => {
47                        ffi::sqlite3_close(raw);
48                        return Err(error);
49                    }
50                    _ => {
51                        ffi::sqlite3_close(raw);
52                        return Err(crate::error::Error {
53                            code: Some(code as isize),
54                            message: None,
55                        });
56                    }
57                },
58            }
59        }
60        Ok(Connection {
61            raw: Raw(raw),
62            busy_callback: None,
63            phantom: PhantomData,
64            statements: Default::default(),
65        })
66    }
67
68    /// Open a thread-safe read-write connection to a new or existing database.
69    pub fn open_thread_safe<T: AsRef<Path>>(path: T) -> Result<ConnectionThreadSafe> {
70        Connection::open_with_flags(
71            path,
72            OpenFlags::new()
73                .with_create()
74                .with_read_write()
75                .with_full_mutex(),
76        )
77        .map(ConnectionThreadSafe)
78    }
79
80    /// Open a thread-safe connection with specific flags.
81    pub fn open_thread_safe_with_flags<T: AsRef<Path>>(
82        path: T,
83        flags: OpenFlags,
84    ) -> Result<ConnectionThreadSafe> {
85        Connection::open_with_flags(path, flags.with_full_mutex()).map(ConnectionThreadSafe)
86    }
87
88    #[doc(hidden)]
89    #[inline]
90    pub fn as_raw(&self) -> *mut ffi::sqlite3 {
91        self.raw.0
92    }
93}
94
95impl Connection {
96    pub fn begin_immediate_transaction(&self) -> Result<Transaction<'_>> {
97        self.execute("BEGIN IMMEDIATE")?;
98        Ok(Transaction {
99            connection: self,
100            tried_commit: false,
101        })
102    }
103
104    /// Use a `StatementHandle` to borrow a prepared statement
105    pub fn borrow_statement(&self, handle: crate::statement::Handle) -> Result<&Statement> {
106        let Some(stmt) = self.statements.get(&handle) else {
107            return Err(crate::Error {
108                code: None,
109                message: Some(
110                    "Tried to borrow a statement Handle with the wrong Connection".into(),
111                ),
112            });
113        };
114        stmt.reset()?;
115        Ok(stmt)
116    }
117
118    /// Drop a prepared statement, finalizing it and freeing its resources
119    pub fn drop_statement(&mut self, handle: crate::statement::Handle) -> Result<()> {
120        self.statements
121            .remove(&handle)
122            .map(|_| ())
123            .ok_or_else(|| crate::Error {
124                code: None,
125                message: Some("Tried to drop a statement Handle with the wrong Connection".into()),
126            })
127    }
128
129    /// Execute a statement without processing the resulting rows if any.
130    #[inline]
131    pub fn execute<T: AsRef<str>>(&self, statement: T) -> Result<()> {
132        unsafe {
133            ok!(
134                self.raw.0,
135                ffi::sqlite3_exec(
136                    self.raw.0,
137                    str_to_cstr!(statement.as_ref()).as_ptr(),
138                    None,
139                    std::ptr::null_mut(),
140                    std::ptr::null_mut(),
141                )
142            );
143        }
144        Ok(())
145    }
146
147    /// Execute a statement and process the resulting rows as plain text.
148    ///
149    /// The callback is triggered for each row. If the callback returns `false`, no more rows will
150    /// be processed. For large queries and non-string data types, prepared statement are highly
151    /// preferable; see `prepare`.
152    #[inline]
153    pub fn iterate<T: AsRef<str>, F>(&self, statement: T, callback: F) -> Result<()>
154    where
155        F: FnMut(&[(&str, Option<&str>)]) -> bool,
156    {
157        unsafe {
158            let callback = Box::new(callback);
159            ok!(
160                self.raw.0,
161                ffi::sqlite3_exec(
162                    self.raw.0,
163                    str_to_cstr!(statement.as_ref()).as_ptr(),
164                    Some(process_callback::<F>),
165                    &*callback as *const F as *mut F as *mut _,
166                    std::ptr::null_mut(),
167                )
168            );
169        }
170        Ok(())
171    }
172
173    /// Create a prepared statement.
174    #[inline]
175    pub fn prepare<T: AsRef<str>>(&mut self, statement: T) -> Result<StatementHandle> {
176        let stmt = crate::statement::new(self.raw.0, statement)?;
177        let handle = crate::statement::Handle(stmt.raw.0);
178        self.statements.insert(handle, stmt);
179        Ok(handle)
180    }
181
182    /// Return the number of rows inserted, updated, or deleted by the most recent INSERT, UPDATE,
183    /// or DELETE statement.
184    #[inline]
185    pub fn change_count(&self) -> usize {
186        unsafe { ffi::sqlite3_changes(self.raw.0) as usize }
187    }
188
189    /// Return the total number of rows inserted, updated, and deleted by all INSERT, UPDATE, and
190    /// DELETE statements since the connection was opened.
191    #[inline]
192    pub fn total_change_count(&self) -> usize {
193        unsafe { ffi::sqlite3_total_changes(self.raw.0) as usize }
194    }
195}
196
197impl Connection {
198    /// Set a callback for handling busy events.
199    ///
200    /// The callback is triggered when the database cannot perform an operation due to processing
201    /// of some other request. If the callback returns `true`, the operation will be repeated.
202    pub fn set_busy_handler<F>(&mut self, callback: F) -> Result<()>
203    where
204        F: FnMut(usize) -> bool + Send + 'static,
205    {
206        self.remove_busy_handler()?;
207        unsafe {
208            let callback = Box::new(callback);
209            let result = ffi::sqlite3_busy_handler(
210                self.raw.0,
211                Some(busy_callback::<F>),
212                &*callback as *const F as *mut F as *mut _,
213            );
214            self.busy_callback = Some(callback);
215            ok!(self.raw.0, result);
216        }
217        Ok(())
218    }
219
220    /// Set an implicit callback for handling busy events that tries to repeat rejected operations
221    /// until a timeout expires.
222    #[inline]
223    pub fn set_busy_timeout(&mut self, milliseconds: usize) -> Result<()> {
224        unsafe {
225            ok!(
226                self.raw.0,
227                ffi::sqlite3_busy_timeout(self.raw.0, milliseconds as c_int)
228            );
229        }
230        Ok(())
231    }
232
233    /// Remove the callback handling busy events.
234    #[inline]
235    pub fn remove_busy_handler(&mut self) -> Result<()> {
236        self.busy_callback = None;
237        unsafe {
238            ok!(
239                self.raw.0,
240                ffi::sqlite3_busy_handler(self.raw.0, None, std::ptr::null_mut())
241            );
242        }
243        Ok(())
244    }
245}
246
247impl Connection {
248    /// Enable loading extensions.
249    #[cfg(feature = "extension")]
250    #[inline]
251    pub fn enable_extension(&self) -> Result<()> {
252        unsafe {
253            ok!(
254                self.raw.0,
255                ffi::sqlite3_enable_load_extension(self.raw.0, 1 as c_int)
256            );
257        }
258        Ok(())
259    }
260
261    /// Disable loading extensions.
262    #[cfg(feature = "extension")]
263    #[inline]
264    pub fn disable_extension(&self) -> Result<()> {
265        unsafe {
266            ok!(
267                self.raw.0,
268                ffi::sqlite3_enable_load_extension(self.raw.0, 0 as c_int)
269            );
270        }
271        Ok(())
272    }
273
274    /// Load an extension.
275    #[cfg(feature = "extension")]
276    #[inline]
277    pub fn load_extension<T: AsRef<str>>(&self, name: T) -> Result<()> {
278        unsafe {
279            ok!(
280                self.raw.0,
281                ffi::sqlite3_load_extension(
282                    self.raw.0,
283                    str_to_cstr!(name.as_ref()).as_ptr() as *const c_char,
284                    std::ptr::null_mut(),
285                    std::ptr::null_mut(),
286                )
287            );
288        }
289        Ok(())
290    }
291}
292
293impl Connection {
294    /// Set the encryption key.
295    #[cfg(feature = "encryption")]
296    #[inline]
297    pub fn set_encryption_key<T: AsRef<str>>(&self, key: T) -> Result<()> {
298        unsafe {
299            ok!(
300                self.raw.0,
301                ffi::sqlite3_key_v2(
302                    self.raw.0,
303                    std::ptr::null() as *const c_char,
304                    str_to_cstr!(key.as_ref()).as_ptr() as *const c_void,
305                    key.as_ref().len() as c_int,
306                )
307            );
308        }
309        Ok(())
310    }
311
312    /// Change the encryption key.
313    #[cfg(feature = "encryption")]
314    #[inline]
315    pub fn change_encryption_key<T: AsRef<str>>(&self, new_key: T) -> Result<()> {
316        unsafe {
317            ok!(
318                self.raw.0,
319                ffi::sqlite3_rekey_v2(
320                    self.raw.0,
321                    std::ptr::null() as *const c_char,
322                    str_to_cstr!(new_key.as_ref()).as_ptr() as *const c_void,
323                    new_key.as_ref().len() as c_int,
324                )
325            );
326        }
327        Ok(())
328    }
329}
330
331impl Drop for Connection {
332    #[inline]
333    #[allow(unused_must_use)]
334    fn drop(&mut self) {
335        self.remove_busy_handler();
336        unsafe { ffi::sqlite3_close(self.raw.0) };
337    }
338}
339
340impl OpenFlags {
341    /// Create flags for opening a database connection.
342    #[inline]
343    pub fn new() -> Self {
344        OpenFlags(0)
345    }
346
347    /// Create the database if it does not already exist.
348    pub fn with_create(mut self) -> Self {
349        self.0 |= ffi::SQLITE_OPEN_CREATE;
350        self
351    }
352
353    /// Open the database in the serialized [threading mode][1].
354    ///
355    /// [1]: https://www.sqlite.org/threadsafe.html
356    pub fn with_full_mutex(mut self) -> Self {
357        self.0 |= ffi::SQLITE_OPEN_FULLMUTEX;
358        self
359    }
360
361    /// Opens the database in the multi-thread [threading mode][1].
362    ///
363    /// [1]: https://www.sqlite.org/threadsafe.html
364    pub fn with_no_mutex(mut self) -> Self {
365        self.0 |= ffi::SQLITE_OPEN_NOMUTEX;
366        self
367    }
368
369    /// Open the database for reading only.
370    pub fn with_read_only(mut self) -> Self {
371        self.0 |= ffi::SQLITE_OPEN_READONLY;
372        self
373    }
374
375    /// Open the database for reading and writing.
376    pub fn with_read_write(mut self) -> Self {
377        self.0 |= ffi::SQLITE_OPEN_READWRITE;
378        self
379    }
380
381    /// Allow the path to be interpreted as a URI.
382    pub fn with_uri(mut self) -> Self {
383        self.0 |= ffi::SQLITE_OPEN_URI;
384        self
385    }
386}
387
388impl Default for OpenFlags {
389    #[inline]
390    fn default() -> Self {
391        Self::new()
392    }
393}
394
395impl Deref for ConnectionThreadSafe {
396    type Target = Connection;
397
398    #[inline]
399    fn deref(&self) -> &Self::Target {
400        &self.0
401    }
402}
403
404impl DerefMut for ConnectionThreadSafe {
405    #[inline]
406    fn deref_mut(&mut self) -> &mut Self::Target {
407        &mut self.0
408    }
409}
410
411unsafe impl Sync for ConnectionThreadSafe {}
412
413unsafe impl Send for Raw {}
414
415extern "C" fn busy_callback<F>(callback: *mut c_void, attempts: c_int) -> c_int
416where
417    F: FnMut(usize) -> bool,
418{
419    unsafe { c_int::from((*(callback as *mut F))(attempts as usize)) }
420}
421
422extern "C" fn process_callback<F>(
423    callback: *mut c_void,
424    count: c_int,
425    values: *mut *mut c_char,
426    columns: *mut *mut c_char,
427) -> c_int
428where
429    F: FnMut(&[(&str, Option<&str>)]) -> bool,
430{
431    unsafe {
432        let mut pairs = Vec::with_capacity(count as usize);
433        for index in 0..(count as isize) {
434            let column = {
435                let pointer = *columns.offset(index);
436                debug_assert!(!pointer.is_null());
437                c_str_to_str!(pointer).unwrap()
438            };
439            let value = {
440                let pointer = *values.offset(index);
441                if pointer.is_null() {
442                    None
443                } else {
444                    Some(c_str_to_str!(pointer).unwrap())
445                }
446            };
447            pairs.push((column, value));
448        }
449        c_int::from(!(*(callback as *mut F))(&pairs))
450    }
451}
452
453pub struct Transaction<'a> {
454    connection: &'a Connection,
455    tried_commit: bool,
456}
457
458impl<'a> Transaction<'a> {
459    pub fn commit(mut self) -> Result<()> {
460        self.tried_commit = true;
461        self.connection.execute("COMMIT")
462    }
463}
464
465impl<'a> Drop for Transaction<'a> {
466    fn drop(&mut self) {
467        if !self.tried_commit {
468            self.connection.execute("ROLLBACK TRANSACTION").unwrap();
469        }
470    }
471}