omry_db/

db.rs

//! Database abstractions.
#![cfg(feature = "connection")]

use deadpool_sqlite::rusqlite;
use deadpool_sqlite::rusqlite::Row;
use deadpool_sqlite::rusqlite::types::Value;
use deadpool_sqlite::{Config, Object, Pool, Runtime};
use rusqlite_migration::{M, Migrations};
use std::path::Path;
use std::rc::Rc;
use thiserror::Error;
use tracing::instrument;

use crate::model::{DbRecord, DbRecordMeta};

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

/// Limit used when selecting all records.
const SELECT_ALL_LIMIT: i64 = 50;

/// 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, $type:ty) => {{
        let rows = $stmt
            .query_map($values, |row| <$type>::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
    }};

    ($stmt:ident, $values:expr) => {
        try_query_db_records!($stmt, $values, DbRecord)
    };
}

const MIGRATIONS_SLICE: &[M<'_>] = &[
    M::up(
        r"CREATE TABLE IF NOT EXISTS records(id INTEGER NOT NULL PRIMARY KEY,
url TEXT NOT NULL,
client_datetime TEXT NOT NULL,
client_datetime_raw TEXT,
timestamp_flora INTEGER NOT NULL,
title TEXT NOT NULL,
language TEXT,
web_archive_serialized BLOB NOT NULL);",
    ),
    M::up("ALTER TABLE records ADD COLUMN updated_at TEXT;"),
    M::up("CREATE INDEX timestamp ON records(timestamp_flora);"),
];
const MIGRATIONS: Migrations<'_> = Migrations::from_slice(MIGRATIONS_SLICE);

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

// 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<DbErrorRepr>,
{
    fn from(error: E) -> Self {
        Self(error.into())
    }
}

/// Database errors.
#[derive(Error, Debug)]
enum DbErrorRepr {
    /// 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),

    /// Error activating write-ahead logging.
    #[error("Couldn't activate the WAL; journaling mode is {0}")]
    WalActivation(String),

    #[error("error performing migration: {0}")]
    Migration(#[from] rusqlite_migration::Error),

    #[error("error performing record update: {0}")]
    Update(String),
}

impl From<deadpool_sqlite::InteractError> for DbErrorRepr {
    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)?,
            document: row.get(7)?,
            updated_at: row.get(8)?,
        })
    }
}

