omry_db/
db.rs

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

omry_db/
db.rs
