diesel/sqlite/connection/
borrowed_connection.rs

1use super::{RawConnection, SqliteConnection};
2use crate::connection::instrumentation::DynInstrumentation;
3use crate::connection::statement_cache::StatementCache;
4use crate::connection::*;
5use crate::expression::QueryMetadata;
6use crate::prelude::{ConnectionResult, QueryResult};
7use crate::query_builder::{Query, QueryFragment, QueryId};
8use crate::sqlite::Sqlite;
9use core::mem::ManuallyDrop;
10
11/// A borrowed SQLite connection
12///
13/// This type is provided to various hook functions called by SQLite internally
14/// It mostly acts as `&mut SqliteConnection` with the additional constraint
15/// that you cannot move out/replace the inner connection handle
16///
17/// Any hook (update, commit, rollback, progress, wal, busy handler, authorizer, trace, collation needed)
18/// is reset after the closure that gets this connection type as argument ends. This specifically means
19/// that any hook registered during [`register_auto_extension`] won't trigger later
20#[allow(missing_debug_implementations)]
21pub struct BorrowedSqliteConnection {
22    // SAFETY invariant: never move the inner connection
23    // out of the container. we rely on the fact that we never
24    // run the destructor for the connection as it's only
25    // borrowed from an outer connection and not actually
26    // owned
27    inner: ManuallyDrop<SqliteConnection>,
28}
29
30impl Drop for BorrowedSqliteConnection {
31    #[allow(unsafe_code)] // manual memory management
32    fn drop(&mut self) {
33        // SAFETY: `self.inner` is not touched again after this take.
34        let conn = unsafe { core::mem::ManuallyDrop::take(&mut self.inner) };
35        let SqliteConnection {
36            statement_cache,
37            raw_connection,
38            transaction_state: _,
39            metadata_lookup: (),
40            instrumentation,
41        } = conn;
42        // wrap the inner connection into manually drop again
43        // to make sure that we don't call the destructor under any condition (even panics)
44        let mut raw_connection = core::mem::ManuallyDrop::new(raw_connection);
45        // we need to make sure to unregister all the hooks here to not leak memory
46        raw_connection.reset_registered_hooks();
47        // Finalize prepared statements, but do not run `RawConnection`'s
48        // `Drop`, which would close a handle we do not own.
49        drop(statement_cache);
50        drop(instrumentation);
51    }
52}
53
54impl SimpleConnection for BorrowedSqliteConnection {
55    fn batch_execute(&mut self, query: &str) -> QueryResult<()> {
56        SqliteConnection::batch_execute(&mut self.inner, query)
57    }
58}
59
60impl Connection for BorrowedSqliteConnection {
61    type Backend = Sqlite;
62
63    type TransactionManager = AnsiTransactionManager;
64
65    fn establish(_database_url: &str) -> ConnectionResult<Self> {
66        Err(crate::ConnectionError::BadConnection(
67            "Cannot create a pinned sqlite connection".into(),
68        ))
69    }
70
71    fn execute_returning_count<T>(&mut self, source: &T) -> QueryResult<usize>
72    where
73        T: QueryFragment<Self::Backend> + QueryId,
74    {
75        SqliteConnection::execute_returning_count(&mut self.inner, source)
76    }
77
78    fn transaction_state(
79        &mut self,
80    ) -> &mut <Self::TransactionManager as TransactionManager<Self>>::TransactionStateData {
81        SqliteConnection::transaction_state(&mut self.inner)
82    }
83
84    fn instrumentation(&mut self) -> &mut dyn Instrumentation {
85        SqliteConnection::instrumentation(&mut self.inner)
86    }
87
88    fn set_instrumentation(&mut self, instrumentation: impl Instrumentation) {
89        SqliteConnection::set_instrumentation(&mut self.inner, instrumentation);
90    }
91
92    fn set_prepared_statement_cache_size(&mut self, size: CacheSize) {
93        SqliteConnection::set_prepared_statement_cache_size(&mut self.inner, size)
94    }
95}
96
97impl LoadConnection for BorrowedSqliteConnection {
98    type Cursor<'conn, 'query>
99        = <SqliteConnection as LoadConnection>::Cursor<'conn, 'query>
100    where
101        Self: 'conn;
102
103    type Row<'conn, 'query>
104        = <SqliteConnection as LoadConnection>::Row<'conn, 'query>
105    where
106        Self: 'conn;
107
108    fn load<'conn, 'query, T>(
109        &'conn mut self,
110        source: T,
111    ) -> QueryResult<Self::Cursor<'conn, 'query>>
112    where
113        T: Query + QueryFragment<Self::Backend> + QueryId + 'query,
114        Self::Backend: QueryMetadata<T::SqlType>,
115    {
116        SqliteConnection::load(&mut self.inner, source)
117    }
118}
119
120impl ConnectionSealed for BorrowedSqliteConnection {}
121
122impl BorrowedSqliteConnection {
123    pub(super) fn new(conn: RawConnection) -> Self {
124        BorrowedSqliteConnection {
125            inner: core::mem::ManuallyDrop::new(SqliteConnection {
126                statement_cache: StatementCache::new(),
127                raw_connection: conn,
128                transaction_state: AnsiTransactionManager::default(),
129                metadata_lookup: (),
130                instrumentation: DynInstrumentation::default_instrumentation(),
131            }),
132        }
133    }
134
135    /// Get access to the underlying SqliteConnection
136    ///
137    /// # Safety
138    ///
139    /// - You are not allowed to move the connection out of the provided reference or otherwise
140    ///   replace the connection
141    #[allow(unsafe_code)] // helper to allow others accessing the connection
142    pub unsafe fn inner_mut(&mut self) -> &mut SqliteConnection {
143        &mut self.inner
144    }
145}
146
147// Dropping the borrowed connection must not unregister hooks of the outer
148// connection that shares the same `sqlite3` handle.
149//
150// These tests rely on sqlite calling the registered hooks, which is not
151// supported when running under miri with a native libsqlite3.
152#[cfg(all(
153    test,
154    not(any(miri, all(target_family = "wasm", target_os = "unknown")))
155))]
156mod tests {
157    use super::*;
158    use crate::connection::SimpleConnection;
159    use std::sync::Arc;
160    use std::sync::atomic::{AtomicU32, Ordering};
161
162    // `tempfile::tempdir()` panics on WASM due to the lack of a filesystem,
163    // and WAL mode does not work with `:memory:`.
164    #[cfg(not(all(target_family = "wasm", target_os = "unknown")))]
165    #[diesel_test_helper::test]
166    fn wal_hook_survives_its_own_call() {
167        let dir = tempfile::tempdir().unwrap();
168        let path = dir.path().join("test.db");
169        let conn = &mut SqliteConnection::establish(path.to_str().unwrap()).unwrap();
170        conn.batch_execute("PRAGMA journal_mode = WAL; CREATE TABLE borrowed_wal (x)")
171            .unwrap();
172
173        let calls = Arc::new(AtomicU32::new(0));
174        let c2 = calls.clone();
175        conn.on_wal(move |_conn, _db_name, _n_pages| {
176            c2.fetch_add(1, Ordering::Relaxed);
177        });
178
179        conn.batch_execute("INSERT INTO borrowed_wal VALUES (1)")
180            .unwrap();
181        conn.batch_execute("INSERT INTO borrowed_wal VALUES (2)")
182            .unwrap();
183        assert_eq!(calls.load(Ordering::Relaxed), 2);
184    }
185
186    // Hooks registered inside a WAL callback are freed when it returns, also
187    // when they fire from a nested borrowed connection inside that callback.
188    #[cfg(not(all(target_family = "wasm", target_os = "unknown")))]
189    #[diesel_test_helper::test]
190    #[allow(unsafe_code)]
191    fn hooks_registered_in_wal_callback_are_freed() {
192        use crate::sqlite::CommitDecision;
193        struct Token(Arc<AtomicU32>);
194        impl Drop for Token {
195            fn drop(&mut self) {
196                self.0.fetch_add(1, Ordering::SeqCst);
197            }
198        }
199        let dir = tempfile::tempdir().unwrap();
200        let path = dir.path().join("test.db");
201        let conn = &mut SqliteConnection::establish(path.to_str().unwrap()).unwrap();
202        conn.batch_execute("PRAGMA journal_mode = WAL; CREATE TABLE borrowed_wal_hooks (x)")
203            .unwrap();
204
205        let dropped = Arc::new(AtomicU32::new(0));
206        let inner_wal_calls = Arc::new(AtomicU32::new(0));
207        let (d, w) = (dropped.clone(), inner_wal_calls.clone());
208        conn.on_wal(move |conn, _db_name, _n_pages| {
209            let conn = unsafe { conn.inner_mut() };
210            let t = Token(d.clone());
211            conn.on_commit(move || {
212                let _ = &t;
213                CommitDecision::Proceed
214            });
215            let (t, w) = (Token(d.clone()), w.clone());
216            // Replaces the outer WAL hook while it runs; fires from a nested borrowed connection.
217            conn.on_wal(move |_conn, _, _| {
218                let _ = &t;
219                w.fetch_add(1, Ordering::SeqCst);
220            });
221            conn.batch_execute("INSERT INTO borrowed_wal_hooks VALUES (0)")
222                .unwrap();
223        });
224
225        conn.batch_execute("INSERT INTO borrowed_wal_hooks VALUES (1)")
226            .unwrap();
227        assert_eq!(inner_wal_calls.load(Ordering::SeqCst), 1);
228        assert_eq!(dropped.load(Ordering::SeqCst), 2);
229    }
230}
diesel/sqlite/connection/borrowed_connection.rs

diesel/sqlite/connection/
borrowed_connection.rs
