diesel/sqlite/connection/
mod.rs

1#[cfg(not(all(target_family = "wasm", target_os = "unknown")))]
2extern crate libsqlite3_sys as ffi;
3
4#[cfg(all(target_family = "wasm", target_os = "unknown"))]
5use sqlite_wasm_rs as ffi;
6
7mod bind_collector;
8mod functions;
9#[cfg(all(test, not(all(target_family = "wasm", target_os = "unknown"))))]
10#[allow(unsafe_code)]
11mod oom_test_support;
12mod owned_row;
13mod raw;
14mod row;
15mod serialized_database;
16mod sqlite_value;
17mod statement_iterator;
18mod stmt;
19
20pub(in crate::sqlite) use self::bind_collector::SqliteBindCollector;
21pub use self::bind_collector::SqliteBindValue;
22pub use self::serialized_database::SerializedDatabase;
23pub use self::sqlite_value::SqliteValue;
24
25use std::os::raw as libc;
26
27use self::raw::RawConnection;
28use self::statement_iterator::*;
29use self::stmt::{Statement, StatementUse};
30use super::SqliteAggregateFunction;
31use crate::connection::instrumentation::{DynInstrumentation, StrQueryHelper};
32use crate::connection::statement_cache::StatementCache;
33use crate::connection::*;
34use crate::deserialize::{FromSqlRow, StaticallySizedRow};
35use crate::expression::QueryMetadata;
36use crate::query_builder::*;
37use crate::result::*;
38use crate::serialize::ToSql;
39use crate::sql_types::{HasSqlType, TypeMetadata};
40use crate::sqlite::Sqlite;
41
42/// Connections for the SQLite backend. Unlike other backends, SQLite supported
43/// connection URLs are:
44///
45/// - File paths (`test.db`)
46/// - [URIs](https://sqlite.org/uri.html) (`file://test.db`)
47/// - Special identifiers (`:memory:`)
48///
49/// # Supported loading model implementations
50///
51/// * [`DefaultLoadingMode`]
52///
53/// As `SqliteConnection` only supports a single loading mode implementation,
54/// it is **not required** to explicitly specify a loading mode
55/// when calling [`RunQueryDsl::load_iter()`] or [`LoadConnection::load`]
56///
57/// [`RunQueryDsl::load_iter()`]: crate::query_dsl::RunQueryDsl::load_iter
58///
59/// ## DefaultLoadingMode
60///
61/// `SqliteConnection` only supports a single loading mode, which loads
62/// values row by row from the result set.
63///
64/// ```rust
65/// # include!("../../doctest_setup.rs");
66/// #
67/// # fn main() {
68/// #     run_test().unwrap();
69/// # }
70/// #
71/// # fn run_test() -> QueryResult<()> {
72/// #     use schema::users;
73/// #     let connection = &mut establish_connection();
74/// use diesel::connection::DefaultLoadingMode;
75/// {
76///     // scope to restrict the lifetime of the iterator
77///     let iter1 = users::table.load_iter::<(i32, String), DefaultLoadingMode>(connection)?;
78///
79///     for r in iter1 {
80///         let (id, name) = r?;
81///         println!("Id: {} Name: {}", id, name);
82///     }
83/// }
84///
85/// // works without specifying the loading mode
86/// let iter2 = users::table.load_iter::<(i32, String), _>(connection)?;
87///
88/// for r in iter2 {
89///     let (id, name) = r?;
90///     println!("Id: {} Name: {}", id, name);
91/// }
92/// #   Ok(())
93/// # }
94/// ```
95///
96/// This mode does **not support** creating
97/// multiple iterators using the same connection.
98///
99/// ```compile_fail
100/// # include!("../../doctest_setup.rs");
101/// #
102/// # fn main() {
103/// #     run_test().unwrap();
104/// # }
105/// #
106/// # fn run_test() -> QueryResult<()> {
107/// #     use schema::users;
108/// #     let connection = &mut establish_connection();
109/// use diesel::connection::DefaultLoadingMode;
110///
111/// let iter1 = users::table.load_iter::<(i32, String), DefaultLoadingMode>(connection)?;
112/// let iter2 = users::table.load_iter::<(i32, String), DefaultLoadingMode>(connection)?;
113///
114/// for r in iter1 {
115///     let (id, name) = r?;
116///     println!("Id: {} Name: {}", id, name);
117/// }
118///
119/// for r in iter2 {
120///     let (id, name) = r?;
121///     println!("Id: {} Name: {}", id, name);
122/// }
123/// #   Ok(())
124/// # }
125/// ```
126///
127/// # Concurrency
128///
129/// By default, when running into a database lock, the operation will abort with a
130/// `Database locked` error. However, it's possible to configure it for greater concurrency,
131/// trading latency for not having to deal with retries yourself.
132///
133/// You can use this example as blue-print for which statements to run after establishing a connection.
134/// It is **important** to run each `PRAGMA` in a single statement to make sure all of them apply
135/// correctly. In addition the order of the `PRAGMA` statements is relevant to prevent timeout
136/// issues for the later `PRAGMA` statements.
137///
138/// ```rust
139/// # include!("../../doctest_setup.rs");
140/// #
141/// # fn main() {
142/// #     run_test().unwrap();
143/// # }
144/// #
145/// # fn run_test() -> QueryResult<()> {
146/// #     use schema::users;
147/// use diesel::connection::SimpleConnection;
148/// let conn = &mut establish_connection();
149/// // see https://fractaledmind.github.io/2023/09/07/enhancing-rails-sqlite-fine-tuning/
150/// // sleep if the database is busy, this corresponds to up to 2 seconds sleeping time.
151/// conn.batch_execute("PRAGMA busy_timeout = 2000;")?;
152/// // better write-concurrency
153/// conn.batch_execute("PRAGMA journal_mode = WAL;")?;
154/// // fsync only in critical moments
155/// conn.batch_execute("PRAGMA synchronous = NORMAL;")?;
156/// // write WAL changes back every 1000 pages, for an in average 1MB WAL file.
157/// // May affect readers if number is increased
158/// conn.batch_execute("PRAGMA wal_autocheckpoint = 1000;")?;
159/// // free some space by truncating possibly massive WAL files from the last run
160/// conn.batch_execute("PRAGMA wal_checkpoint(TRUNCATE);")?;
161/// #   Ok(())
162/// # }
163/// ```
164#[allow(missing_debug_implementations)]
165#[cfg(feature = "sqlite")]
166pub struct SqliteConnection {
167    // statement_cache needs to be before raw_connection
168    // otherwise we will get errors about open statements before closing the
169    // connection itself
170    statement_cache: StatementCache<Sqlite, Statement>,
171    raw_connection: RawConnection,
172    transaction_state: AnsiTransactionManager,
173    // this exists for the sole purpose of implementing `WithMetadataLookup` trait
174    // and avoiding static mut which will be deprecated in 2024 edition
175    metadata_lookup: (),
176    instrumentation: DynInstrumentation,
177    // We potentially need to store a serialized
178    // database in here to make sure the database bytes
179    // live as long as the connection
180    // This is used by SqliteConnection::deserialize_readonly_database_from_buffer
181    // only
182    // This field needs to come after the RawConnection
183    // as we need to make sure the data are still there until the
184    // connection is dropped
185    //
186    // We are not allowed to modify the inner buffer until the database connection is dropped
187    serialized_data: Vec<Vec<u8>>,
188}
189
190// This relies on the invariant that RawConnection or Statement are never
191// leaked. If a reference to one of those was held on a different thread, this
192// would not be thread safe.
193#[allow(unsafe_code)]
194unsafe impl Send for SqliteConnection {}
195
196impl SimpleConnection for SqliteConnection {
197    fn batch_execute(&mut self, query: &str) -> QueryResult<()> {
198        self.instrumentation
199            .on_connection_event(InstrumentationEvent::StartQuery {
200                query: &StrQueryHelper::new(query),
201            });
202        let resp = self.raw_connection.exec(query);
203        self.instrumentation
204            .on_connection_event(InstrumentationEvent::FinishQuery {
205                query: &StrQueryHelper::new(query),
206                error: resp.as_ref().err(),
207            });
208        resp
209    }
210}
211
212impl ConnectionSealed for SqliteConnection {}
213
214impl Connection for SqliteConnection {
215    type Backend = Sqlite;
216    type TransactionManager = AnsiTransactionManager;
217
218    /// Establish a connection to the database specified by `database_url`.
219    ///
220    /// See [SqliteConnection] for supported `database_url`.
221    ///
222    /// If the database does not exist, this method will try to
223    /// create a new database and then establish a connection to it.
224    ///
225    /// ## WASM support
226    ///
227    /// If you plan to use this connection type on the `wasm32-unknown-unknown` target please
228    /// make sure to read the following notes:
229    ///
230    /// * The database is stored in memory by default.
231    /// * Persistent VFS (Virtual File Systems) is optional,
232    ///   see <https://github.com/Spxg/sqlite-wasm-rs> for details
233    fn establish(database_url: &str) -> ConnectionResult<Self> {
234        let mut instrumentation = DynInstrumentation::default_instrumentation();
235        instrumentation.on_connection_event(InstrumentationEvent::StartEstablishConnection {
236            url: database_url,
237        });
238
239        let establish_result = Self::establish_inner(database_url);
240        instrumentation.on_connection_event(InstrumentationEvent::FinishEstablishConnection {
241            url: database_url,
242            error: establish_result.as_ref().err(),
243        });
244        let mut conn = establish_result?;
245        conn.instrumentation = instrumentation;
246        Ok(conn)
247    }
248
249    fn execute_returning_count<T>(&mut self, source: &T) -> QueryResult<usize>
250    where
251        T: QueryFragment<Self::Backend> + QueryId,
252    {
253        let statement_use = self.prepared_query(source)?;
254        statement_use.run().and_then(|_| {
255            self.raw_connection
256                .rows_affected_by_last_query()
257                .map_err(Error::DeserializationError)
258        })
259    }
260
261    fn transaction_state(&mut self) -> &mut AnsiTransactionManager
262    where
263        Self: Sized,
264    {
265        &mut self.transaction_state
266    }
267
268    fn instrumentation(&mut self) -> &mut dyn Instrumentation {
269        &mut *self.instrumentation
270    }
271
272    fn set_instrumentation(&mut self, instrumentation: impl Instrumentation) {
273        self.instrumentation = instrumentation.into();
274    }
275
276    fn set_prepared_statement_cache_size(&mut self, size: CacheSize) {
277        self.statement_cache.set_cache_size(size);
278    }
279}
280
281impl LoadConnection<DefaultLoadingMode> for SqliteConnection {
282    type Cursor<'conn, 'query> = StatementIterator<'conn, 'query>;
283    type Row<'conn, 'query> = self::row::SqliteRow<'conn, 'query>;
284
285    fn load<'conn, 'query, T>(
286        &'conn mut self,
287        source: T,
288    ) -> QueryResult<Self::Cursor<'conn, 'query>>
289    where
290        T: Query + QueryFragment<Self::Backend> + QueryId + 'query,
291        Self::Backend: QueryMetadata<T::SqlType>,
292    {
293        let statement = self.prepared_query(source)?;
294
295        Ok(StatementIterator::new(statement))
296    }
297}
298
299impl WithMetadataLookup for SqliteConnection {
300    fn metadata_lookup(&mut self) -> &mut <Sqlite as TypeMetadata>::MetadataLookup {
301        &mut self.metadata_lookup
302    }
303}
304
305#[cfg(feature = "r2d2")]
306impl crate::r2d2::R2D2Connection for crate::sqlite::SqliteConnection {
307    fn ping(&mut self) -> QueryResult<()> {
308        use crate::RunQueryDsl;
309
310        crate::r2d2::CheckConnectionQuery.execute(self).map(|_| ())
311    }
312
313    fn is_broken(&mut self) -> bool {
314        AnsiTransactionManager::is_broken_transaction_manager(self)
315    }
316}
317
318impl MultiConnectionHelper for SqliteConnection {
319    fn to_any<'a>(
320        lookup: &mut <Self::Backend as crate::sql_types::TypeMetadata>::MetadataLookup,
321    ) -> &mut (dyn std::any::Any + 'a) {
322        lookup
323    }
324
325    fn from_any(
326        lookup: &mut dyn std::any::Any,
327    ) -> Option<&mut <Self::Backend as crate::sql_types::TypeMetadata>::MetadataLookup> {
328        lookup.downcast_mut()
329    }
330}
331
332impl SqliteConnection {
333    /// Run a transaction with `BEGIN IMMEDIATE`
334    ///
335    /// This method will return an error if a transaction is already open.
336    ///
337    /// # Example
338    ///
339    /// ```rust
340    /// # include!("../../doctest_setup.rs");
341    /// #
342    /// # fn main() {
343    /// #     run_test().unwrap();
344    /// # }
345    /// #
346    /// # fn run_test() -> QueryResult<()> {
347    /// #     let mut conn = SqliteConnection::establish(":memory:").unwrap();
348    /// conn.immediate_transaction(|conn| {
349    ///     // Do stuff in a transaction
350    ///     Ok(())
351    /// })
352    /// # }
353    /// ```
354    pub fn immediate_transaction<T, E, F>(&mut self, f: F) -> Result<T, E>
355    where
356        F: FnOnce(&mut Self) -> Result<T, E>,
357        E: From<Error>,
358    {
359        self.transaction_sql(f, "BEGIN IMMEDIATE")
360    }
361
362    /// Run a transaction with `BEGIN EXCLUSIVE`
363    ///
364    /// This method will return an error if a transaction is already open.
365    ///
366    /// # Example
367    ///
368    /// ```rust
369    /// # include!("../../doctest_setup.rs");
370    /// #
371    /// # fn main() {
372    /// #     run_test().unwrap();
373    /// # }
374    /// #
375    /// # fn run_test() -> QueryResult<()> {
376    /// #     let mut conn = SqliteConnection::establish(":memory:").unwrap();
377    /// conn.exclusive_transaction(|conn| {
378    ///     // Do stuff in a transaction
379    ///     Ok(())
380    /// })
381    /// # }
382    /// ```
383    pub fn exclusive_transaction<T, E, F>(&mut self, f: F) -> Result<T, E>
384    where
385        F: FnOnce(&mut Self) -> Result<T, E>,
386        E: From<Error>,
387    {
388        self.transaction_sql(f, "BEGIN EXCLUSIVE")
389    }
390
391    fn transaction_sql<T, E, F>(&mut self, f: F, sql: &str) -> Result<T, E>
392    where
393        F: FnOnce(&mut Self) -> Result<T, E>,
394        E: From<Error>,
395    {
396        AnsiTransactionManager::begin_transaction_sql(&mut *self, sql)?;
397        match f(&mut *self) {
398            Ok(value) => {
399                AnsiTransactionManager::commit_transaction(&mut *self)?;
400                Ok(value)
401            }
402            Err(e) => {
403                AnsiTransactionManager::rollback_transaction(&mut *self)?;
404                Err(e)
405            }
406        }
407    }
408
409    fn prepared_query<'conn, 'query, T>(
410        &'conn mut self,
411        source: T,
412    ) -> QueryResult<StatementUse<'conn, 'query>>
413    where
414        T: QueryFragment<Sqlite> + QueryId + 'query,
415    {
416        self.instrumentation
417            .on_connection_event(InstrumentationEvent::StartQuery {
418                query: &crate::debug_query(&source),
419            });
420        let raw_connection = &self.raw_connection;
421        let cache = &mut self.statement_cache;
422        let statement = match cache.cached_statement(
423            &source,
424            &Sqlite,
425            &[],
426            raw_connection,
427            Statement::prepare,
428            &mut *self.instrumentation,
429        ) {
430            Ok(statement) => statement,
431            Err(e) => {
432                self.instrumentation
433                    .on_connection_event(InstrumentationEvent::FinishQuery {
434                        query: &crate::debug_query(&source),
435                        error: Some(&e),
436                    });
437
438                return Err(e);
439            }
440        };
441
442        StatementUse::bind(statement, source, &mut *self.instrumentation)
443    }
444
445    #[doc(hidden)]
446    pub fn register_sql_function<ArgsSqlType, RetSqlType, Args, Ret, F>(
447        &mut self,
448        fn_name: &str,
449        deterministic: bool,
450        mut f: F,
451    ) -> QueryResult<()>
452    where
453        F: FnMut(Args) -> Ret + std::panic::UnwindSafe + Send + 'static,
454        Args: FromSqlRow<ArgsSqlType, Sqlite> + StaticallySizedRow<ArgsSqlType, Sqlite>,
455        Ret: ToSql<RetSqlType, Sqlite>,
456        Sqlite: HasSqlType<RetSqlType>,
457    {
458        functions::register(
459            &self.raw_connection,
460            fn_name,
461            deterministic,
462            move |_, args| f(args),
463        )
464    }
465
466    #[doc(hidden)]
467    pub fn register_noarg_sql_function<RetSqlType, Ret, F>(
468        &self,
469        fn_name: &str,
470        deterministic: bool,
471        f: F,
472    ) -> QueryResult<()>
473    where
474        F: FnMut() -> Ret + std::panic::UnwindSafe + Send + 'static,
475        Ret: ToSql<RetSqlType, Sqlite>,
476        Sqlite: HasSqlType<RetSqlType>,
477    {
478        functions::register_noargs(&self.raw_connection, fn_name, deterministic, f)
479    }
480
481    #[doc(hidden)]
482    pub fn register_aggregate_function<ArgsSqlType, RetSqlType, Args, Ret, A>(
483        &mut self,
484        fn_name: &str,
485    ) -> QueryResult<()>
486    where
487        A: SqliteAggregateFunction<Args, Output = Ret> + 'static + Send + std::panic::UnwindSafe,
488        Args: FromSqlRow<ArgsSqlType, Sqlite> + StaticallySizedRow<ArgsSqlType, Sqlite>,
489        Ret: ToSql<RetSqlType, Sqlite>,
490        Sqlite: HasSqlType<RetSqlType>,
491    {
492        functions::register_aggregate::<_, _, _, _, A>(&self.raw_connection, fn_name)
493    }
494
495    /// Register a collation function.
496    ///
497    /// `collation` must always return the same answer given the same inputs.
498    /// If `collation` panics and unwinds the stack, the process is aborted, since it is used
499    /// across a C FFI boundary, which cannot be unwound across and there is no way to
500    /// signal failures via the SQLite interface in this case..
501    ///
502    /// If the name is already registered it will be overwritten.
503    ///
504    /// This method will return an error if registering the function fails, either due to an
505    /// out-of-memory situation or because a collation with that name already exists and is
506    /// currently being used in parallel by a query.
507    ///
508    /// The collation needs to be specified when creating a table:
509    /// `CREATE TABLE my_table ( str TEXT COLLATE MY_COLLATION )`,
510    /// where `MY_COLLATION` corresponds to name passed as `collation_name`.
511    ///
512    /// # Example
513    ///
514    /// ```rust
515    /// # include!("../../doctest_setup.rs");
516    /// #
517    /// # fn main() {
518    /// #     run_test().unwrap();
519    /// # }
520    /// #
521    /// # fn run_test() -> QueryResult<()> {
522    /// #     let mut conn = SqliteConnection::establish(":memory:").unwrap();
523    /// // sqlite NOCASE only works for ASCII characters,
524    /// // this collation allows handling UTF-8 (barring locale differences)
525    /// conn.register_collation("RUSTNOCASE", |rhs, lhs| {
526    ///     rhs.to_lowercase().cmp(&lhs.to_lowercase())
527    /// })
528    /// # }
529    /// ```
530    pub fn register_collation<F>(&mut self, collation_name: &str, collation: F) -> QueryResult<()>
531    where
532        F: Fn(&str, &str) -> std::cmp::Ordering + Send + 'static + std::panic::UnwindSafe,
533    {
534        self.raw_connection
535            .register_collation_function(collation_name, collation)
536    }
537
538    /// Serialize the current SQLite database into a byte buffer.
539    ///
540    /// The serialized data is identical to the data that would be written to disk if the database
541    /// was saved in a file.
542    ///
543    /// # Returns
544    ///
545    /// This function returns a [`SerializedDatabase`] wrapping the serialized
546    /// bytes. If SQLite fails to allocate the buffer holding them, the failure
547    /// is reported by [`SerializedDatabase::try_as_slice`].
548    pub fn serialize_database_to_buffer(&mut self) -> SerializedDatabase {
549        self.raw_connection.serialize()
550    }
551
552    /// Deserialize an SQLite database from a byte buffer.
553    ///
554    /// This function takes a byte slice and attempts to deserialize it into a SQLite database.
555    /// If successful, the database is loaded into the connection. If the deserialization fails,
556    /// an error is returned.
557    ///
558    /// The database is opened in READONLY mode.
559    ///
560    /// # Example
561    ///
562    /// ```no_run
563    /// # use diesel::sqlite::SerializedDatabase;
564    /// # use diesel::sqlite::SqliteConnection;
565    /// # use diesel::result::QueryResult;
566    /// # use diesel::sql_query;
567    /// # use diesel::Connection;
568    /// # use diesel::RunQueryDsl;
569    /// # fn main() {
570    /// let connection = &mut SqliteConnection::establish(":memory:").unwrap();
571    ///
572    /// sql_query("CREATE TABLE users (id INTEGER PRIMARY KEY, name TEXT, email TEXT)")
573    ///     .execute(connection).unwrap();
574    /// sql_query("INSERT INTO users (name, email) VALUES ('John Doe', 'john.doe@example.com'), ('Jane Doe', 'jane.doe@example.com')")
575    ///     .execute(connection).unwrap();
576    ///
577    /// // Serialize the database to a byte vector
578    /// let serialized_db: SerializedDatabase = connection.serialize_database_to_buffer();
579    ///
580    /// // Create a new in-memory SQLite database
581    /// let connection = &mut SqliteConnection::establish(":memory:").unwrap();
582    ///
583    /// // Deserialize the byte vector into the new database
584    /// connection.deserialize_readonly_database_from_buffer(serialized_db.try_as_slice().unwrap()).unwrap();
585    /// #
586    /// # }
587    /// ```
588    // TODO: Diesel 3.0 This signature needs to change, we want to expose more options (schema name, readonly)
589    // and also ensure that this is not as unsafe as the current construct anymore. Maybe just accept a owned buffer or static pointer
590    // only instead? (So `Cow<'static, [u8]>`?)
591    #[allow(unsafe_code)]
592    pub fn deserialize_readonly_database_from_buffer(&mut self, data: &[u8]) -> QueryResult<()> {
593        // we copy the buffer here
594        // to make sure the underlying buffer lives as long as the connection
595        self.serialized_data.push(data.to_vec());
596        let last = self
597            .serialized_data
598            .last()
599            .expect("We literally pushed it above, so it's there");
600        unsafe {
601            // SAFETY: We store the buffer inside of the connection and we never touch it until
602            // we drop the connection
603            self.raw_connection.deserialize(last)
604        }
605    }
606
607    fn register_diesel_sql_functions(&self) -> QueryResult<()> {
608        use crate::sql_types::{Integer, Text};
609
610        // This function has side effects (creates triggers), so it should not
611        // be deterministic. We use DIRECTONLY to prevent it from being called
612        // from malicious schema objects in untrusted databases.
613        functions::register::<Text, Integer, _, _, _>(
614            &self.raw_connection,
615            "diesel_manage_updated_at",
616            false,
617            |conn, table_name: String| {
618                conn.exec(&::alloc::__export::must_use({
        ::alloc::fmt::format(format_args!("CREATE TRIGGER __diesel_manage_updated_at_{0}\nAFTER UPDATE ON {0}\nFOR EACH ROW WHEN\n  old.updated_at IS NULL AND\n  new.updated_at IS NULL OR\n  old.updated_at == new.updated_at\nBEGIN\n  UPDATE {0}\n  SET updated_at = CURRENT_TIMESTAMP\n  WHERE ROWID = new.ROWID;\nEND\n",
                table_name))
    })format!(
619                    include_str!("diesel_manage_updated_at.sql"),
620                    table_name = table_name
621                ))
622                .expect("Failed to create trigger");
623                0 // have to return *something*
624            },
625        )
626    }
627
628    fn establish_inner(database_url: &str) -> Result<SqliteConnection, ConnectionError> {
629        use crate::result::ConnectionError::CouldntSetupConfiguration;
630        let raw_connection = RawConnection::establish(database_url)?;
631        let conn = Self {
632            statement_cache: StatementCache::new(),
633            raw_connection,
634            transaction_state: AnsiTransactionManager::default(),
635            metadata_lookup: (),
636            instrumentation: DynInstrumentation::none(),
637            serialized_data: Vec::new(),
638        };
639        conn.register_diesel_sql_functions()
640            .map_err(CouldntSetupConfiguration)?;
641        Ok(conn)
642    }
643}
644
645fn error_message(err_code: libc::c_int) -> &'static str {
646    ffi::code_to_str(err_code)
647}
648
649#[cfg(test)]
650mod tests {
651    use super::*;
652    use crate::dsl::sql;
653    use crate::prelude::*;
654    use crate::sql_types::{Integer, Text};
655
656    fn connection() -> SqliteConnection {
657        SqliteConnection::establish(":memory:").unwrap()
658    }
659
660    #[declare_sql_function]
661    extern "SQL" {
662        fn fun_case(x: Text) -> Text;
663        fn my_add(x: Integer, y: Integer) -> Integer;
664        fn answer() -> Integer;
665        fn add_counter(x: Integer) -> Integer;
666
667        #[aggregate]
668        fn my_sum(expr: Integer) -> Integer;
669        #[aggregate]
670        fn range_max(expr1: Integer, expr2: Integer, expr3: Integer) -> Nullable<Integer>;
671    }
672
673    #[diesel_test_helper::test]
674    fn database_serializes_and_deserializes_successfully() {
675        let expected_users = vec![
676            (
677                1,
678                "John Doe".to_string(),
679                "john.doe@example.com".to_string(),
680            ),
681            (
682                2,
683                "Jane Doe".to_string(),
684                "jane.doe@example.com".to_string(),
685            ),
686        ];
687
688        let conn1 = &mut connection();
689        let _ =
690            crate::sql_query("CREATE TABLE users (id INTEGER PRIMARY KEY, name TEXT, email TEXT)")
691                .execute(conn1);
692        let _ = crate::sql_query("INSERT INTO users (name, email) VALUES ('John Doe', 'john.doe@example.com'), ('Jane Doe', 'jane.doe@example.com')")
693            .execute(conn1);
694
695        for _i in 0..2 {
696            let serialized_database = conn1.serialize_database_to_buffer();
697            let conn2 = &mut connection();
698            conn2
699                .deserialize_readonly_database_from_buffer(
700                    serialized_database.try_as_slice().unwrap(),
701                )
702                .unwrap();
703
704            let query =
705                sql::<(Integer, Text, Text)>("SELECT id, name, email FROM users ORDER BY id");
706            let actual_users = query.load::<(i32, String, String)>(conn2).unwrap();
707
708            assert_eq!(expected_users, actual_users);
709            // drop the database here
710            // and requery the database to make sure the database owns
711            // required data
712            std::mem::drop(serialized_database);
713            let query =
714                sql::<(Integer, Text, Text)>("SELECT id, name, email FROM users ORDER BY id");
715            let actual_users = query.load::<(i32, String, String)>(conn2).unwrap();
716
717            assert_eq!(expected_users, actual_users);
718        }
719    }
720
721    #[diesel_test_helper::test]
722    fn database_deserialize_random_bytes() {
723        let buffer = vec![0, 1, 2, 3, 4];
724        let conn = &mut SqliteConnection::establish(":memory:").unwrap();
725
726        conn.deserialize_readonly_database_from_buffer(&buffer)
727            .unwrap();
728
729        let r = sql::<Integer>("SELECT id FROM users").load::<i32>(conn);
730
731        assert!(r.is_err());
732        assert_eq!(r.unwrap_err().to_string(), "file is not a database");
733
734        let conn = &mut SqliteConnection::establish(":memory:").unwrap();
735
736        let _ =
737            crate::sql_query("CREATE TABLE users (id INTEGER PRIMARY KEY, name TEXT, email TEXT)")
738                .execute(conn);
739        let _ = crate::sql_query("INSERT INTO users (name, email) VALUES ('John Doe', 'john.doe@example.com'), ('Jane Doe', 'jane.doe@example.com')")
740            .execute(conn);
741
742        let db = conn.serialize_database_to_buffer();
743        // only get a valid header, but append garbage
744        let mut bad_buffer = db[..100].to_vec();
745        bad_buffer.extend(b"whatever");
746        conn.deserialize_readonly_database_from_buffer(&bad_buffer)
747            .unwrap();
748
749        let r = sql::<Integer>("SELECT id FROM users").load::<i32>(conn);
750
751        assert!(r.is_err());
752        assert_eq!(
753            r.unwrap_err().to_string(),
754            "database disk image is malformed"
755        );
756
757        // only get a valid header, but append garbage
758        let mut size_fitting_bad_buffer = db[..100].to_vec();
759        size_fitting_bad_buffer.extend(
760            core::iter::repeat(b"abcdefghij")
761                .flatten()
762                .take(db.len() - 100),
763        );
764        let r = conn.deserialize_readonly_database_from_buffer(&size_fitting_bad_buffer);
765
766        assert!(r.is_err());
767        assert_eq!(
768            r.unwrap_err().to_string(),
769            "database disk image is malformed"
770        );
771    }
772
773    #[diesel_test_helper::test]
774    fn database_serializes_empty_deserialized_database() {
775        let conn = &mut SqliteConnection::establish(":memory:").unwrap();
776        conn.deserialize_readonly_database_from_buffer(&[]).unwrap();
777
778        let serialized = conn.serialize_database_to_buffer();
779
780        assert!(serialized.is_empty());
781        assert!(serialized.try_as_slice().unwrap().is_empty());
782    }
783
784    #[cfg(not(all(target_family = "wasm", target_os = "unknown")))]
785    #[allow(unsafe_code)]
786    mod sqlite_serialize_oom {
787        use super::super::oom_test_support::{panic_message, run_in_child, with_heap_limit};
788        use super::super::{ffi, SerializedDatabase};
789        use crate::connection::{Connection, SimpleConnection};
790        use crate::sqlite::SqliteConnection;
791
792        const MIN_DATABASE_BYTES: i64 = 1_048_576;
793
794        // 64 KiB covers statement setup but cannot hold the 1 MiB serialization,
795        // pinning the failure to output allocation after SQLite reports its size.
796        fn with_failing_serialize<R>(f: impl FnOnce() -> R) -> R {
797            with_heap_limit(65_536, f)
798        }
799
800        #[test]
801        fn sqlite_serialize_oom_is_contained() {
802            run_in_child(|| {
803                let mut conn = large_database();
804
805                let (baseline_size, baseline) = serialize_direct(&conn);
806                assert!(
807                    baseline_size >= MIN_DATABASE_BYTES,
808                    "the serialized database is smaller than 1 MiB"
809                );
810                assert!(
811                    !baseline.is_null(),
812                    "SQLite refused to serialize a valid database"
813                );
814                // SAFETY: `sqlite3_serialize` returned this buffer and no wrapper owns it.
815                unsafe { ffi::sqlite3_free(baseline as _) };
816
817                let (reported_size, data) = with_failing_serialize(|| serialize_direct(&conn));
818                if !data.is_null() {
819                    // SAFETY: `sqlite3_serialize` returned this buffer and no wrapper owns it.
820                    unsafe { ffi::sqlite3_free(data as _) };
821                }
822                assert!(
823                    data.is_null(),
824                    "SQLite did not fail the output allocation of the serialization"
825                );
826                // SQLite reports the required size before attempting output allocation.
827                assert!(
828                    reported_size >= MIN_DATABASE_BYTES,
829                    "SQLite reported a serialization size of {reported_size} with a null buffer"
830                );
831
832                let serialized: SerializedDatabase =
833                    with_failing_serialize(|| conn.serialize_database_to_buffer());
834                let error = serialized
835                    .try_as_slice()
836                    .expect_err("the failed output allocation must surface as an error");
837                assert_eq!(error.to_string(), "out of memory");
838
839                let payload = std::panic::catch_unwind(core::panic::AssertUnwindSafe(|| {
840                    core::hint::black_box(serialized[0]);
841                }))
842                .expect_err("the serialized database access did not panic");
843                let message = panic_message(&*payload);
844                assert!(
845                    message.contains("Cannot access the serialized database: out of memory"),
846                    "SQLite serialization allocation failure surfaced as `{message}` instead \
847                     of a caught allocation panic"
848                );
849            });
850        }
851
852        fn large_database() -> SqliteConnection {
853            let mut conn = SqliteConnection::establish(":memory:").unwrap();
854            conn.batch_execute(&format!(
855                "CREATE TABLE blobs (id INTEGER PRIMARY KEY, payload BLOB);
856                 INSERT INTO blobs (payload) VALUES (zeroblob({MIN_DATABASE_BYTES}));"
857            ))
858            .unwrap();
859            conn
860        }
861
862        fn serialize_direct(conn: &SqliteConnection) -> (ffi::sqlite3_int64, *mut u8) {
863            // SAFETY: The connection is live, a null schema selects `main`, and `size` is a writable out-parameter.
864            unsafe {
865                let mut size: ffi::sqlite3_int64 = 0;
866                let data = ffi::sqlite3_serialize(
867                    conn.raw_connection.internal_connection.as_ptr(),
868                    core::ptr::null(),
869                    &mut size as *mut _,
870                    0,
871                );
872                (size, data)
873            }
874        }
875    }
876
877    #[diesel_test_helper::test]
878    fn register_custom_function() {
879        let connection = &mut connection();
880        fun_case_utils::register_impl(connection, |x: String| {
881            x.chars()
882                .enumerate()
883                .map(|(i, c)| {
884                    if i % 2 == 0 {
885                        c.to_lowercase().to_string()
886                    } else {
887                        c.to_uppercase().to_string()
888                    }
889                })
890                .collect::<String>()
891        })
892        .unwrap();
893
894        let mapped_string = crate::select(fun_case("foobar"))
895            .get_result::<String>(connection)
896            .unwrap();
897        assert_eq!("fOoBaR", mapped_string);
898    }
899
900    #[diesel_test_helper::test]
901    fn register_multiarg_function() {
902        let connection = &mut connection();
903        my_add_utils::register_impl(connection, |x: i32, y: i32| x + y).unwrap();
904
905        let added = crate::select(my_add(1, 2)).get_result::<i32>(connection);
906        assert_eq!(Ok(3), added);
907    }
908
909    #[diesel_test_helper::test]
910    fn register_noarg_function() {
911        let connection = &mut connection();
912        answer_utils::register_impl(connection, || 42).unwrap();
913
914        let answer = crate::select(answer()).get_result::<i32>(connection);
915        assert_eq!(Ok(42), answer);
916    }
917
918    #[diesel_test_helper::test]
919    fn register_nondeterministic_noarg_function() {
920        let connection = &mut connection();
921        answer_utils::register_nondeterministic_impl(connection, || 42).unwrap();
922
923        let answer = crate::select(answer()).get_result::<i32>(connection);
924        assert_eq!(Ok(42), answer);
925    }
926
927    #[diesel_test_helper::test]
928    fn register_nondeterministic_function() {
929        let connection = &mut connection();
930        let mut y = 0;
931        add_counter_utils::register_nondeterministic_impl(connection, move |x: i32| {
932            y += 1;
933            x + y
934        })
935        .unwrap();
936
937        let added = crate::select((add_counter(1), add_counter(1), add_counter(1)))
938            .get_result::<(i32, i32, i32)>(connection);
939        assert_eq!(Ok((2, 3, 4)), added);
940    }
941
942    #[derive(Default)]
943    struct MySum {
944        sum: i32,
945    }
946
947    impl SqliteAggregateFunction<i32> for MySum {
948        type Output = i32;
949
950        fn step(&mut self, expr: i32) {
951            self.sum += expr;
952        }
953
954        fn finalize(aggregator: Option<Self>) -> Self::Output {
955            aggregator.map(|a| a.sum).unwrap_or_default()
956        }
957    }
958
959    table! {
960        my_sum_example {
961            id -> Integer,
962            value -> Integer,
963        }
964    }
965
966    #[diesel_test_helper::test]
967    fn register_aggregate_function() {
968        use self::my_sum_example::dsl::*;
969
970        let connection = &mut connection();
971        crate::sql_query(
972            "CREATE TABLE my_sum_example (id integer primary key autoincrement, value integer)",
973        )
974        .execute(connection)
975        .unwrap();
976        crate::sql_query("INSERT INTO my_sum_example (value) VALUES (1), (2), (3)")
977            .execute(connection)
978            .unwrap();
979
980        my_sum_utils::register_impl::<MySum, _>(connection).unwrap();
981
982        let result = my_sum_example
983            .select(my_sum(value))
984            .get_result::<i32>(connection);
985        assert_eq!(Ok(6), result);
986    }
987
988    #[diesel_test_helper::test]
989    fn register_aggregate_function_returns_finalize_default_on_empty_set() {
990        use self::my_sum_example::dsl::*;
991
992        let connection = &mut connection();
993        crate::sql_query(
994            "CREATE TABLE my_sum_example (id integer primary key autoincrement, value integer)",
995        )
996        .execute(connection)
997        .unwrap();
998
999        my_sum_utils::register_impl::<MySum, _>(connection).unwrap();
1000
1001        let result = my_sum_example
1002            .select(my_sum(value))
1003            .get_result::<i32>(connection);
1004        assert_eq!(Ok(0), result);
1005    }
1006
1007    #[derive(Default)]
1008    struct RangeMax<T> {
1009        max_value: Option<T>,
1010    }
1011
1012    impl<T: Default + Ord + Copy + Clone> SqliteAggregateFunction<(T, T, T)> for RangeMax<T> {
1013        type Output = Option<T>;
1014
1015        fn step(&mut self, (x0, x1, x2): (T, T, T)) {
1016            let max = if x0 >= x1 && x0 >= x2 {
1017                x0
1018            } else if x1 >= x0 && x1 >= x2 {
1019                x1
1020            } else {
1021                x2
1022            };
1023
1024            self.max_value = match self.max_value {
1025                Some(current_max_value) if max > current_max_value => Some(max),
1026                None => Some(max),
1027                _ => self.max_value,
1028            };
1029        }
1030
1031        fn finalize(aggregator: Option<Self>) -> Self::Output {
1032            aggregator?.max_value
1033        }
1034    }
1035
1036    table! {
1037        range_max_example {
1038            id -> Integer,
1039            value1 -> Integer,
1040            value2 -> Integer,
1041            value3 -> Integer,
1042        }
1043    }
1044
1045    #[diesel_test_helper::test]
1046    fn register_aggregate_multiarg_function() {
1047        use self::range_max_example::dsl::*;
1048
1049        let connection = &mut connection();
1050        crate::sql_query(
1051            r#"CREATE TABLE range_max_example (
1052                id integer primary key autoincrement,
1053                value1 integer,
1054                value2 integer,
1055                value3 integer
1056            )"#,
1057        )
1058        .execute(connection)
1059        .unwrap();
1060        crate::sql_query(
1061            "INSERT INTO range_max_example (value1, value2, value3) VALUES (3, 2, 1), (2, 2, 2)",
1062        )
1063        .execute(connection)
1064        .unwrap();
1065
1066        range_max_utils::register_impl::<RangeMax<i32>, _, _, _>(connection).unwrap();
1067        let result = range_max_example
1068            .select(range_max(value1, value2, value3))
1069            .get_result::<Option<i32>>(connection)
1070            .unwrap();
1071        assert_eq!(Some(3), result);
1072    }
1073
1074    table! {
1075        my_collation_example {
1076            id -> Integer,
1077            value -> Text,
1078        }
1079    }
1080
1081    #[diesel_test_helper::test]
1082    fn register_collation_function() {
1083        use self::my_collation_example::dsl::*;
1084
1085        let connection = &mut connection();
1086
1087        connection
1088            .register_collation("RUSTNOCASE", |rhs, lhs| {
1089                rhs.to_lowercase().cmp(&lhs.to_lowercase())
1090            })
1091            .unwrap();
1092
1093        crate::sql_query(
1094                "CREATE TABLE my_collation_example (id integer primary key autoincrement, value text collate RUSTNOCASE)",
1095            ).execute(connection)
1096            .unwrap();
1097        crate::sql_query(
1098            "INSERT INTO my_collation_example (value) VALUES ('foo'), ('FOo'), ('f00')",
1099        )
1100        .execute(connection)
1101        .unwrap();
1102
1103        let result = my_collation_example
1104            .filter(value.eq("foo"))
1105            .select(value)
1106            .load::<String>(connection);
1107        assert_eq!(
1108            Ok(&["foo".to_owned(), "FOo".to_owned()][..]),
1109            result.as_ref().map(|vec| vec.as_ref())
1110        );
1111
1112        let result = my_collation_example
1113            .filter(value.eq("FOO"))
1114            .select(value)
1115            .load::<String>(connection);
1116        assert_eq!(
1117            Ok(&["foo".to_owned(), "FOo".to_owned()][..]),
1118            result.as_ref().map(|vec| vec.as_ref())
1119        );
1120
1121        let result = my_collation_example
1122            .filter(value.eq("f00"))
1123            .select(value)
1124            .load::<String>(connection);
1125        assert_eq!(
1126            Ok(&["f00".to_owned()][..]),
1127            result.as_ref().map(|vec| vec.as_ref())
1128        );
1129
1130        let result = my_collation_example
1131            .filter(value.eq("F00"))
1132            .select(value)
1133            .load::<String>(connection);
1134        assert_eq!(
1135            Ok(&["f00".to_owned()][..]),
1136            result.as_ref().map(|vec| vec.as_ref())
1137        );
1138
1139        let result = my_collation_example
1140            .filter(value.eq("oof"))
1141            .select(value)
1142            .load::<String>(connection);
1143        assert_eq!(Ok(&[][..]), result.as_ref().map(|vec| vec.as_ref()));
1144    }
1145
1146    // regression test for https://github.com/diesel-rs/diesel/issues/3425
1147    #[diesel_test_helper::test]
1148    fn test_correct_serialization_of_owned_strings() {
1149        use crate::prelude::*;
1150
1151        #[derive(Debug, crate::expression::AsExpression)]
1152        #[diesel(sql_type = diesel::sql_types::Text)]
1153        struct CustomWrapper(String);
1154
1155        impl crate::serialize::ToSql<Text, Sqlite> for CustomWrapper {
1156            fn to_sql<'b>(
1157                &'b self,
1158                out: &mut crate::serialize::Output<'b, '_, Sqlite>,
1159            ) -> crate::serialize::Result {
1160                out.set_value(self.0.to_string());
1161                Ok(crate::serialize::IsNull::No)
1162            }
1163        }
1164
1165        let connection = &mut connection();
1166
1167        let res = crate::select(
1168            CustomWrapper("".into())
1169                .into_sql::<crate::sql_types::Text>()
1170                .nullable(),
1171        )
1172        .get_result::<Option<String>>(connection)
1173        .unwrap();
1174        assert_eq!(res, Some(String::new()));
1175    }
1176
1177    #[diesel_test_helper::test]
1178    fn test_correct_serialization_of_owned_bytes() {
1179        use crate::prelude::*;
1180
1181        #[derive(Debug, crate::expression::AsExpression)]
1182        #[diesel(sql_type = diesel::sql_types::Binary)]
1183        struct CustomWrapper(Vec<u8>);
1184
1185        impl crate::serialize::ToSql<crate::sql_types::Binary, Sqlite> for CustomWrapper {
1186            fn to_sql<'b>(
1187                &'b self,
1188                out: &mut crate::serialize::Output<'b, '_, Sqlite>,
1189            ) -> crate::serialize::Result {
1190                out.set_value(self.0.clone());
1191                Ok(crate::serialize::IsNull::No)
1192            }
1193        }
1194
1195        let connection = &mut connection();
1196
1197        let res = crate::select(
1198            CustomWrapper(Vec::new())
1199                .into_sql::<crate::sql_types::Binary>()
1200                .nullable(),
1201        )
1202        .get_result::<Option<Vec<u8>>>(connection)
1203        .unwrap();
1204        assert_eq!(res, Some(Vec::new()));
1205    }
1206
1207    #[diesel_test_helper::test]
1208    fn correctly_handle_empty_query() {
1209        let check_empty_query_error = |r: crate::QueryResult<usize>| {
1210            assert!(r.is_err());
1211            let err = r.unwrap_err();
1212            assert!(
1213                matches!(err, crate::result::Error::QueryBuilderError(ref b) if b.is::<crate::result::EmptyQuery>()),
1214                "Expected a query builder error, but got {err}"
1215            );
1216        };
1217        let connection = &mut SqliteConnection::establish(":memory:").unwrap();
1218        check_empty_query_error(crate::sql_query("").execute(connection));
1219        check_empty_query_error(crate::sql_query("   ").execute(connection));
1220        check_empty_query_error(crate::sql_query("\n\t").execute(connection));
1221        check_empty_query_error(crate::sql_query("-- SELECT 1;").execute(connection));
1222    }
1223
1224    #[diesel_test_helper::test]
1225    fn aggregate_function_works_with_aligned_data() {
1226        #[derive(Debug, Default)]
1227        #[repr(align(64))]
1228        struct OverAligned;
1229
1230        impl SqliteAggregateFunction<i32> for OverAligned {
1231            type Output = i64;
1232
1233            fn step(&mut self, _value: i32) {
1234                let need = core::mem::align_of::<Self>();
1235                let got = core::mem::align_of_val(self);
1236                assert_eq!(need, got);
1237            }
1238
1239            fn finalize(_agg: Option<Self>) -> i64 {
1240                0
1241            }
1242        }
1243        #[declare_sql_function]
1244        extern "SQL" {
1245            #[aggregate]
1246            fn over_aligned_sum(x: Integer) -> diesel::sql_types::BigInt;
1247        }
1248
1249        let mut conn = SqliteConnection::establish(":memory:").unwrap();
1250        over_aligned_sum_utils::register_impl::<OverAligned, _>(&mut conn).unwrap();
1251
1252        diesel::select(over_aligned_sum(1))
1253            .execute(&mut conn)
1254            .unwrap();
1255    }
1256
1257    #[diesel_test_helper::test]
1258    fn sum_twice() {
1259        #[derive(Default)]
1260        struct Sum(i32);
1261
1262        impl SqliteAggregateFunction<i32> for Sum {
1263            type Output = i32;
1264
1265            fn step(&mut self, value: i32) {
1266                self.0 += value;
1267            }
1268
1269            fn finalize(agg: Option<Self>) -> i32 {
1270                agg.map(|s| s.0).unwrap_or_default()
1271            }
1272        }
1273
1274        #[declare_sql_function]
1275        extern "SQL" {
1276            #[aggregate]
1277            fn my_sum(x: Integer) -> Integer;
1278        }
1279
1280        let mut conn = SqliteConnection::establish(":memory:").unwrap();
1281        my_sum_utils::register_impl::<Sum, _>(&mut conn).unwrap();
1282
1283        conn.batch_execute(
1284            "
1285            CREATE TABLE test(key1 INTEGER, key2 INTEGER);
1286            INSERT INTO test(key1, key2) VALUES (1, 2), (2, 4), (3, 6);
1287",
1288        )
1289        .unwrap();
1290
1291        table! {
1292            test (key1, key2) {
1293                key1 -> Integer,
1294                key2 -> Integer,
1295            }
1296        }
1297
1298        let (first_res, second_res) = test::table
1299            .select((my_sum(test::key1), my_sum(test::key2)))
1300            .get_result::<(i32, i32)>(&mut conn)
1301            .unwrap();
1302
1303        assert_eq!(first_res, 6);
1304        assert_eq!(second_res, 12);
1305
1306        conn.batch_execute("DELETE FROM test").unwrap();
1307        let (first_res, second_res) = test::table
1308            .select((my_sum(test::key1), my_sum(test::key2)))
1309            .get_result::<(i32, i32)>(&mut conn)
1310            .unwrap();
1311
1312        assert_eq!(first_res, 0);
1313        assert_eq!(second_res, 0);
1314    }
1315
1316    #[diesel_test_helper::test]
1317    fn test_injection() {
1318        diesel::table! {
1319            #[sql_name = "quote'table"]
1320            quote_table (id) {
1321                id -> Nullable<Integer>,
1322                name -> Nullable<Text>,
1323            }
1324        }
1325
1326        let mut conn = SqliteConnection::establish(":memory:").unwrap();
1327
1328        conn.batch_execute("CREATE TABLE \"quote'table\" (id INTEGER PRIMARY KEY, name TEXT);")
1329            .unwrap();
1330
1331        diesel::insert_into(quote_table::table)
1332            .values((quote_table::id.eq(1), quote_table::name.eq("Jane")))
1333            .execute(&mut conn)
1334            .unwrap();
1335
1336        let data = quote_table::table
1337            .load::<(Option<i32>, Option<String>)>(&mut conn)
1338            .unwrap();
1339        assert_eq!(data, [(Some(1), Some("Jane".to_owned()))]);
1340    }
1341}
diesel/sqlite/connection/mod.rs

diesel/sqlite/connection/
mod.rs
