omry_db/
db.rs

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

omry_db/
db.rs
