1#![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
15const MAX_BIND_PARAMS: usize = 32766;
18
19macro_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#[derive(Debug, Error)]
43#[error(transparent)]
44pub struct DatabaseError(DbErrorRepr);
45
46impl<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#[derive(Error, Debug)]
60enum DbErrorRepr {
61 #[error("error creating the DB pool: {0}")]
63 CreatePool(#[from] deadpool_sqlite::CreatePoolError),
64
65 #[error("error interacting with the DB pool: {0}")]
69 PoolInteract(String),
70
71 #[error("error getting connection from the DB pool: {0}")]
73 Pool(#[from] deadpool_sqlite::PoolError),
74
75 #[error("error in rusqlite: {0}")]
77 Rusqlite(#[from] rusqlite::Error),
78
79 #[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("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
97pub 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#[derive(Debug, Clone)]
133pub struct Database {
134 pool: Pool,
136}
137
138impl Database {
139 #[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 #[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 #[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 #[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 #[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 #[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 #[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 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 #[instrument(skip(self))]
315 pub async fn delete(&self, id: i64) -> Result<usize> {
316 const QUERY: &str = r"DELETE FROM records WHERE id = ?1;";
317
318 let conn = self.pool.get().await?;
319 let rows_changed = conn
320 .interact(move |conn| {
321 let mut stmt = conn.prepare(QUERY)?;
322 stmt.execute([id])
323 })
324 .await??;
325
326 Ok(rows_changed)
327 }
328
329 pub async fn get_record(&self, id: i64) -> Result<DbRecord> {
334 let query = r"SELECT id, url, client_datetime, client_datetime_raw, timestamp_flora, title,
335 language, web_archive_serialized
336 FROM records WHERE id = ?1;";
337
338 let conn = self.pool.get().await?;
339 let result = conn
340 .interact(move |conn| {
341 let mut stmt = conn.prepare(query)?;
342 stmt.query_row([id], |row| {
343 let record = row.try_into()?;
344 Ok(record)
345 })
346 })
347 .await??;
348
349 Ok(result)
350 }
351
352 #[instrument]
357 pub async fn get_many_records(&self, ids: &[i64]) -> Result<Vec<DbRecord>> {
358 if ids.len() > MAX_BIND_PARAMS {
359 return Err(DbErrorRepr::TooManyBindParameters(ids.len()).into());
360 }
361
362 tracing::debug!("Getting records.");
363
364 let query = r"SELECT id, url, client_datetime, client_datetime_raw, timestamp_flora, title,
365 language, web_archive_serialized
366 FROM records
367 WHERE id IN rarray(?1);";
368
369 let conn = self
370 .get_conn_array()
371 .await
372 .inspect_err(|e| tracing::error!("Couldn't get a connection from the pool: {e}"))?;
373
374 let values = ids.iter().copied().map(Value::from).collect::<Vec<Value>>();
375
376 let records = conn
377 .interact(|conn| -> Result<Vec<DbRecord>> {
378 let values = Rc::new(values);
379
380 let mut stmt = conn
381 .prepare(query)
382 .inspect_err(|e| tracing::error!("Couldn't prepare the query: {e}"))?;
383
384 let records = try_query_db_records!(stmt, [values]);
385 Ok(records)
386 })
387 .await
388 .inspect_err(|e| {
389 tracing::error!("Couldn't interact with the database connection: {e}");
390 })??;
391
392 tracing::debug!("Got {} records", records.len());
393
394 Ok(records)
395 }
396
397 #[instrument(skip(self))]
403 pub async fn get_all_records(&self) -> Result<Vec<DbRecord>> {
404 tracing::debug!("Getting all records");
405
406 let query = r"SELECT * FROM records";
407 let conn = self.pool.get().await?;
408 let records = conn
409 .interact(move |conn| -> Result<Vec<DbRecord>> {
410 let mut stmt = conn.prepare(query)?;
411 let records = try_query_db_records!(stmt, []);
412 Ok(records)
413 })
414 .await??;
415
416 Ok(records)
417 }
418
419 #[instrument(skip(self))]
427 pub async fn get_all_meta(&self) -> Result<Vec<DbRecordMeta>> {
428 tracing::debug!("Getting metadata for all records");
429
430 let query = r"SELECT id, url, client_datetime, timestamp_flora, title,
431 language FROM records;";
432
433 let conn = self.pool.get().await?;
434 let records = conn
435 .interact(move |conn| -> Result<Vec<DbRecordMeta>> {
436 let mut stmt = conn.prepare(query)?;
437 let records = try_query_db_records!(stmt, [], DbRecordMeta);
438 Ok(records)
439 })
440 .await??;
441
442 Ok(records)
443 }
444
445 #[instrument(skip(self))]
451 pub async fn get_all_ids(&self) -> Result<Vec<i64>> {
452 tracing::debug!("Getting all record IDs");
453
454 let query = r"SELECT id FROM records;";
455
456 let conn = self.pool.get().await?;
457 let ids = conn
458 .interact(move |conn| -> Result<Vec<i64>> {
459 let mut stmt = conn.prepare(query)?;
460 let rows = stmt.query_map([], |row| row.get(0))?;
461
462 let (lower_bound, upper_bound) = rows.size_hint();
463 let n_records = upper_bound.unwrap_or(lower_bound);
464 let mut ids = Vec::with_capacity(n_records);
465
466 for row in rows {
467 ids.push(row?);
468 }
469
470 Ok(ids)
471 })
472 .await??;
473
474 Ok(ids)
475 }
476
477 #[allow(clippy::missing_errors_doc)]
480 #[instrument(skip(self))]
481 pub async fn most_recent_meta(&self) -> Result<Vec<DbRecordMeta>> {
482 tracing::trace!("Selecting recent RecordMeta");
483
484 let query = r"
485 SELECT id, url, client_datetime, timestamp_flora, title, language
486 FROM records
487 ORDER BY timestamp_flora DESC;";
488
489 let conn = self.pool.get().await?;
490 let records = conn
491 .interact(move |conn| -> Result<Vec<DbRecordMeta>> {
492 let mut stmt = conn.prepare(query)?;
493 let records = try_query_db_records!(stmt, [], DbRecordMeta);
494 Ok(records)
495 })
496 .await??;
497
498 Ok(records)
499 }
500}
501
502#[cfg(test)]
503mod tests {
504 use super::*;
505 use crate::model::tests::test_record_params;
506 use anyhow::Result;
507 use chrono::{DateTime, FixedOffset, Utc};
508 use omry_archiving::record::Record;
509 use pretty_assertions::assert_eq;
510 use serde::Serialize;
511 use std::fs::File;
512 use std::ops::Deref;
513 use tempfile::NamedTempFile;
514
515 struct TempDb {
518 #[allow(dead_code)]
519 temp_db_file: NamedTempFile<File>,
520 db: Database,
521 }
522
523 impl TempDb {
524 async fn new() -> Result<Self> {
525 let temp_db_file = NamedTempFile::with_suffix(".db")?;
526 let db = Database::new_with_records_table(temp_db_file.path()).await?;
527 Ok(Self { temp_db_file, db })
528 }
529 }
530
531 impl Deref for TempDb {
532 type Target = Database;
533
534 fn deref(&self) -> &Self::Target {
535 &self.db
536 }
537 }
538
539 #[derive(Debug, Serialize)]
541 struct DbRecordDef {
542 id: Option<i64>,
543 url: String,
544 client_datetime: DateTime<FixedOffset>,
545 client_datetime_raw: Option<String>,
546 timestamp_flora: i64,
547 title: String,
548 language: Option<String>,
549 web_archive_serialized: Vec<u8>,
550 }
551
552 impl From<DbRecord> for DbRecordDef {
553 fn from(value: DbRecord) -> DbRecordDef {
554 let DbRecord {
555 id,
556 url,
557 client_datetime,
558 client_datetime_raw,
559 timestamp_flora,
560 title,
561 language,
562 web_archive_serialized,
563 } = value;
564
565 Self {
566 id,
567 url,
568 client_datetime,
569 client_datetime_raw,
570 timestamp_flora,
571 title,
572 language,
573 web_archive_serialized,
574 }
575 }
576 }
577
578 #[derive(Debug, Serialize)]
579 struct DbRecordMetaDef {
580 id: i64,
581 url: String,
582 client_datetime: DateTime<FixedOffset>,
583 timestamp_flora: i64,
584 title: String,
585 language: Option<String>,
586 }
587
588 impl From<DbRecordMeta> for DbRecordMetaDef {
589 fn from(value: DbRecordMeta) -> Self {
590 let DbRecordMeta {
591 id,
592 url,
593 client_datetime,
594 timestamp_flora,
595 title,
596 language,
597 } = value;
598
599 Self {
600 id,
601 url,
602 client_datetime,
603 timestamp_flora,
604 title,
605 language,
606 }
607 }
608 }
609
610 fn test_db_record() -> Result<DbRecord> {
611 let record = Record::new(test_record_params()?);
612 Ok(record.try_into()?)
613 }
614
615 #[tokio::test]
616 async fn can_create_table() -> Result<()> {
617 let db = TempDb::new().await?;
618 assert!(db.create_table(&DbRecord::schema()).await.is_ok());
619 Ok(())
621 }
622
623 #[tokio::test]
624 async fn can_add_record() -> Result<()> {
625 let db = TempDb::new().await?;
626 let record = test_db_record()?;
627 let result = db.insert(record).await?;
628 assert_eq!(result, 1);
629 Ok(())
630 }
631
632 #[tokio::test]
633 async fn can_get_record() -> Result<()> {
634 let db = TempDb::new().await?;
635 let record = test_db_record()?;
636 assert!(record.id.is_none());
637 db.insert(record).await?;
638
639 let expected_id = 1i64;
640 let retrieved_record = db.get_record(expected_id).await?;
641
642 let retrieved_record: DbRecordDef = retrieved_record.into();
643 insta::assert_ron_snapshot!(retrieved_record, {
644 ".client_datetime" => "[client_datetime]",
645 ".timestamp_flora" => "[timestamp]"
646 });
647
648 Ok(())
649 }
650
651 #[tokio::test]
652 async fn can_get_many_records() -> Result<()> {
653 let db = TempDb::new().await?;
654
655 let num_records_to_insert = 7;
656 let mut record_ids = Vec::with_capacity(num_records_to_insert);
657
658 for _ in 0..num_records_to_insert {
659 let record = test_db_record()?;
660 let record_id = db.insert(record).await?;
661 record_ids.push(record_id);
662 }
663
664 let selected_records = db
665 .get_many_records(&record_ids)
666 .await?
667 .into_iter()
668 .map(DbRecordDef::from)
669 .collect::<Vec<_>>();
670
671 assert_eq!(selected_records.len(), num_records_to_insert);
672
673 insta::assert_ron_snapshot!(selected_records, {
674 "[].client_datetime" => "[client_datetime]",
675 "[].timestamp_flora" => "[timestamp]"
676 });
677
678 Ok(())
679 }
680
681 #[tokio::test]
682 async fn passing_too_many_bind_parameters_returns_error() -> Result<()> {
683 let db = TempDb::new().await?;
684
685 let num_records_to_insert = MAX_BIND_PARAMS + 1;
686 db.insert(test_db_record()?).await?;
687
688 let mut ids = Vec::with_capacity(num_records_to_insert);
689 for index in 0..num_records_to_insert {
690 ids.push(i64::try_from(index + 1)?);
691 }
692
693 let result = db.get_many_records(&ids).await;
694 assert!(result.is_err());
695
696 #[allow(clippy::unwrap_used)]
698 let error = result.unwrap_err();
699 insta::assert_snapshot!(error);
700
701 Ok(())
702 }
703
704 #[tokio::test]
705 async fn can_get_count() -> Result<()> {
706 let db = TempDb::new().await?;
707
708 let num_records_to_insert = 7;
709 let mut record_ids = Vec::with_capacity(num_records_to_insert);
710
711 for _ in 0..num_records_to_insert {
712 let record = test_db_record()?;
713 let record_id = db.insert(record).await?;
714 record_ids.push(record_id);
715 }
716
717 let count: usize = db.count().await?.try_into()?;
718
719 assert_eq!(count, num_records_to_insert);
720 Ok(())
721 }
722
723 #[tokio::test]
724 async fn can_delete_record() -> Result<()> {
725 let db = TempDb::new().await?;
726 let record = test_db_record()?;
727 let id = db.insert(record).await?;
728 let rows_changed = db.delete(id).await?;
729 assert_eq!(rows_changed, 1);
730 Ok(())
731 }
732
733 #[tokio::test]
734 async fn can_get_all_records() -> Result<()> {
735 let db = TempDb::new().await?;
736
737 let num_records_to_insert = 7;
738 let mut record_ids = Vec::with_capacity(num_records_to_insert);
739
740 for _ in 0..num_records_to_insert {
741 let record = test_db_record()?;
742 let record_id = db.insert(record).await?;
743 record_ids.push(record_id);
744 }
745
746 let selected_records = db
747 .get_all_records()
748 .await?
749 .into_iter()
750 .map(DbRecordDef::from)
751 .collect::<Vec<_>>();
752
753 assert_eq!(selected_records.len(), num_records_to_insert);
754
755 insta::assert_ron_snapshot!(selected_records, {
756 "[].client_datetime" => "[client_datetime]",
757 "[].timestamp_flora" => "[timestamp]",
758 "[].web_archive_serialized" => "[data]"
759 });
760
761 Ok(())
762 }
763
764 #[tokio::test]
765 async fn can_get_all_meta() -> Result<()> {
766 let db = TempDb::new().await?;
767
768 let num_records_to_insert = 7;
769 let mut record_ids = Vec::with_capacity(num_records_to_insert);
770
771 for _ in 0..num_records_to_insert {
772 let record = test_db_record()?;
773 let record_id = db.insert(record).await?;
774 record_ids.push(record_id);
775 }
776
777 let selected_records = db
778 .get_all_meta()
779 .await?
780 .into_iter()
781 .map(DbRecordMetaDef::from)
782 .collect::<Vec<_>>();
783
784 assert_eq!(selected_records.len(), num_records_to_insert);
785
786 insta::assert_ron_snapshot!(selected_records, {
787 "[].client_datetime" => "[client_datetime]",
788 "[].timestamp_flora" => "[timestamp]",
789 });
790
791 Ok(())
792 }
793
794 #[tokio::test]
795 async fn can_get_all_ids() -> Result<()> {
796 let db = TempDb::new().await?;
797
798 let num_records_to_insert = 7;
799 let mut record_ids = Vec::with_capacity(num_records_to_insert);
800
801 for _ in 0..num_records_to_insert {
802 let record = test_db_record()?;
803 let record_id = db.insert(record).await?;
804 record_ids.push(record_id);
805 }
806
807 let db_ids = db.get_all_ids().await?;
808 assert_eq!(db_ids, record_ids);
809 Ok(())
810 }
811
812 #[tokio::test]
813 async fn can_get_most_recent_meta() -> Result<()> {
814 let db = TempDb::new().await?;
815
816 let num_records_to_insert = 5;
817 let ts_incr = 5;
818 let mut record_ids = Vec::with_capacity(num_records_to_insert);
819
820 let mut cur_timestamp = Utc::now().timestamp();
821 for _ in 0..num_records_to_insert {
822 let mut rec = test_db_record()?;
823 rec.timestamp_flora = cur_timestamp;
824 let record_id = db.insert(rec).await?;
825 record_ids.push(record_id);
826 cur_timestamp += ts_incr;
827 }
828
829 let recs = db
830 .most_recent_meta()
831 .await?
832 .into_iter()
833 .map(DbRecordMetaDef::from)
834 .collect::<Vec<_>>();
835
836 assert_eq!(recs.len(), num_records_to_insert);
837 let timestamps = recs.iter().map(|r| r.timestamp_flora).collect::<Vec<_>>();
838 assert!(
839 recs.is_sorted_by(|a, b| a.timestamp_flora > b.timestamp_flora),
840 "{timestamps:#?}"
841 );
842 Ok(())
843 }
844}