omry_db/

db.rs

//! Database abstractions.
use crate::model::{DbRecord, Schema};
use deadpool_sqlite::rusqlite;
use deadpool_sqlite::rusqlite::Row;
use deadpool_sqlite::rusqlite::types::Value;
use deadpool_sqlite::{Config, Object, Pool, Runtime};
use std::path::Path;
use std::rc::Rc;
use thiserror::Error;
use tracing::instrument;

/// Maximum number of 'bound' (programmatically populated)
/// parameters allowed in a query.
const MAX_BIND_PARAMS: usize = 32766;

/// Tries to query for rows in `records` using the given statement, and convert them to [`DbRecord`].
macro_rules! try_query_db_records {
    ($stmt:ident, $values:expr) => {{
        let rows = $stmt
            .query_map($values, |row| DbRecord::try_from(row))
            .inspect_err(|e| tracing::error!("Couldn't execute the query: {e}"))?;

        let mut records = vec![];
        for maybe_record in rows {
            let record = maybe_record
                .inspect_err(|e| tracing::error!("Couldn't get the mapped record: {e}"))?;
            records.push(record);
        }

        records
    }};
}

/// Errors that may occur when interacting with database APIs.
#[derive(Debug, Error)]
#[error(transparent)]
pub struct DatabaseError(DatabaseErrorRepr);

// Implement this manually so we can still conveniently use ? with errors
// that map onto variants of DatabaseErrorRepr
// (using #[from] would conflict with this).
impl<E> From<E> for DatabaseError
where
    E: Into<DatabaseErrorRepr>,
{
    fn from(error: E) -> Self {
        Self(error.into())
    }
}

/// Database errors.
#[derive(Error, Debug)]
enum DatabaseErrorRepr {
    /// Error when creating the database connection pool.
    #[error("error creating the DB pool: {0}")]
    CreatePool(#[from] deadpool_sqlite::CreatePoolError),

    /// Error when interacting with the database pool.
    // This error has a string payload and a manual implementation of [`From`]
    // because `deadpool_sqlite::InteractError` isn't Sync.
    #[error("error interacting with the DB pool: {0}")]
    PoolInteract(String),

    /// Error when getting a new connection from the database pool.
    #[error("error getting connection from the DB pool: {0}")]
    Pool(#[from] deadpool_sqlite::PoolError),

    /// Any unspecified error returned by [`rusqlite`].
    #[error("error in rusqlite: {0}")]
    Rusqlite(#[from] rusqlite::Error),

    /// Error returned when a query function was given too many
    /// parameters to bind.
    #[error(
        "The caller passed too many parameters to bind (got {0}, max allowed is {MAX_BIND_PARAMS})"
    )]
    TooManyBindParameters(usize),
}

impl From<deadpool_sqlite::InteractError> for DatabaseErrorRepr {
    fn from(error: deadpool_sqlite::InteractError) -> Self {
        Self::PoolInteract(error.to_string())
    }
}

/// Alias for results returned by the database APIs.
pub type Result<T> = std::result::Result<T, DatabaseError>;

impl TryFrom<&Row<'_>> for DbRecord {
    type Error = rusqlite::Error;

    fn try_from(row: &Row) -> std::result::Result<Self, Self::Error> {
        Ok(Self {
            id: row.get(0)?,
            url: row.get(1)?,
            client_datetime: row.get(2)?,
            client_datetime_raw: row.get(3)?,
            timestamp_flora: row.get(4)?,
            title: row.get(5)?,
            language: row.get(6)?,
            web_archive_serialized: row.get(7)?,
        })
    }
}

/// Abstraction for flora's primary datastore (SQLite).
#[derive(Debug, Clone)]
pub struct Database {
    /// The connection pool.
    pool: Pool,
}

impl Database {
    /// Connects to the SQLite database at the given URL and creates a new `Database`.
    ///
    /// ## Errors
    /// Returns an instance of [`DatabaseError`] on error.
    #[instrument(skip(path))]
    async fn new(path: impl AsRef<Path>) -> Result<Self> {
        let cfg = Config::new(path.as_ref());
        let pool = cfg.create_pool(Runtime::Tokio1)?;
        Ok(Self { pool })
    }