impl TryFrom<&Row<'_>> for DbRecordMeta {
    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)?,
            timestamp_flora: row.get(3)?,
            title: row.get(4)?,
            language: row.get(5)?,
            updated_at: row.get(6)?,
        })
    }
}
/// 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) -> Result<()> {
        let conn = self.pool.get().await?;
        conn.interact(move |conn| MIGRATIONS.to_latest(conn))
            .await??;

        Ok(())
    }

    /// Enables the WAL
    ///
    /// See <https://sqlite.org/wal.html>
    ///
    /// # Errors
    ///
    /// Returns [`DatabaseError`] if WAL activation fails.
    #[instrument(skip(self))]
    async fn activate_wal(&self) -> Result<()> {
        const SUCCESS_VALUE: &str = "wal";

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

        tracing::debug!("Enabling the WAL.");

        let journal_mode: String = conn
            .interact(move |conn| {
                conn.pragma_update_and_check(None, "journal_mode", "WAL", |v| v.get(0))
            })
            .await??;

        if journal_mode.as_str() == SUCCESS_VALUE {
            Ok(())
        } else {
            Err(DbErrorRepr::WalActivation(journal_mode).into())
        }
    }

    /// Initializes the database, activates WAL, creates a 'records' table, and creates indexes.
    ///
    /// # 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?;
        db.activate_wal().await?;
        db.create_table().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.
    #[instrument(skip_all, fields(record_url = record.url))]
    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, :updated_at)
        RETURNING id;";

        let DbRecord {
            id,
            url,
            client_datetime,
            client_datetime_raw,
            timestamp_flora,
            title,
            language,
            document,
            updated_at,
        } = record;

        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": id,
                        ":url": url,
                        ":client_datetime": client_datetime,
                        ":client_datetime_raw": client_datetime_raw,
                        ":timestamp_flora": timestamp_flora,
                        ":title": title,
                        ":language": language,
                        ":web_archive_serialized": document,
                        ":updated_at": updated_at,
                    },
                    |row| row.get(0),
                )
            })
            .await?
            .inspect_err(|e| tracing::error!("Error on insert: {e}"))?;

        Ok(row_id)
    }

    /// Deletes the record corresponding to the given id.
    ///
    /// Returns the number of documents affected
    /// (should be 1 if a document matching the id was deleted,
    /// and 0 if no matching document was found in the DB).
    ///
    /// # 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, updated_at
            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(DbErrorRepr::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, updated_at
            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, start_row: Option<i64>) -> Result<Vec<DbRecord>> {
        tracing::debug!("Getting all records");

        // pagination adapted from <https://gist.github.com/ssokolow/262503>
        let query = r"SELECT * FROM records
WHERE id NOT IN (SELECT id FROM records
                 ORDER BY id ASC LIMIT ?1)
ORDER BY id ASC LIMIT ?2;";

        // start at 0 if we aren't given an explicit start row
        let start_row = start_row.unwrap_or_default();
        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, [start_row, SELECT_ALL_LIMIT]);
                Ok(records)
            })
            .await??;

        Ok(records)
    }

    /// Selects all rows in the `records` table, without any pagination,
    /// and returns their metadata (all fields except serialized data
    /// of the page content).
    ///
    /// # Errors
    ///
    /// Returns [`DatabaseError`] if the operation fails.
    #[instrument(skip(self))]
    pub async fn get_all_meta(&self, start_row: Option<i64>) -> Result<Vec<DbRecordMeta>> {
        tracing::debug!("Getting metadata for all records");

        let query = r"
SELECT id, url, client_datetime, timestamp_flora, title, language, updated_at
FROM records
WHERE id NOT IN (SELECT id FROM records
                 ORDER BY id ASC LIMIT ?1)
ORDER BY id ASC LIMIT ?2;";

        let start_row = start_row.unwrap_or_default();
        let conn = self.pool.get().await?;
        let records = conn
            .interact(move |conn| -> Result<Vec<DbRecordMeta>> {
                let mut stmt = conn.prepare(query)?;
                let records =
                    try_query_db_records!(stmt, [start_row, SELECT_ALL_LIMIT], DbRecordMeta);
                Ok(records)
            })
            .await??;

        Ok(records)
    }

    /// Returns all record IDs in SQLite.
    ///
    /// # Errors
    ///
    /// Returns [`DatabaseError`] if the operation fails.
    #[instrument(skip(self))]
    pub async fn get_all_ids(&self) -> Result<Vec<i64>> {
        tracing::debug!("Getting all record IDs");

        let query = r"SELECT id FROM records;";

        let conn = self.pool.get().await?;
        let ids = conn
            .interact(move |conn| -> Result<Vec<i64>> {
                let mut stmt = conn.prepare(query)?;
                let rows = stmt.query_map([], |row| row.get(0))?;
                let mut ids = vec![];
                for row in rows {
                    ids.push(row?);
                }

                Ok(ids)
            })
            .await??;

        Ok(ids)
    }

    /// Returns all record IDs in descending order.
    #[instrument(skip(self))]
    async fn get_all_ids_by_recent(&self) -> Result<Vec<i64>> {
        // order by timestamp because in rare cases IDs might not
        // be monotonically increasing
        // (see <https://sqlite.org/autoinc.html>)
        let query = r"SELECT id FROM records
ORDER BY timestamp_flora DESC;";

        let conn = self.pool.get().await?;
        let ids = conn
            .interact(move |conn| -> Result<Vec<i64>> {
                let mut stmt = conn.prepare(query)?;
                let rows = stmt.query_map([], |row| row.get(0))?;
                let mut ids = vec![];
                for row in rows {
                    ids.push(row?);
                }

                Ok(ids)
            })
            .await??;

        Ok(ids)
    }

    /// Similar to [`Database::get_all_meta`], except orders records by
    /// descending order of `timestamp_flora`.
    ///
    /// Returns an empty list if the given index is out of bounds.
    #[allow(clippy::missing_errors_doc)]
    #[instrument(skip(self))]
    pub async fn most_recent_meta(&self, start_idx: Option<usize>) -> Result<Vec<DbRecordMeta>> {
        tracing::trace!("Selecting recent RecordMeta");

        let mut ids = self.get_all_ids_by_recent().await?;
        let remaining_ids = if let Some(start_idx) = start_idx {
            let start_id = ids.get(start_idx);

            let Some(start_id) = start_id else {
                tracing::trace!("requested index is out of bounds");
                return Ok(vec![]);
            };

            // all IDs we haven't yet seen will be to the left of this point
            let cur_point = ids.partition_point(|id| id > start_id);
            ids.split_off(cur_point)
        } else {
            ids
        };

        let ids = remaining_ids
            .iter()
            .copied()
            .map(Value::from)
            .collect::<Vec<Value>>();
        let query = r"
SELECT id, url, client_datetime, timestamp_flora, title, language, updated_at
FROM records
WHERE id IN rarray(?1)
ORDER BY timestamp_flora DESC LIMIT ?2;";

        let conn = self.get_conn_array().await?;
        let records = conn
            .interact(move |conn| -> Result<Vec<DbRecordMeta>> {
                let ids = Rc::new(ids);
                let mut stmt = conn.prepare(query)?;
                let records = try_query_db_records!(
                    stmt,
                    rusqlite::params![ids, SELECT_ALL_LIMIT],
                    DbRecordMeta
                );
                Ok(records)
            })
            .await??;

        Ok(records)
    }

    /// Updates the record by ID.
    ///
    /// # Errors
    /// Returns [`DatabaseError`] if the update operation fails.
    #[instrument(skip_all, fields(record_id = record.id))]
    pub async fn update_record_data(&self, record: DbRecord) -> Result<i64> {
        let query = r"UPDATE records
SET web_archive_serialized = :document, updated_at = :updated_at
WHERE records.id = :id
RETURNING id;";

        let DbRecord {
            id,
            document,
            updated_at,
            ..
        } = record;
        let id = id.ok_or_else(|| DbErrorRepr::Update("missing id".into()))?;
        let updated_at =
            updated_at.ok_or_else(|| DbErrorRepr::Update("missing updated_at".into()))?;

        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": id,
                        ":document": document,
                        ":updated_at": updated_at,
                    },
                    |row| row.get(0),
                )
            })
            .await?
            .inspect_err(|e| tracing::error!("Error on update: {e}"))?;
        Ok(row_id)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::model::tests::test_record_params;
    use anyhow::{Context, Result};
    use chrono::{DateTime, FixedOffset, Utc};
    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>,
        document: Vec<u8>,
        updated_at: Option<i64>,
    }

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

            Self {
                id,
                url,
                client_datetime,
                client_datetime_raw,
                timestamp_flora,
                title,
                language,
                document,
                updated_at,
            }
        }
    }

    #[derive(Debug, Serialize)]
    struct DbRecordMetaDef {
        id: i64,
        url: String,
        client_datetime: DateTime<FixedOffset>,
        timestamp_flora: i64,
        title: String,
        language: Option<String>,
        updated_at: Option<i64>,
    }

    impl From<DbRecordMeta> for DbRecordMetaDef {
        fn from(value: DbRecordMeta) -> Self {
            let DbRecordMeta {
                id,
                url,
                client_datetime,
                timestamp_flora,
                title,
                language,
                updated_at,
            } = value;

            Self {
                id,
                url,
                client_datetime,
                timestamp_flora,
                title,
                language,
                updated_at,
            }
        }
    }

    #[allow(dead_code, reason = "Used in snapshot.")]
    #[derive(Debug)]
    struct RowInfo {
        name: String,
        type_: String,
        notnull: bool,
    }

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

        fn try_from(row: &Row) -> std::result::Result<Self, Self::Error> {
            let (name, type_, notnull) = (row.get(0)?, row.get(1)?, row.get(2)?);
            Ok(Self {
                name,
                type_,
                notnull,
            })
        }
    }

    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?;
        let query = r#"SELECT name, type, "notnull" FROM pragma_table_info('records');"#;
        let conn = db.pool.get().await?;
        let row_infos = conn
            .interact(move |conn| -> Result<Vec<RowInfo>> {
                let mut stmt = conn.prepare(query)?;
                let rows = stmt.query_map([], |row| RowInfo::try_from(row))?;
                let mut values = vec![];
                for row in rows {
                    values.push(row?);
                }
                Ok(values)
            })
            .await
            .map_err(|e| anyhow::anyhow!("{e}"))??;
        insta::assert_debug_snapshot!(row_infos);
        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?;

        // Insert more than the limit to test pagination.
        let num_records_to_insert = (SELECT_ALL_LIMIT + 5).try_into()?;
        for _ in 0..num_records_to_insert {
            let record = test_db_record()?;
            let _record_id = db.insert(record).await?;
        }

        let mut db_recs = Vec::with_capacity(num_records_to_insert);
        let mut start_row = None;
        loop {
            let mut selected_records = db
                .get_all_records(start_row)
                .await?
                .into_iter()
                .map(DbRecordDef::from)
                .collect::<Vec<_>>();

            if selected_records.is_empty() {
                break;
            }

            db_recs.append(&mut selected_records);
            start_row = Some(db_recs.len().try_into()?);
        }

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

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

        Ok(())
    }

    /// This tests that we can still get all records
    /// after deleting some from the middle (causing IDs to become incontiguous).
    #[tokio::test]
    async fn can_get_all_records_after_deleting_some() -> Result<()> {
        let db = TempDb::new().await?;

        // Insert more than the limit to test pagination.
        let num_records_to_insert = (SELECT_ALL_LIMIT + 5).try_into()?;
        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);
        }

        // remove IDs from `record_ids` to assert that
        // `record_ids` length matches the number of records
        // we get back in the end
        let idx_to_delete = [3, 7, 11];
        for idx in idx_to_delete {
            let id = record_ids.remove(idx);
            db.delete(id).await?;
        }

        let mut db_recs = Vec::with_capacity(num_records_to_insert);
        let mut start_row = None;
        loop {
            let mut selected_records = db
                .get_all_records(start_row)
                .await?
                .into_iter()
                .map(DbRecordDef::from)
                .collect::<Vec<_>>();

            if selected_records.is_empty() {
                break;
            }

            db_recs.append(&mut selected_records);
            start_row = Some(db_recs.len().try_into()?);
        }

        assert_eq!(db_recs.len(), record_ids.len());
        insta::assert_ron_snapshot!(db_recs, {
            "[].client_datetime" => "[client_datetime]",
            "[].timestamp_flora" => "[timestamp]",
            "[].document" => "[data]"
        });

        Ok(())
    }

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

        let num_records_to_insert = (SELECT_ALL_LIMIT + 1).try_into()?;
        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 mut start_row = None;
        let mut db_recs = Vec::with_capacity(num_records_to_insert);
        loop {
            let mut selected_recs = db
                .get_all_meta(start_row)
                .await?
                .into_iter()
                .map(DbRecordMetaDef::from)
                .collect::<Vec<_>>();

            if selected_recs.is_empty() {
                break;
            }

            db_recs.append(&mut selected_recs);
            start_row = Some(db_recs.len().try_into()?);
        }

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

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

        Ok(())
    }

    #[tokio::test]
    async fn can_get_all_ids() -> 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 db_ids = db.get_all_ids().await?;
        assert_eq!(db_ids, record_ids);
        Ok(())
    }

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

        let num_recs_to_insert = (SELECT_ALL_LIMIT + 5).try_into()?;
        let ts_incr = 5;
        let mut record_ids = Vec::with_capacity(num_recs_to_insert);

        let mut cur_timestamp = Utc::now().timestamp();
        for _ in 0..num_recs_to_insert {
            let mut rec = test_db_record()?;
            rec.timestamp_flora = cur_timestamp;
            let record_id = db.insert(rec).await?;
            record_ids.push(record_id);
            cur_timestamp += ts_incr;
        }

        let mut recs = Vec::with_capacity(num_recs_to_insert);
        let mut start_idx = None;
        loop {
            let mut selected_recs = db
                .most_recent_meta(start_idx)
                .await?
                .into_iter()
                .map(DbRecordMetaDef::from)
                .collect::<Vec<_>>();

            if selected_recs.is_empty() {
                break;
            }

            recs.append(&mut selected_recs);
            start_idx = Some(recs.len());
        }

        assert_eq!(recs.len(), num_recs_to_insert);
        let timestamps = recs.iter().map(|r| r.timestamp_flora).collect::<Vec<_>>();
        assert!(
            recs.is_sorted_by(|a, b| a.timestamp_flora > b.timestamp_flora),
            "{timestamps:#?}"
        );
        Ok(())
    }

    #[tokio::test]
    async fn can_update_record() -> Result<()> {
        let db = TempDb::new().await?;
        let record = test_db_record()?;
        let id = db.insert(record).await.context("error inserting record")?;

        // workaround for DbRecord not being `Clone`
        let mut record = db.get_record(id).await.context("error getting record")?;
        record.updated_at = Some(Utc::now().timestamp());
        let updated_id = db
            .update_record_data(record)
            .await
            .context("error updating record data")?;
        assert_eq!(id, updated_id);
        Ok(())
    }
}