Skip to main content

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 rusqlite_migration::{M, Migrations};
9use std::path::Path;
10use std::rc::Rc;
11use thiserror::Error;
12use tracing::instrument;
13
14use crate::model::{DbRecord, DbRecordMeta};
15
16/// Maximum number of 'bound' (programmatically populated)
17/// parameters allowed in a query.
18const MAX_BIND_PARAMS: usize = 32766;
19
20/// Limit used when selecting all records.
21const SELECT_ALL_LIMIT: i64 = 50;
22
23/// Tries to query for rows in `records` using the given statement, and convert them to [`DbRecord`].
24macro_rules! try_query_db_records {
25    ($stmt:ident, $values:expr, $type:ty) => {{
26        $stmt
27            .query_map($values, |row| <$type>::try_from(row))
28            .inspect_err(|e| tracing::error!("Couldn't execute the query: {e}"))?
29            .collect::<std::result::Result<Vec<$type>, rusqlite::Error>>()?
30    }};
31
32    ($stmt:ident, $values:expr) => {
33        try_query_db_records!($stmt, $values, DbRecord)
34    };
35}
36
37const MIGRATIONS_SLICE: &[M<'_>] = &[
38    M::up(
39        r"CREATE TABLE IF NOT EXISTS records(id INTEGER NOT NULL PRIMARY KEY,
40url TEXT NOT NULL,
41client_datetime TEXT NOT NULL,
42client_datetime_raw TEXT,
43timestamp_flora INTEGER NOT NULL,
44title TEXT NOT NULL,
45language TEXT,
46web_archive_serialized BLOB NOT NULL);",
47    ),
48    M::up("ALTER TABLE records ADD COLUMN updated_at INTEGER;"),
49    M::up("CREATE INDEX timestamp ON records(timestamp_flora);"),
50];
51const MIGRATIONS: Migrations<'_> = Migrations::from_slice(MIGRATIONS_SLICE);
52
53/// Errors that may occur when interacting with database APIs.
54#[derive(Debug, Error)]
55#[error(transparent)]
56pub struct DatabaseError(DbErrorRepr);
57
58// Implement this manually so we can still conveniently use ? with errors
59// that map onto variants of DatabaseErrorRepr
60// (using #[from] would conflict with this).
61impl<E> From<E> for DatabaseError
62where
63    E: Into<DbErrorRepr>,
64{
65    fn from(error: E) -> Self {
66        Self(error.into())
67    }
68}
69
70/// Database errors.
71#[derive(Error, Debug)]
72enum DbErrorRepr {
73    /// Error when creating the database connection pool.
74    #[error("error creating the DB pool: {0}")]
75    CreatePool(#[from] deadpool_sqlite::CreatePoolError),
76
77    /// Error when interacting with the database pool.
78    // This error has a string payload and a manual implementation of [`From`]
79    // because `deadpool_sqlite::InteractError` isn't Sync.
80    #[error("error interacting with the DB pool: {0}")]
81    PoolInteract(String),
82
83    /// Error when getting a new connection from the database pool.
84    #[error("error getting connection from the DB pool: {0}")]
85    Pool(#[from] deadpool_sqlite::PoolError),
86
87    /// Any unspecified error returned by [`rusqlite`].
88    #[error("error in rusqlite: {0}")]
89    Rusqlite(#[from] rusqlite::Error),
90
91    /// Error returned when a query function was given too many
92    /// parameters to bind.
93    #[error(
94        "The caller passed too many parameters to bind (got {0}, max allowed is {MAX_BIND_PARAMS})"
95    )]
96    TooManyBindParameters(usize),
97
98    /// Error activating write-ahead logging.
99    #[error("Couldn't activate the WAL; journaling mode is {0}")]
100    WalActivation(String),
101
102    #[error("error performing migration: {0}")]
103    Migration(#[from] rusqlite_migration::Error),
104
105    #[error("error performing record update: {0}")]
106    Update(String),
107
108    #[error("error converting int: {0}")]
109    IntConversion(#[from] std::num::TryFromIntError),
110}
111
112impl From<deadpool_sqlite::InteractError> for DbErrorRepr {
113    fn from(error: deadpool_sqlite::InteractError) -> Self {
114        Self::PoolInteract(error.to_string())
115    }
116}
117
118/// Alias for results returned by the database APIs.
119pub type Result<T> = std::result::Result<T, DatabaseError>;
120
121impl TryFrom<&Row<'_>> for DbRecord {
122    type Error = rusqlite::Error;
123
124    fn try_from(row: &Row) -> std::result::Result<Self, Self::Error> {
125        Ok(Self {
126            id: row.get(0)?,
127            url: row.get(1)?,
128            client_datetime: row.get(2)?,
129            client_datetime_raw: row.get(3)?,
130            timestamp_flora: row.get(4)?,
131            title: row.get(5)?,
132            language: row.get(6)?,
133            document: row.get(7)?,
134            updated_at: row.get(8)?,
135        })
136    }
137}
138
139impl TryFrom<&Row<'_>> for DbRecordMeta {
140    type Error = rusqlite::Error;
141
142    fn try_from(row: &Row) -> std::result::Result<Self, Self::Error> {
143        Ok(Self {
144            id: row.get(0)?,
145            url: row.get(1)?,
146            client_datetime: row.get(2)?,
147            timestamp_flora: row.get(3)?,
148            title: row.get(4)?,
149            language: row.get(5)?,
150            updated_at: row.get(6)?,
151        })
152    }
153}
154/// Abstraction for flora's primary datastore (SQLite).
155#[derive(Debug, Clone)]
156pub struct Database {
157    /// The connection pool.
158    pool: Pool,
159}
160
161impl Database {
162    /// Connects to the SQLite database at the given URL and creates a new `Database`.
163    ///
164    /// # Errors
165    /// Returns an instance of [`DatabaseError`] on error.
166    #[instrument(skip(path))]
167    async fn new(path: impl AsRef<Path>) -> Result<Self> {
168        let cfg = Config::new(path.as_ref());
169        let pool = cfg.create_pool(Runtime::Tokio1)?;
170        Ok(Self { pool })
171    }
172
173    /// Gets a connection, and calls [`rusqlite::vtab::array::load_module`] with it.
174    ///
175    /// This method should be called any time you need to get
176    /// a connection that will be use the `array` extension of [`rusqlite`].
177    /// If you don't use this method, the pool may return a connection that didn't
178    /// have `array` loaded, even if you had called this method previously.
179    ///
180    ///
181    /// See <https://docs.rs/rusqlite/latest/rusqlite/vtab/array/index.html>.
182    #[instrument(skip(self))]
183    async fn get_conn_array(&self) -> Result<Object> {
184        let conn = self.pool.get().await?;
185        conn.interact(|conn| {
186            tracing::debug!("Loading rusqlite::vtab::array");
187            rusqlite::vtab::array::load_module(conn)
188        })
189        .await?
190        .inspect(|()| tracing::debug!("Loaded rusqlite::vtab::array"))?;
191
192        Ok(conn)
193    }
194
195    /// Creates a table to store the records.
196    #[instrument(skip(self))]
197    async fn create_table(&self) -> Result<()> {
198        let conn = self.pool.get().await?;
199        conn.interact(move |conn| MIGRATIONS.to_latest(conn))
200            .await??;
201
202        Ok(())
203    }
204
205    /// Enables the WAL
206    ///
207    /// See <https://sqlite.org/wal.html>
208    ///
209    /// # Errors
210    ///
211    /// Returns [`DatabaseError`] if WAL activation fails.
212    #[instrument(skip(self))]
213    async fn activate_wal(&self) -> Result<()> {
214        const SUCCESS_VALUE: &str = "wal";
215
216        let conn = self.pool.get().await?;
217
218        tracing::debug!("Enabling the WAL.");
219
220        let journal_mode: String = conn
221            .interact(move |conn| {
222                conn.pragma_update_and_check(None, "journal_mode", "WAL", |v| v.get(0))
223            })
224            .await??;
225
226        if journal_mode.as_str() == SUCCESS_VALUE {
227            Ok(())
228        } else {
229            Err(DbErrorRepr::WalActivation(journal_mode).into())
230        }
231    }
232
233    /// Initializes the database, activates WAL, creates a 'records' table, and creates indexes.
234    ///
235    /// # Errors
236    /// Returns [`DatabaseError`] if the operation fails.
237    #[instrument(skip(path))]
238    pub async fn new_with_records_table(path: impl AsRef<Path>) -> Result<Self> {
239        let db = Self::new(path).await?;
240        db.activate_wal().await?;
241        db.create_table().await?;
242        Ok(db)
243    }
244
245    /// Queries the database to get a count of records it contains.
246    ///
247    /// # Errors
248    /// Returns [`DatabaseError`] if the operation fails.
249    #[instrument]
250    pub async fn count(&self) -> Result<i64> {
251        let query = r"SELECT count(*) FROM records;";
252        let conn = self.pool.get().await?;
253
254        let record_count = conn
255            .interact(move |conn| {
256                let mut stmt = conn.prepare(query)?;
257                stmt.query_one([], |row| row.get(0))
258            })
259            .await??;
260
261        Ok(record_count)
262    }
263
264    /// Inserts the given record into the `Database`.
265    ///
266    /// Returns the ID of the inserted row.
267    ///
268    /// # Errors
269    /// Returns [`DatabaseError`] if the insert operation fails.
270    #[instrument(skip_all, fields(record_url = record.url))]
271    pub async fn insert(&self, record: DbRecord) -> Result<i64> {
272        let query = r"INSERT INTO records
273        VALUES (:id, :url, :client_datetime, :client_datetime_raw, :timestamp_flora, :title,
274        :language, :web_archive_serialized, :updated_at)
275        RETURNING id;";
276
277        let DbRecord {
278            id,
279            url,
280            client_datetime,
281            client_datetime_raw,
282            timestamp_flora,
283            title,
284            language,
285            document,
286            updated_at,
287        } = record;
288
289        let conn = self.pool.get().await?;
290        let row_id = conn
291            .interact(move |conn| {
292                let mut stmt = conn.prepare(query)?;
293                stmt.query_one(
294                    rusqlite::named_params! {
295                        ":id": id,
296                        ":url": url,
297                        ":client_datetime": client_datetime,
298                        ":client_datetime_raw": client_datetime_raw,
299                        ":timestamp_flora": timestamp_flora,
300                        ":title": title,
301                        ":language": language,
302                        ":web_archive_serialized": document,
303                        ":updated_at": updated_at,
304                    },
305                    |row| row.get(0),
306                )
307            })
308            .await?
309            .inspect_err(|e| tracing::error!("Error on insert: {e}"))?;
310
311        Ok(row_id)
312    }
313
314    /// Deletes the record corresponding to the given id.
315    ///
316    /// Returns the number of documents affected
317    /// (should be 1 if a document matching the id was deleted,
318    /// and 0 if no matching document was found in the DB).
319    ///
320    /// # Errors
321    /// Returns [`DatabaseError`] if the operation fails.
322    #[instrument(skip(self))]
323    pub async fn delete(&self, id: i64) -> Result<usize> {
324        const QUERY: &str = r"DELETE FROM records WHERE id = ?1;";
325
326        let conn = self.pool.get().await?;
327        let rows_changed = conn
328            .interact(move |conn| {
329                let mut stmt = conn.prepare(QUERY)?;
330                stmt.execute([id])
331            })
332            .await??;
333
334        Ok(rows_changed)
335    }
336
337    /// Get a record from SQLite given an id.
338    ///
339    /// # Errors
340    /// Returns [`DatabaseError`] if the operation fails.
341    pub async fn get_record(&self, id: i64) -> Result<DbRecord> {
342        let query = r"SELECT id, url, client_datetime, client_datetime_raw, timestamp_flora, title,
343            language, web_archive_serialized, updated_at
344            FROM records WHERE id = ?1;";
345
346        let conn = self.pool.get().await?;
347        let result = conn
348            .interact(move |conn| {
349                let mut stmt = conn.prepare(query)?;
350                stmt.query_row([id], |row| {
351                    let record = row.try_into()?;
352                    Ok(record)
353                })
354            })
355            .await??;
356
357        Ok(result)
358    }
359
360    /// Get multiple records from SQLite given their IDs.
361    ///
362    /// # Errors
363    /// Returns [`DatabaseError`] if the operation fails.
364    #[instrument]
365    pub async fn get_many_records(&self, ids: &[i64]) -> Result<Vec<DbRecord>> {
366        if ids.len() > MAX_BIND_PARAMS {
367            return Err(DbErrorRepr::TooManyBindParameters(ids.len()).into());
368        }
369
370        tracing::debug!("Getting records.");
371
372        let query = r"SELECT id, url, client_datetime, client_datetime_raw, timestamp_flora, title,
373            language, web_archive_serialized, updated_at
374            FROM records
375            WHERE id IN rarray(?1);";
376
377        let conn = self
378            .get_conn_array()
379            .await
380            .inspect_err(|e| tracing::error!("Couldn't get a connection from the pool: {e}"))?;
381
382        let values = ids.iter().copied().map(Value::from).collect::<Vec<Value>>();
383
384        let records = conn
385            .interact(|conn| -> Result<Vec<DbRecord>> {
386                let values = Rc::new(values);
387
388                let mut stmt = conn
389                    .prepare(query)
390                    .inspect_err(|e| tracing::error!("Couldn't prepare the query: {e}"))?;
391
392                let records = try_query_db_records!(stmt, [values]);
393                Ok(records)
394            })
395            .await
396            .inspect_err(|e| {
397                tracing::error!("Couldn't interact with the database connection: {e}");
398            })??;
399
400        tracing::debug!("Got {} records", records.len());
401
402        Ok(records)
403    }
404
405    /// Selects all rows in the `records` table, without any pagination.
406    ///
407    /// # Errors
408    ///
409    /// Returns [`DatabaseError`] if the operation fails.
410    #[instrument(skip(self))]
411    pub async fn get_all_records(&self, start_row: Option<i64>) -> Result<Vec<DbRecord>> {
412        tracing::debug!("Getting all records");
413
414        // pagination adapted from <https://gist.github.com/ssokolow/262503>
415        let query = r"SELECT * FROM records
416WHERE id NOT IN (SELECT id FROM records
417                 ORDER BY id ASC LIMIT ?1)
418ORDER BY id ASC LIMIT ?2;";
419
420        // start at 0 if we aren't given an explicit start row
421        let start_row = start_row.unwrap_or_default();
422        let conn = self.pool.get().await?;
423        let records = conn
424            .interact(move |conn| -> Result<Vec<DbRecord>> {
425                let mut stmt = conn.prepare(query)?;
426                let records = try_query_db_records!(stmt, [start_row, SELECT_ALL_LIMIT]);
427                Ok(records)
428            })
429            .await??;
430
431        Ok(records)
432    }
433
434    /// Selects all rows in the `records` table, without any pagination,
435    /// and returns their metadata (all fields except serialized data
436    /// of the page content).
437    ///
438    /// # Errors
439    ///
440    /// Returns [`DatabaseError`] if the operation fails.
441    #[instrument(skip(self))]
442    pub async fn get_all_meta(&self, start_row: Option<i64>) -> Result<Vec<DbRecordMeta>> {
443        tracing::debug!("Getting metadata for all records");
444
445        let query = r"
446SELECT id, url, client_datetime, timestamp_flora, title, language, updated_at
447FROM records
448WHERE id NOT IN (SELECT id FROM records
449                 ORDER BY id ASC LIMIT ?1)
450ORDER BY id ASC LIMIT ?2;";
451
452        let start_row = start_row.unwrap_or_default();
453        let conn = self.pool.get().await?;
454        let records = conn
455            .interact(move |conn| -> Result<Vec<DbRecordMeta>> {
456                let mut stmt = conn.prepare(query)?;
457                let records =
458                    try_query_db_records!(stmt, [start_row, SELECT_ALL_LIMIT], DbRecordMeta);
459                Ok(records)
460            })
461            .await??;
462
463        Ok(records)
464    }
465
466    /// Returns all record IDs in SQLite.
467    ///
468    /// # Errors
469    ///
470    /// Returns [`DatabaseError`] if the operation fails.
471    #[instrument(skip(self))]
472    pub async fn get_all_ids(&self) -> Result<Vec<i64>> {
473        tracing::debug!("Getting all record IDs");
474
475        let query = r"SELECT id FROM records;";
476
477        let conn = self.pool.get().await?;
478        let ids = conn
479            .interact(move |conn| -> Result<Vec<i64>> {
480                let mut stmt = conn.prepare(query)?;
481                let rows = stmt.query_map([], |row| row.get(0))?;
482                let mut ids = vec![];
483                for row in rows {
484                    ids.push(row?);
485                }
486
487                Ok(ids)
488            })
489            .await??;
490
491        Ok(ids)
492    }
493
494    /// Returns all record IDs in descending order.
495    #[instrument(skip(self))]
496    async fn get_ids_by_recent(&self, limit: i64) -> Result<Vec<i64>> {
497        // order by timestamp because in rare cases IDs might not
498        // be monotonically increasing
499        // (see <https://sqlite.org/autoinc.html>)
500        let query = r"SELECT id FROM records
501ORDER BY timestamp_flora DESC
502LIMIT ?1;";
503
504        let conn = self.pool.get().await?;
505        let ids = conn
506            .interact(move |conn| -> Result<Vec<i64>> {
507                let mut stmt = conn.prepare(query)?;
508                let rows = stmt.query_map([limit], |row| row.get(0))?;
509                let mut ids = vec![];
510                for row in rows {
511                    ids.push(row?);
512                }
513
514                Ok(ids)
515            })
516            .await??;
517
518        Ok(ids)
519    }
520
521    /// Similar to [`Database::get_all_meta`], except orders records by
522    /// descending order of `timestamp_flora`.
523    ///
524    /// Returns an empty list if the given index is out of bounds.
525    #[allow(clippy::missing_errors_doc)]
526    #[instrument(skip(self))]
527    pub async fn most_recent_meta(&self, start_idx: Option<usize>) -> Result<Vec<DbRecordMeta>> {
528        tracing::trace!("Selecting recent RecordMeta");
529
530        // idea: use row_number() window function to do this in SQL instead
531        let limit = start_idx.unwrap_or_default();
532        let prev_ids = self.get_ids_by_recent(limit.try_into()?).await?;
533        let prev_ids = prev_ids
534            .iter()
535            .copied()
536            .map(Value::from)
537            .collect::<Vec<Value>>();
538
539        let query = r"
540SELECT id, url, client_datetime, timestamp_flora, title, language, updated_at
541FROM records
542WHERE id NOT IN rarray(?1)
543ORDER BY timestamp_flora DESC LIMIT ?2;";
544
545        let conn = self.get_conn_array().await?;
546        let records = conn
547            .interact(move |conn| -> Result<Vec<DbRecordMeta>> {
548                let ids = Rc::new(prev_ids);
549                let mut stmt = conn.prepare(query)?;
550                let records = try_query_db_records!(
551                    stmt,
552                    rusqlite::params![ids, SELECT_ALL_LIMIT],
553                    DbRecordMeta
554                );
555                Ok(records)
556            })
557            .await??;
558
559        Ok(records)
560    }
561
562    /// Updates the record by ID.
563    ///
564    /// # Errors
565    /// Returns [`DatabaseError`] if the update operation fails.
566    #[instrument(skip_all, fields(record_id = record.id))]
567    pub async fn update_record_data(&self, record: DbRecord) -> Result<i64> {
568        let query = r"UPDATE records
569SET web_archive_serialized = :document, updated_at = :updated_at
570WHERE records.id = :id
571RETURNING id;";
572
573        let DbRecord {
574            id,
575            document,
576            updated_at,
577            ..
578        } = record;
579        let id = id.ok_or_else(|| DbErrorRepr::Update("missing id".into()))?;
580        let updated_at =
581            updated_at.ok_or_else(|| DbErrorRepr::Update("missing updated_at".into()))?;
582
583        let conn = self.pool.get().await?;
584        let row_id = conn
585            .interact(move |conn| {
586                let mut stmt = conn.prepare(query)?;
587                stmt.query_one(
588                    rusqlite::named_params! {
589                        ":id": id,
590                        ":document": document,
591                        ":updated_at": updated_at,
592                    },
593                    |row| row.get(0),
594                )
595            })
596            .await?
597            .inspect_err(|e| tracing::error!("Error on update: {e}"))?;
598        Ok(row_id)
599    }
600}
601
602#[cfg(test)]
603mod tests {
604    use super::*;
605    use crate::model::tests::test_record_params;
606    use anyhow::{Context, Result};
607    use chrono::{DateTime, FixedOffset, Utc};
608    use omry_archiving::record::Record;
609    use pretty_assertions::assert_eq;
610    use serde::Serialize;
611    use std::fs::File;
612    use std::ops::Deref;
613    use tempfile::NamedTempFile;
614
615    /// Data structure to keep the temporary file holding the database alive,
616    /// because it will get unlinked on drop.
617    struct TempDb {
618        #[allow(dead_code)]
619        temp_db_file: NamedTempFile<File>,
620        db: Database,
621    }
622
623    impl TempDb {
624        async fn new() -> Result<Self> {
625            let temp_db_file = NamedTempFile::with_suffix(".db")?;
626            let db = Database::new_with_records_table(temp_db_file.path()).await?;
627            Ok(Self { temp_db_file, db })
628        }
629    }
630
631    impl Deref for TempDb {
632        type Target = Database;
633
634        fn deref(&self) -> &Self::Target {
635            &self.db
636        }
637    }
638
639    /// Wrapper for [`DbRecord`] that supports [`Serialize`] for redacting snapshots.
640    #[derive(Debug, Serialize)]
641    struct DbRecordDef {
642        id: Option<i64>,
643        url: String,
644        client_datetime: DateTime<FixedOffset>,
645        client_datetime_raw: Option<String>,
646        timestamp_flora: i64,
647        title: String,
648        language: Option<String>,
649        document: Vec<u8>,
650        updated_at: Option<i64>,
651    }
652
653    impl From<DbRecord> for DbRecordDef {
654        fn from(value: DbRecord) -> DbRecordDef {
655            let DbRecord {
656                id,
657                url,
658                client_datetime,
659                client_datetime_raw,
660                timestamp_flora,
661                title,
662                language,
663                document,
664                updated_at,
665            } = value;
666
667            Self {
668                id,
669                url,
670                client_datetime,
671                client_datetime_raw,
672                timestamp_flora,
673                title,
674                language,
675                document,
676                updated_at,
677            }
678        }
679    }
680
681    #[derive(Debug, Serialize)]
682    struct DbRecordMetaDef {
683        id: i64,
684        url: String,
685        client_datetime: DateTime<FixedOffset>,
686        timestamp_flora: i64,
687        title: String,
688        language: Option<String>,
689        updated_at: Option<i64>,
690    }
691
692    impl From<DbRecordMeta> for DbRecordMetaDef {
693        fn from(value: DbRecordMeta) -> Self {
694            let DbRecordMeta {
695                id,
696                url,
697                client_datetime,
698                timestamp_flora,
699                title,
700                language,
701                updated_at,
702            } = value;
703
704            Self {
705                id,
706                url,
707                client_datetime,
708                timestamp_flora,
709                title,
710                language,
711                updated_at,
712            }
713        }
714    }
715
716    #[allow(dead_code, reason = "Used in snapshot.")]
717    #[derive(Debug)]
718    struct RowInfo {
719        name: String,
720        type_: String,
721        notnull: bool,
722    }
723
724    impl TryFrom<&Row<'_>> for RowInfo {
725        type Error = rusqlite::Error;
726
727        fn try_from(row: &Row) -> std::result::Result<Self, Self::Error> {
728            let (name, type_, notnull) = (row.get(0)?, row.get(1)?, row.get(2)?);
729            Ok(Self {
730                name,
731                type_,
732                notnull,
733            })
734        }
735    }
736
737    fn test_db_record() -> Result<DbRecord> {
738        let record = Record::new(test_record_params()?);
739        Ok(record.try_into()?)
740    }
741
742    #[tokio::test]
743    async fn can_create_table() -> Result<()> {
744        let db = TempDb::new().await?;
745        let query = r#"SELECT name, type, "notnull" FROM pragma_table_info('records');"#;
746        let conn = db.pool.get().await?;
747        let row_infos = conn
748            .interact(move |conn| -> Result<Vec<RowInfo>> {
749                let mut stmt = conn.prepare(query)?;
750                let rows = stmt.query_map([], |row| RowInfo::try_from(row))?;
751                let mut values = vec![];
752                for row in rows {
753                    values.push(row?);
754                }
755                Ok(values)
756            })
757            .await
758            .map_err(|e| anyhow::anyhow!("{e}"))??;
759        insta::assert_debug_snapshot!(row_infos);
760        Ok(())
761    }
762
763    #[tokio::test]
764    async fn can_add_record() -> Result<()> {
765        let db = TempDb::new().await?;
766        let record = test_db_record()?;
767        let result = db.insert(record).await?;
768        assert_eq!(result, 1);
769        Ok(())
770    }
771
772    #[tokio::test]
773    async fn can_get_record() -> Result<()> {
774        let db = TempDb::new().await?;
775        let record = test_db_record()?;
776        assert!(record.id.is_none());
777        db.insert(record).await?;
778
779        let expected_id = 1i64;
780        let retrieved_record = db.get_record(expected_id).await?;
781
782        let retrieved_record: DbRecordDef = retrieved_record.into();
783        insta::assert_ron_snapshot!(retrieved_record, {
784            ".client_datetime" => "[client_datetime]",
785            ".timestamp_flora" => "[timestamp]"
786        });
787
788        Ok(())
789    }
790
791    #[tokio::test]
792    async fn can_get_many_records() -> Result<()> {
793        let db = TempDb::new().await?;
794
795        let num_records_to_insert = 7;
796        let mut record_ids = Vec::with_capacity(num_records_to_insert);
797
798        for _ in 0..num_records_to_insert {
799            let record = test_db_record()?;
800            let record_id = db.insert(record).await?;
801            record_ids.push(record_id);
802        }
803
804        let selected_records = db
805            .get_many_records(&record_ids)
806            .await?
807            .into_iter()
808            .map(DbRecordDef::from)
809            .collect::<Vec<_>>();
810
811        assert_eq!(selected_records.len(), num_records_to_insert);
812
813        insta::assert_ron_snapshot!(selected_records, {
814            "[].client_datetime" => "[client_datetime]",
815            "[].timestamp_flora" => "[timestamp]"
816        });
817
818        Ok(())
819    }
820
821    #[tokio::test]
822    async fn passing_too_many_bind_parameters_returns_error() -> Result<()> {
823        let db = TempDb::new().await?;
824
825        let num_records_to_insert = MAX_BIND_PARAMS + 1;
826        db.insert(test_db_record()?).await?;
827
828        let mut ids = Vec::with_capacity(num_records_to_insert);
829        for index in 0..num_records_to_insert {
830            ids.push(i64::try_from(index + 1)?);
831        }
832
833        let result = db.get_many_records(&ids).await;
834        assert!(result.is_err());
835
836        // At this point, we know this is definitely an error.
837        #[allow(clippy::unwrap_used)]
838        let error = result.unwrap_err();
839        insta::assert_snapshot!(error);
840
841        Ok(())
842    }
843
844    #[tokio::test]
845    async fn can_get_count() -> Result<()> {
846        let db = TempDb::new().await?;
847
848        let num_records_to_insert = 7;
849        let mut record_ids = Vec::with_capacity(num_records_to_insert);
850
851        for _ in 0..num_records_to_insert {
852            let record = test_db_record()?;
853            let record_id = db.insert(record).await?;
854            record_ids.push(record_id);
855        }
856
857        let count: usize = db.count().await?.try_into()?;
858
859        assert_eq!(count, num_records_to_insert);
860        Ok(())
861    }
862
863    #[tokio::test]
864    async fn can_delete_record() -> Result<()> {
865        let db = TempDb::new().await?;
866        let record = test_db_record()?;
867        let id = db.insert(record).await?;
868        let rows_changed = db.delete(id).await?;
869        assert_eq!(rows_changed, 1);
870        Ok(())
871    }
872
873    #[tokio::test]
874    async fn can_get_all_records() -> Result<()> {
875        let db = TempDb::new().await?;
876
877        // Insert more than the limit to test pagination.
878        let num_records_to_insert = (SELECT_ALL_LIMIT + 5).try_into()?;
879        for _ in 0..num_records_to_insert {
880            let record = test_db_record()?;
881            let _record_id = db.insert(record).await?;
882        }
883
884        let mut db_recs = Vec::with_capacity(num_records_to_insert);
885        let mut start_row = None;
886        loop {
887            let mut selected_records = db
888                .get_all_records(start_row)
889                .await?
890                .into_iter()
891                .map(DbRecordDef::from)
892                .collect::<Vec<_>>();
893
894            if selected_records.is_empty() {
895                break;
896            }
897
898            db_recs.append(&mut selected_records);
899            start_row = Some(db_recs.len().try_into()?);
900        }
901
902        assert_eq!(db_recs.len(), num_records_to_insert);
903
904        insta::assert_ron_snapshot!(db_recs, {
905            "[].client_datetime" => "[client_datetime]",
906            "[].timestamp_flora" => "[timestamp]",
907            "[].document" => "[data]"
908        });
909
910        Ok(())
911    }
912
913    /// This tests that we can still get all records
914    /// after deleting some from the middle (causing IDs to become incontiguous).
915    #[tokio::test]
916    async fn can_get_all_records_after_deleting_some() -> Result<()> {
917        let db = TempDb::new().await?;
918
919        // Insert more than the limit to test pagination.
920        let num_records_to_insert = (SELECT_ALL_LIMIT + 5).try_into()?;
921        let mut record_ids = Vec::with_capacity(num_records_to_insert);
922
923        for _ in 0..num_records_to_insert {
924            let record = test_db_record()?;
925            let record_id = db.insert(record).await?;
926            record_ids.push(record_id);
927        }
928
929        // remove IDs from `record_ids` to assert that
930        // `record_ids` length matches the number of records
931        // we get back in the end
932        let idx_to_delete = [3, 7, 11];
933        for idx in idx_to_delete {
934            let id = record_ids.remove(idx);
935            db.delete(id).await?;
936        }
937
938        let mut db_recs = Vec::with_capacity(num_records_to_insert);
939        let mut start_row = None;
940        loop {
941            let mut selected_records = db
942                .get_all_records(start_row)
943                .await?
944                .into_iter()
945                .map(DbRecordDef::from)
946                .collect::<Vec<_>>();
947
948            if selected_records.is_empty() {
949                break;
950            }
951
952            db_recs.append(&mut selected_records);
953            start_row = Some(db_recs.len().try_into()?);
954        }
955
956        assert_eq!(db_recs.len(), record_ids.len());
957        insta::assert_ron_snapshot!(db_recs, {
958            "[].client_datetime" => "[client_datetime]",
959            "[].timestamp_flora" => "[timestamp]",
960            "[].document" => "[data]"
961        });
962
963        Ok(())
964    }
965
966    #[tokio::test]
967    async fn can_get_all_meta() -> Result<()> {
968        let db = TempDb::new().await?;
969
970        let num_records_to_insert = (SELECT_ALL_LIMIT + 1).try_into()?;
971        let mut record_ids = Vec::with_capacity(num_records_to_insert);
972
973        for _ in 0..num_records_to_insert {
974            let record = test_db_record()?;
975            let record_id = db.insert(record).await?;
976            record_ids.push(record_id);
977        }
978
979        let mut start_row = None;
980        let mut db_recs = Vec::with_capacity(num_records_to_insert);
981        loop {
982            let mut selected_recs = db
983                .get_all_meta(start_row)
984                .await?
985                .into_iter()
986                .map(DbRecordMetaDef::from)
987                .collect::<Vec<_>>();
988
989            if selected_recs.is_empty() {
990                break;
991            }
992
993            db_recs.append(&mut selected_recs);
994            start_row = Some(db_recs.len().try_into()?);
995        }
996
997        assert_eq!(db_recs.len(), num_records_to_insert);
998
999        insta::assert_ron_snapshot!(db_recs, {
1000            "[].client_datetime" => "[client_datetime]",
1001            "[].timestamp_flora" => "[timestamp]",
1002        });
1003
1004        Ok(())
1005    }
1006
1007    #[tokio::test]
1008    async fn can_get_all_ids() -> Result<()> {
1009        let db = TempDb::new().await?;
1010
1011        let num_records_to_insert = 7;
1012        let mut record_ids = Vec::with_capacity(num_records_to_insert);
1013
1014        for _ in 0..num_records_to_insert {
1015            let record = test_db_record()?;
1016            let record_id = db.insert(record).await?;
1017            record_ids.push(record_id);
1018        }
1019
1020        let db_ids = db.get_all_ids().await?;
1021        assert_eq!(db_ids, record_ids);
1022        Ok(())
1023    }
1024
1025    #[tokio::test]
1026    async fn can_get_most_recent_meta() -> Result<()> {
1027        let db = TempDb::new().await?;
1028
1029        let num_recs_to_insert = (SELECT_ALL_LIMIT + 5).try_into()?;
1030        let ts_incr = 5;
1031        let mut record_ids = Vec::with_capacity(num_recs_to_insert);
1032
1033        let mut cur_timestamp = Utc::now().timestamp();
1034        for _ in 0..num_recs_to_insert {
1035            let mut rec = test_db_record()?;
1036            rec.timestamp_flora = cur_timestamp;
1037            let record_id = db.insert(rec).await?;
1038            record_ids.push(record_id);
1039            cur_timestamp += ts_incr;
1040        }
1041
1042        let mut recs = Vec::with_capacity(num_recs_to_insert);
1043        let mut start_idx = None;
1044        loop {
1045            let mut selected_recs = db
1046                .most_recent_meta(start_idx)
1047                .await?
1048                .into_iter()
1049                .map(DbRecordMetaDef::from)
1050                .collect::<Vec<_>>();
1051
1052            if selected_recs.is_empty() {
1053                break;
1054            }
1055
1056            recs.append(&mut selected_recs);
1057            start_idx = Some(recs.len());
1058        }
1059
1060        assert_eq!(recs.len(), num_recs_to_insert);
1061        let timestamps = recs.iter().map(|r| r.timestamp_flora).collect::<Vec<_>>();
1062        assert!(
1063            recs.is_sorted_by(|a, b| a.timestamp_flora > b.timestamp_flora),
1064            "{timestamps:#?}"
1065        );
1066        Ok(())
1067    }
1068
1069    #[tokio::test]
1070    async fn can_update_record() -> Result<()> {
1071        let db = TempDb::new().await?;
1072        let record = test_db_record()?;
1073        let id = db.insert(record).await.context("error inserting record")?;
1074
1075        // workaround for DbRecord not being `Clone`
1076        let mut record = db.get_record(id).await.context("error getting record")?;
1077        record.updated_at = Some(Utc::now().timestamp());
1078        let updated_id = db
1079            .update_record_data(record)
1080            .await
1081            .context("error updating record data")?;
1082        assert_eq!(id, updated_id);
1083        Ok(())
1084    }
1085}