    /// Gets a connection, and calls [`rusqlite::vtab::array::load_module`] with it.
    ///
    /// This method should be called to get any connection that will be used
    /// to perform queries using the `array` extension of [`rusqlite`].
    ///
    /// See <https://docs.rs/rusqlite/latest/rusqlite/vtab/array/index.html>.
    #[instrument(skip(self))]
    async fn get_conn_array(&self) -> Result<Object> {
        let conn = self.pool.get().await?;
        conn.interact(|conn| {
            tracing::debug!("Loading rusqlite::vtab::array");
            rusqlite::vtab::array::load_module(conn)
        })
        .await?
        .inspect(|()| tracing::debug!("Loaded rusqlite::vtab::array"))?;

        Ok(conn)
    }

    /// Creates a table to store the records.
    #[instrument(skip(self))]
    async fn create_table(&self, schema: &Schema) -> Result<()> {
        let create_table_query = format!("CREATE TABLE IF NOT EXISTS {schema};");
        let conn = self.pool.get().await?;
        let _rows: usize = conn
            .interact(move |conn| {
                let mut stmt = conn.prepare(&create_table_query)?;
                stmt.execute(())
            })
            .await??;

        Ok(())
    }

    /// Creates indexes for the `records` table.
    #[instrument(skip(self))]
    async fn create_indexes(&self, schema: &Schema, columns: &[&str]) -> Result<()> {
        let table_name = schema.table_name();
        let conn = self.pool.get().await?;

        for column in columns {
            let idx_name = format!("idx_{column}");
            let query =
                format!("CREATE INDEX IF NOT EXISTS {idx_name} ON {table_name} ({column});");
            tracing::debug!("Query: {query}");

            tracing::info!("Creating index '{idx_name}' on '{table_name}'");
            conn.interact(move |conn| {
                let mut stmt = conn.prepare(&query)?;
                stmt.execute([])
            })
            .await??;
        }

        Ok(())
    }

    /// Initializes the database with a 'records' table.
    ///
    /// ## Errors
    /// Returns [`DatabaseError`] if the operation fails.
    #[instrument(skip(path))]
    pub async fn new_with_records_table(path: impl AsRef<Path>) -> Result<Self> {
        let db = Self::new(path).await?;
        let schema = DbRecord::schema();
        db.create_table(&schema).await?;
        db.create_indexes(&schema, &["url"]).await?;
        Ok(db)
    }

    /// Queries the database to get a count of records it contains.
    ///
    /// ## Errors
    /// Returns [`DatabaseError`] if the operation fails.
    #[instrument]
    pub async fn count(&self) -> Result<i64> {
        let query = r"SELECT count(*) FROM records;";
        let conn = self.pool.get().await?;

        let record_count = conn
            .interact(move |conn| {
                let mut stmt = conn.prepare(query)?;
                stmt.query_one([], |row| row.get(0))
            })
            .await??;

        Ok(record_count)
    }

    /// Inserts the given record into the `Database`.
    ///
    /// Returns the ID of the inserted row.
    ///
    /// ## Errors
    /// Returns [`DatabaseError`] if the insert operation fails.
    pub async fn insert(&self, record: DbRecord) -> Result<i64> {
        let query = r"INSERT INTO records
        VALUES (:id, :url, :client_datetime, :client_datetime_raw, :timestamp_flora, :title,
        :language, :web_archive_serialized)
        RETURNING id;";

        let conn = self.pool.get().await?;

        let row_id = conn
            .interact(move |conn| {
                let mut stmt = conn.prepare(query)?;
                stmt.query_one(
                    rusqlite::named_params! {
                        ":id": record.id,
                        ":url": record.url,
                        ":client_datetime": record.client_datetime,
                        ":client_datetime_raw": record.client_datetime_raw,
                        ":timestamp_flora": record.timestamp_flora,
                        ":title": record.title,
                        ":language": record.language,
                        ":web_archive_serialized": record.web_archive_serialized,
                    },
                    |row| row.get(0),
                )
            })
            .await??;

        Ok(row_id)
    }

    /// Deletes the record corresponding to the given id.
    ///
    /// ## Errors
    /// Returns [`DatabaseError`] if the operation fails.
    #[instrument(skip(self))]
    pub async fn delete(&self, id: i64) -> Result<usize> {
        const QUERY: &str = r"DELETE FROM records WHERE id = ?1;";

        let conn = self.pool.get().await?;
        let rows_changed = conn
            .interact(move |conn| {
                let mut stmt = conn.prepare(QUERY)?;
                stmt.execute([id])
            })
            .await??;

        Ok(rows_changed)
    }

