omry_db/
db.rs

1//! Database abstractions.
2#![cfg(feature = "connection")]
3
4use crate::model::{DbRecord, Schema};
5use deadpool_sqlite::rusqlite;
6use deadpool_sqlite::rusqlite::Row;
7use deadpool_sqlite::rusqlite::types::Value;
8use deadpool_sqlite::{Config, Object, Pool, Runtime};
9use std::path::Path;
10use std::rc::Rc;
11use thiserror::Error;
12use tracing::instrument;
13
14/// Maximum number of 'bound' (programmatically populated)
15/// parameters allowed in a query.
16const MAX_BIND_PARAMS: usize = 32766;
17
18/// Tries to query for rows in `records` using the given statement, and convert them to [`DbRecord`].
19macro_rules! try_query_db_records {
20    ($stmt:ident, $values:expr) => {{
21        let rows = $stmt
22            .query_map($values, |row| DbRecord::try_from(row))
23            .inspect_err(|e| tracing::error!("Couldn't execute the query: {e}"))?;
24
25        let mut records = vec![];
26        for maybe_record in rows {
27            let record = maybe_record
28                .inspect_err(|e| tracing::error!("Couldn't get the mapped record: {e}"))?;
29            records.push(record);
30        }
31
32        records
33    }};
34}
35
36/// Errors that may occur when interacting with database APIs.
37#[derive(Debug, Error)]
38#[error(transparent)]
39pub struct DatabaseError(DatabaseErrorRepr);
40
41// Implement this manually so we can still conveniently use ? with errors
42// that map onto variants of DatabaseErrorRepr
43// (using #[from] would conflict with this).
44impl<E> From<E> for DatabaseError
45where
46    E: Into<DatabaseErrorRepr>,
47{
48    fn from(error: E) -> Self {
49        Self(error.into())
50    }
51}
52
53/// Database errors.
54#[derive(Error, Debug)]
55enum DatabaseErrorRepr {
56    /// Error when creating the database connection pool.
57    #[error("error creating the DB pool: {0}")]
58    CreatePool(#[from] deadpool_sqlite::CreatePoolError),
59
60    /// Error when interacting with the database pool.
61    // This error has a string payload and a manual implementation of [`From`]
62    // because `deadpool_sqlite::InteractError` isn't Sync.
63    #[error("error interacting with the DB pool: {0}")]
64    PoolInteract(String),
65
66    /// Error when getting a new connection from the database pool.
67    #[error("error getting connection from the DB pool: {0}")]
68    Pool(#[from] deadpool_sqlite::PoolError),
69
70    /// Any unspecified error returned by [`rusqlite`].
71    #[error("error in rusqlite: {0}")]
72    Rusqlite(#[from] rusqlite::Error),
73
74    /// Error returned when a query function was given too many
75    /// parameters to bind.
76    #[error(
77        "The caller passed too many parameters to bind (got {0}, max allowed is {MAX_BIND_PARAMS})"
78    )]
79    TooManyBindParameters(usize),
80}
81
82impl From<deadpool_sqlite::InteractError> for DatabaseErrorRepr {
83    fn from(error: deadpool_sqlite::InteractError) -> Self {
84        Self::PoolInteract(error.to_string())
85    }
86}
87
88/// Alias for results returned by the database APIs.
89pub type Result<T> = std::result::Result<T, DatabaseError>;
90
91impl TryFrom<&Row<'_>> for DbRecord {
92    type Error = rusqlite::Error;
93
94    fn try_from(row: &Row) -> std::result::Result<Self, Self::Error> {
95        Ok(Self {
96            id: row.get(0)?,
97            url: row.get(1)?,
98            client_datetime: row.get(2)?,
99            client_datetime_raw: row.get(3)?,
100            timestamp_flora: row.get(4)?,
101            title: row.get(5)?,
102            language: row.get(6)?,
103            web_archive_serialized: row.get(7)?,
104        })
105    }
106}
107
108/// Abstraction for flora's primary datastore (SQLite).
109#[derive(Debug, Clone)]
110pub struct Database {
111    /// The connection pool.
112    pool: Pool,
113}
114
115impl Database {
116    /// Connects to the SQLite database at the given URL and creates a new `Database`.
117    ///
118    /// ## Errors
119    /// Returns an instance of [`DatabaseError`] on error.
120    #[instrument(skip(path))]
121    async fn new(path: impl AsRef<Path>) -> Result<Self> {
122        let cfg = Config::new(path.as_ref());
123        let pool = cfg.create_pool(Runtime::Tokio1)?;
124        Ok(Self { pool })
125    }
126
127    /// Gets a connection, and calls [`rusqlite::vtab::array::load_module`] with it.
128    ///
129    /// This method should be called to get any connection that will be used
130    /// to perform queries using the `array` extension of [`rusqlite`].
131    ///
132    /// See <https://docs.rs/rusqlite/latest/rusqlite/vtab/array/index.html>.
133    #[instrument(skip(self))]
134    async fn get_conn_array(&self) -> Result<Object> {
135        let conn = self.pool.get().await?;
136        conn.interact(|conn| {
137            tracing::debug!("Loading rusqlite::vtab::array");
138            rusqlite::vtab::array::load_module(conn)
139        })
140        .await?
141        .inspect(|()| tracing::debug!("Loaded rusqlite::vtab::array"))?;
142
143        Ok(conn)
144    }
145
146    /// Creates a table to store the records.
147    #[instrument(skip(self))]
148    async fn create_table(&self, schema: &Schema) -> Result<()> {
149        let create_table_query = format!("CREATE TABLE IF NOT EXISTS {schema};");
150        let conn = self.pool.get().await?;
151        let _rows: usize = conn
152            .interact(move |conn| {
153                let mut stmt = conn.prepare(&create_table_query)?;
154                stmt.execute(())
155            })
156            .await??;
157
158        Ok(())
159    }
160
161    /// Creates indexes for the `records` table.
162    #[instrument(skip(self))]
163    async fn create_indexes(&self, schema: &Schema, columns: &[&str]) -> Result<()> {
164        let table_name = schema.table_name();
165        let conn = self.pool.get().await?;
166
167        for column in columns {
168            let idx_name = format!("idx_{column}");
169            let query =
170                format!("CREATE INDEX IF NOT EXISTS {idx_name} ON {table_name} ({column});");
171            tracing::debug!("Query: {query}");
172
173            tracing::info!("Creating index '{idx_name}' on '{table_name}'");
174            conn.interact(move |conn| {
175                let mut stmt = conn.prepare(&query)?;
176                stmt.execute([])
177            })
178            .await??;
179        }
180
181        Ok(())
182    }
183
184    /// Initializes the database with a 'records' table.
185    ///
186    /// ## Errors
187    /// Returns [`DatabaseError`] if the operation fails.
188    #[instrument(skip(path))]
189    pub async fn new_with_records_table(path: impl AsRef<Path>) -> Result<Self> {
190        let db = Self::new(path).await?;
191        let schema = DbRecord::schema();
192        db.create_table(&schema).await?;
193        db.create_indexes(&schema, &["url"]).await?;
194        Ok(db)
195    }
196
197    /// Queries the database to get a count of records it contains.
198    ///
199    /// ## Errors
200    /// Returns [`DatabaseError`] if the operation fails.
201    #[instrument]
202    pub async fn count(&self) -> Result<i64> {
203        let query = r"SELECT count(*) FROM records;";
204        let conn = self.pool.get().await?;
205
206        let record_count = conn
207            .interact(move |conn| {
208                let mut stmt = conn.prepare(query)?;
209                stmt.query_one([], |row| row.get(0))
210            })
211            .await??;
212
213        Ok(record_count)
214    }
215
216    /// Inserts the given record into the `Database`.
217    ///
218    /// Returns the ID of the inserted row.
219    ///
220    /// ## Errors
221    /// Returns [`DatabaseError`] if the insert operation fails.
222    pub async fn insert(&self, record: DbRecord) -> Result<i64> {
223        let query = r"INSERT INTO records
224        VALUES (:id, :url, :client_datetime, :client_datetime_raw, :timestamp_flora, :title,
225        :language, :web_archive_serialized)
226        RETURNING id;";
227
228        let conn = self.pool.get().await?;
229
230        let row_id = conn
231            .interact(move |conn| {
232                let mut stmt = conn.prepare(query)?;
233                stmt.query_one(
234                    rusqlite::named_params! {
235                        ":id": record.id,
236                        ":url": record.url,
237                        ":client_datetime": record.client_datetime,
238                        ":client_datetime_raw": record.client_datetime_raw,
239                        ":timestamp_flora": record.timestamp_flora,
240                        ":title": record.title,
241                        ":language": record.language,
242                        ":web_archive_serialized": record.web_archive_serialized,
243                    },
244                    |row| row.get(0),
245                )
246            })
247            .await??;
248
249        Ok(row_id)
250    }
251
252    /// Deletes the record corresponding to the given id.
253    ///
254    /// ## Errors
255    /// Returns [`DatabaseError`] if the operation fails.
256    #[instrument(skip(self))]
257    pub async fn delete(&self, id: i64) -> Result<usize> {
258        const QUERY: &str = r"DELETE FROM records WHERE id = ?1;";
259
260        let conn = self.pool.get().await?;
261        let rows_changed = conn
262            .interact(move |conn| {
263                let mut stmt = conn.prepare(QUERY)?;
264                stmt.execute([id])
265            })
266            .await??;
267
268        Ok(rows_changed)
269    }
270
271    /// Get a record from SQLite given an id.
272    ///
273    /// ## Errors
274    /// Returns [`DatabaseError`] if the operation fails.
275    pub async fn get_record(&self, id: i64) -> Result<DbRecord> {
276        let query = r"SELECT id, url, client_datetime, client_datetime_raw, timestamp_flora, title,
277            language, web_archive_serialized
278            FROM records WHERE id = ?1;";
279
280        let conn = self.pool.get().await?;
281        let result = conn
282            .interact(move |conn| {
283                let mut stmt = conn.prepare(query)?;
284                stmt.query_row([id], |row| {
285                    let record = row.try_into()?;
286                    Ok(record)
287                })
288            })
289            .await??;
290
291        Ok(result)
292    }
293
294    /// Get multiple records from SQLite given their IDs.
295    ///
296    /// ## Errors
297    /// Returns [`DatabaseError`] if the operation fails.
298    #[instrument]
299    pub async fn get_many_records(&self, ids: &[i64]) -> Result<Vec<DbRecord>> {
300        if ids.len() > MAX_BIND_PARAMS {
301            return Err(DatabaseErrorRepr::TooManyBindParameters(ids.len()).into());
302        }
303
304        tracing::debug!("Getting records.");
305
306        let query = r"SELECT id, url, client_datetime, client_datetime_raw, timestamp_flora, title,
307            language, web_archive_serialized
308            FROM records
309            WHERE id IN rarray(?1);";
310
311        let conn = self
312            .get_conn_array()
313            .await
314            .inspect_err(|e| tracing::error!("Couldn't get a connection from the pool: {e}"))?;
315
316        let values = ids.iter().copied().map(Value::from).collect::<Vec<Value>>();
317
318        let records = conn
319            .interact(|conn| -> Result<Vec<DbRecord>> {
320                let values = Rc::new(values);
321
322                let mut stmt = conn
323                    .prepare(query)
324                    .inspect_err(|e| tracing::error!("Couldn't prepare the query: {e}"))?;
325
326                let records = try_query_db_records!(stmt, [values]);
327                Ok(records)
328            })
329            .await
330            .inspect_err(|e| {
331                tracing::error!("Couldn't interact with the database connection: {e}");
332            })??;
333
334        tracing::debug!("Got {} records", records.len());
335
336        Ok(records)
337    }
338
339    /// Selects all rows in the `records` table, without any pagination.
340    ///
341    /// ## Errors
342    ///
343    /// Returns [`DatabaseError`] if the operation fails.
344    #[instrument(skip(self))]
345    pub async fn get_all_records(&self) -> Result<Vec<DbRecord>> {
346        tracing::debug!("Getting all records");
347
348        let query = r"SELECT * FROM records";
349        let conn = self.pool.get().await?;
350        let records = conn
351            .interact(move |conn| -> Result<Vec<DbRecord>> {
352                let mut stmt = conn.prepare(query)?;
353                let records = try_query_db_records!(stmt, []);
354                Ok(records)
355            })
356            .await??;
357
358        Ok(records)
359    }
360}
361
362#[cfg(test)]
363mod tests {
364    use super::*;
365    use crate::model::tests::test_record_params;
366    use anyhow::Result;
367    use chrono::{DateTime, FixedOffset};
368    use omry_archiving::record::Record;
369    use pretty_assertions::assert_eq;
370    use serde::Serialize;
371    use std::fs::File;
372    use std::ops::Deref;
373    use tempfile::NamedTempFile;
374
375    /// Data structure to keep the temporary file holding the database alive,
376    /// because it will get unlinked on drop.
377    struct TempDb {
378        #[allow(dead_code)]
379        temp_db_file: NamedTempFile<File>,
380        db: Database,
381    }
382
383    impl TempDb {
384        async fn new() -> Result<Self> {
385            let temp_db_file = NamedTempFile::with_suffix(".db")?;
386            let db = Database::new_with_records_table(temp_db_file.path()).await?;
387            Ok(Self { temp_db_file, db })
388        }
389    }
390
391    impl Deref for TempDb {
392        type Target = Database;
393
394        fn deref(&self) -> &Self::Target {
395            &self.db
396        }
397    }
398
399    /// Wrapper for [`DbRecord`] that supports [`Serialize`] for redacting snapshots.
400    #[derive(Debug, Serialize)]
401    struct DbRecordDef {
402        id: Option<i64>,
403        url: String,
404        client_datetime: DateTime<FixedOffset>,
405        client_datetime_raw: Option<String>,
406        timestamp_flora: i64,
407        title: String,
408        language: Option<String>,
409        web_archive_serialized: Vec<u8>,
410    }
411
412    impl From<DbRecord> for DbRecordDef {
413        fn from(value: DbRecord) -> DbRecordDef {
414            let DbRecord {
415                id,
416                url,
417                client_datetime,
418                client_datetime_raw,
419                timestamp_flora,
420                title,
421                language,
422                web_archive_serialized,
423            } = value;
424
425            Self {
426                id,
427                url,
428                client_datetime,
429                client_datetime_raw,
430                timestamp_flora,
431                title,
432                language,
433                web_archive_serialized,
434            }
435        }
436    }
437
438    fn test_db_record() -> Result<DbRecord> {
439        let record = Record::new(test_record_params()?);
440        Ok(record.try_into()?)
441    }
442
443    #[tokio::test]
444    async fn can_create_table() -> Result<()> {
445        let db = TempDb::new().await?;
446        assert!(db.create_table(&DbRecord::schema()).await.is_ok());
447        // TODO: Check the resulting DB schema matches what we gave to create_table
448        Ok(())
449    }
450
451    #[tokio::test]
452    async fn can_add_record() -> Result<()> {
453        let db = TempDb::new().await?;
454        let record = test_db_record()?;
455        let result = db.insert(record).await?;
456        assert_eq!(result, 1);
457        Ok(())
458    }
459
460    #[tokio::test]
461    async fn can_get_record() -> Result<()> {
462        let db = TempDb::new().await?;
463        let record = test_db_record()?;
464        assert!(record.id.is_none());
465        db.insert(record).await?;
466
467        let expected_id = 1i64;
468        let retrieved_record = db.get_record(expected_id).await?;
469
470        let retrieved_record: DbRecordDef = retrieved_record.into();
471        insta::assert_ron_snapshot!(retrieved_record, {
472            ".client_datetime" => "[client_datetime]",
473            ".timestamp_flora" => "[timestamp]"
474        });
475
476        Ok(())
477    }
478
479    #[tokio::test]
480    async fn can_get_many_records() -> Result<()> {
481        let db = TempDb::new().await?;
482
483        let num_records_to_insert = 7;
484        let mut record_ids = Vec::with_capacity(num_records_to_insert);
485
486        for _ in 0..num_records_to_insert {
487            let record = test_db_record()?;
488            let record_id = db.insert(record).await?;
489            record_ids.push(record_id);
490        }
491
492        let selected_records = db
493            .get_many_records(&record_ids)
494            .await?
495            .into_iter()
496            .map(DbRecordDef::from)
497            .collect::<Vec<_>>();
498
499        assert_eq!(selected_records.len(), num_records_to_insert);
500
501        insta::assert_ron_snapshot!(selected_records, {
502            "[].client_datetime" => "[client_datetime]",
503            "[].timestamp_flora" => "[timestamp]"
504        });
505
506        Ok(())
507    }
508
509    #[tokio::test]
510    async fn passing_too_many_bind_parameters_returns_error() -> Result<()> {
511        let db = TempDb::new().await?;
512
513        let num_records_to_insert = MAX_BIND_PARAMS + 1;
514        db.insert(test_db_record()?).await?;
515
516        let mut ids = Vec::with_capacity(num_records_to_insert);
517        for index in 0..num_records_to_insert {
518            ids.push(i64::try_from(index + 1)?);
519        }
520
521        let result = db.get_many_records(&ids).await;
522        assert!(result.is_err());
523
524        // At this point, we know this is definitely an error.
525        #[allow(clippy::unwrap_used)]
526        let error = result.unwrap_err();
527        insta::assert_snapshot!(error);
528
529        Ok(())
530    }
531
532    #[tokio::test]
533    async fn can_get_count() -> Result<()> {
534        let db = TempDb::new().await?;
535
536        let num_records_to_insert = 7;
537        let mut record_ids = Vec::with_capacity(num_records_to_insert);
538
539        for _ in 0..num_records_to_insert {
540            let record = test_db_record()?;
541            let record_id = db.insert(record).await?;
542            record_ids.push(record_id);
543        }
544
545        let count: usize = db.count().await?.try_into()?;
546
547        assert_eq!(count, num_records_to_insert);
548        Ok(())
549    }
550
551    #[tokio::test]
552    async fn can_delete_record() -> Result<()> {
553        let db = TempDb::new().await?;
554        let record = test_db_record()?;
555        let id = db.insert(record).await?;
556        let rows_changed = db.delete(id).await?;
557        assert_eq!(rows_changed, 1);
558        Ok(())
559    }
560
561    #[tokio::test]
562    async fn can_get_all_records() -> Result<()> {
563        let db = TempDb::new().await?;
564
565        let num_records_to_insert = 7;
566        let mut record_ids = Vec::with_capacity(num_records_to_insert);
567
568        for _ in 0..num_records_to_insert {
569            let record = test_db_record()?;
570            let record_id = db.insert(record).await?;
571            record_ids.push(record_id);
572        }
573
574        let selected_records = db
575            .get_all_records()
576            .await?
577            .into_iter()
578            .map(DbRecordDef::from)
579            .collect::<Vec<_>>();
580
581        assert_eq!(selected_records.len(), num_records_to_insert);
582
583        insta::assert_ron_snapshot!(selected_records, {
584            "[].client_datetime" => "[client_datetime]",
585            "[].timestamp_flora" => "[timestamp]",
586            "[].web_archive_serialized" => "[data]"
587        });
588
589        Ok(())
590    }
591}
omry_db/db.rs

omry_db/
db.rs
