crates/db_schema/src/impls/federation_allowlist.rs

   1 use crate::{
   2   schema::federation_allowlist,
   3   source::{
   4     federation_allowlist::{FederationAllowList, FederationAllowListForm},
   5     instance::Instance,
   6   },
   7   utils::{get_conn, DbPool},
   8 };
   9 use diesel::{dsl::insert_into, result::Error};
  10 use diesel_async::{AsyncPgConnection, RunQueryDsl};
  11
  12 impl FederationAllowList {
  13   pub async fn replace(pool: &mut DbPool<'_>, list_opt: Option<Vec<String>>) -> Result<(), Error> {
  14     let conn = &mut get_conn(pool).await?;
  15     conn
  16       .build_transaction()
  17       .run(|conn| {
  18         Box::pin(async move {
  19           if let Some(list) = list_opt {
  20             Self::clear(conn).await?;
  21
  22             for domain in list {
  23               // Upsert all of these as instances
  24               let instance = Instance::read_or_create(&mut conn.into(), domain).await?;
  25
  26               let form = FederationAllowListForm {
  27                 instance_id: instance.id,
  28                 updated: None,
  29               };
  30               insert_into(federation_allowlist::table)
  31                 .values(form)
  32                 .get_result::<Self>(conn)
  33                 .await?;
  34             }
  35             Ok(())
  36           } else {
  37             Ok(())
  38           }
  39         }) as _
  40       })
  41       .await
  42   }
  43
  44   async fn clear(conn: &mut AsyncPgConnection) -> Result<usize, Error> {
  45     diesel::delete(federation_allowlist::table)
  46       .execute(conn)
  47       .await
  48   }
  49 }
  50 #[cfg(test)]
  51 mod tests {
  52   #![allow(clippy::unwrap_used)]
  53   #![allow(clippy::indexing_slicing)]
  54
  55   use crate::{
  56     source::{federation_allowlist::FederationAllowList, instance::Instance},
  57     utils::build_db_pool_for_tests,
  58   };
  59   use serial_test::serial;
  60
  61   #[tokio::test]
  62   #[serial]
  63   async fn test_allowlist_insert_and_clear() {
  64     let pool = &build_db_pool_for_tests().await;
  65     let pool = &mut pool.into();
  66     let domains = vec![
  67       "tld1.xyz".to_string(),
  68       "tld2.xyz".to_string(),
  69       "tld3.xyz".to_string(),
  70     ];
  71
  72     let allowed = Some(domains.clone());
  73
  74     FederationAllowList::replace(pool, allowed).await.unwrap();
  75
  76     let allows = Instance::allowlist(pool).await.unwrap();
  77     let allows_domains = allows
  78       .iter()
  79       .map(|i| i.domain.clone())
  80       .collect::<Vec<String>>();
  81
  82     assert_eq!(3, allows.len());
  83     assert_eq!(domains, allows_domains);
  84
  85     // Now test clearing them via Some(empty vec)
  86     let clear_allows = Some(Vec::new());
  87
  88     FederationAllowList::replace(pool, clear_allows)
  89       .await
  90       .unwrap();
  91     let allows = Instance::allowlist(pool).await.unwrap();
  92
  93     assert_eq!(0, allows.len());
  94
  95     Instance::delete_all(pool).await.unwrap();
  96   }
  97 }