    /// Get a record from SQLite given an id.
    ///
    /// ## Errors
    /// Returns [`DatabaseError`] if the operation fails.
    pub async fn get_record(&self, id: i64) -> Result<DbRecord> {
        let query = r"SELECT id, url, client_datetime, client_datetime_raw, timestamp_flora, title,
            language, web_archive_serialized
            FROM records WHERE id = ?1;";

        let conn = self.pool.get().await?;
        let result = conn
            .interact(move |conn| {
                let mut stmt = conn.prepare(query)?;
                stmt.query_row([id], |row| {
                    let record = row.try_into()?;
                    Ok(record)
                })
            })
            .await??;

        Ok(result)
    }

    /// Get multiple records from SQLite given their IDs.
    ///
    /// ## Errors
    /// Returns [`DatabaseError`] if the operation fails.
    #[instrument]
    pub async fn get_many_records(&self, ids: &[i64]) -> Result<Vec<DbRecord>> {
        if ids.len() > MAX_BIND_PARAMS {
            return Err(DatabaseErrorRepr::TooManyBindParameters(ids.len()).into());
        }

        tracing::debug!("Getting records.");

        let query = r"SELECT id, url, client_datetime, client_datetime_raw, timestamp_flora, title,
            language, web_archive_serialized
            FROM records
            WHERE id IN rarray(?1);";

        let conn = self
            .get_conn_array()
            .await
            .inspect_err(|e| tracing::error!("Couldn't get a connection from the pool: {e}"))?;

        let values = ids.iter().copied().map(Value::from).collect::<Vec<Value>>();

        let records = conn
            .interact(|conn| -> Result<Vec<DbRecord>> {
                let values = Rc::new(values);

                let mut stmt = conn
                    .prepare(query)
                    .inspect_err(|e| tracing::error!("Couldn't prepare the query: {e}"))?;

                let records = try_query_db_records!(stmt, [values]);
                Ok(records)
            })
            .await
            .inspect_err(|e| {
                tracing::error!("Couldn't interact with the database connection: {e}");
            })??;

        tracing::debug!("Got {} records", records.len());

        Ok(records)
    }

    /// Selects all rows in the `records` table, without any pagination.
    ///
    /// ## Errors
    ///
    /// Returns [`DatabaseError`] if the operation fails.
    #[instrument(skip(self))]
    pub async fn get_all_records(&self) -> Result<Vec<DbRecord>> {
        tracing::debug!("Getting all records");

        let query = r"SELECT * FROM records";
        let conn = self.pool.get().await?;
        let records = conn
            .interact(move |conn| -> Result<Vec<DbRecord>> {
                let mut stmt = conn.prepare(query)?;
                let records = try_query_db_records!(stmt, []);
                Ok(records)
            })
            .await??;

        Ok(records)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::model::tests::test_record_params;
    use anyhow::Result;
    use chrono::{DateTime, FixedOffset};
    use omry_archiving::record::Record;
    use pretty_assertions::assert_eq;
    use serde::Serialize;
    use std::fs::File;
    use std::ops::Deref;
    use tempfile::NamedTempFile;

    /// Data structure to keep the temporary file holding the database alive,
    /// because it will get unlinked on drop.
    struct TempDb {
        #[allow(dead_code)]
        temp_db_file: NamedTempFile<File>,
        db: Database,
    }

    impl TempDb {
        async fn new() -> Result<Self> {
            let temp_db_file = NamedTempFile::with_suffix(".db")?;
            let db = Database::new_with_records_table(temp_db_file.path()).await?;
            Ok(Self { temp_db_file, db })
        }
    }

    impl Deref for TempDb {
        type Target = Database;

        fn deref(&self) -> &Self::Target {
            &self.db
        }
    }

    /// Wrapper for [`DbRecord`] that supports [`Serialize`] for redacting snapshots.
    #[derive(Debug, Serialize)]
    struct DbRecordDef {
        id: Option<i64>,
        url: String,
        client_datetime: DateTime<FixedOffset>,
        client_datetime_raw: Option<String>,
        timestamp_flora: i64,
        title: String,
        language: Option<String>,
        web_archive_serialized: Vec<u8>,
    }

    impl From<DbRecord> for DbRecordDef {
        fn from(value: DbRecord) -> DbRecordDef {
            let DbRecord {
                id,
                url,
                client_datetime,
                client_datetime_raw,
                timestamp_flora,
                title,
                language,
                web_archive_serialized,
            } = value;

            Self {
                id,
                url,
                client_datetime,
                client_datetime_raw,
                timestamp_flora,
                title,
                language,
                web_archive_serialized,
            }
        }
    }

    fn test_db_record() -> Result<DbRecord> {
        let record = Record::new(test_record_params()?);
        Ok(record.try_into()?)
    }

    #[tokio::test]
    async fn can_create_table() -> Result<()> {
        let db = TempDb::new().await?;
        assert!(db.create_table(&DbRecord::schema()).await.is_ok());
        // TODO: Check the resulting DB schema matches what we gave to create_table
        Ok(())
    }

    #[tokio::test]
    async fn can_add_record() -> Result<()> {
        let db = TempDb::new().await?;
        let record = test_db_record()?;
        let result = db.insert(record).await?;
        assert_eq!(result, 1);
        Ok(())
    }

    #[tokio::test]
    async fn can_get_record() -> Result<()> {
        let db = TempDb::new().await?;
        let record = test_db_record()?;
        assert!(record.id.is_none());
        db.insert(record).await?;

        let expected_id = 1i64;
        let retrieved_record = db.get_record(expected_id).await?;

        let retrieved_record: DbRecordDef = retrieved_record.into();
        insta::assert_ron_snapshot!(retrieved_record, {
            ".client_datetime" => "[client_datetime]",
            ".timestamp_flora" => "[timestamp]"
        });

        Ok(())
    }

    #[tokio::test]
    async fn can_get_many_records() -> Result<()> {
        let db = TempDb::new().await?;

        let num_records_to_insert = 7;
        let mut record_ids = Vec::with_capacity(num_records_to_insert);

        for _ in 0..num_records_to_insert {
            let record = test_db_record()?;
            let record_id = db.insert(record).await?;
            record_ids.push(record_id);
        }

        let selected_records = db
            .get_many_records(&record_ids)
            .await?
            .into_iter()
            .map(DbRecordDef::from)
            .collect::<Vec<_>>();

        assert_eq!(selected_records.len(), num_records_to_insert);

        insta::assert_ron_snapshot!(selected_records, {
            "[].client_datetime" => "[client_datetime]",
            "[].timestamp_flora" => "[timestamp]"
        });

        Ok(())
    }

    #[tokio::test]
    async fn passing_too_many_bind_parameters_returns_error() -> Result<()> {
        let db = TempDb::new().await?;

        let num_records_to_insert = MAX_BIND_PARAMS + 1;
        db.insert(test_db_record()?).await?;

        let mut ids = Vec::with_capacity(num_records_to_insert);
        for index in 0..num_records_to_insert {
            ids.push(i64::try_from(index + 1)?);
        }

        let result = db.get_many_records(&ids).await;
        assert!(result.is_err());

        // At this point, we know this is definitely an error.
        #[allow(clippy::unwrap_used)]
        let error = result.unwrap_err();
        insta::assert_snapshot!(error);

        Ok(())
    }

    #[tokio::test]
    async fn can_get_count() -> Result<()> {
        let db = TempDb::new().await?;

        let num_records_to_insert = 7;
        let mut record_ids = Vec::with_capacity(num_records_to_insert);

        for _ in 0..num_records_to_insert {
            let record = test_db_record()?;
            let record_id = db.insert(record).await?;
            record_ids.push(record_id);
        }

        let count: usize = db.count().await?.try_into()?;

        assert_eq!(count, num_records_to_insert);
        Ok(())
    }

    #[tokio::test]
    async fn can_delete_record() -> Result<()> {
        let db = TempDb::new().await?;
        let record = test_db_record()?;
        let id = db.insert(record).await?;
        let rows_changed = db.delete(id).await?;
        assert_eq!(rows_changed, 1);
        Ok(())
    }

    #[tokio::test]
    async fn can_get_all_records() -> Result<()> {
        let db = TempDb::new().await?;

        let num_records_to_insert = 7;
        let mut record_ids = Vec::with_capacity(num_records_to_insert);

        for _ in 0..num_records_to_insert {
            let record = test_db_record()?;
            let record_id = db.insert(record).await?;
            record_ids.push(record_id);
        }

        let selected_records = db
            .get_all_records()
            .await?
            .into_iter()
            .map(DbRecordDef::from)
            .collect::<Vec<_>>();

        assert_eq!(selected_records.len(), num_records_to_insert);

        insta::assert_ron_snapshot!(selected_records, {
            "[].client_datetime" => "[client_datetime]",
            "[].timestamp_flora" => "[timestamp]",
            "[].web_archive_serialized" => "[data]"
        });

        Ok(())
    }
}