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 };
   8 use diesel::{dsl::*, result::Error, *};
   9
  10 impl FederationAllowList {
  11   pub fn replace(conn: &mut PgConnection, list_opt: Option<Vec<String>>) -> Result<(), Error> {
  12     conn.build_transaction().read_write().run(|conn| {
  13       if let Some(list) = list_opt {
  14         Self::clear(conn)?;
  15
  16         for domain in list {
  17           // Upsert all of these as instances
  18           let instance = Instance::create(conn, &domain)?;
  19
  20           let form = FederationAllowListForm {
  21             instance_id: instance.id,
  22             updated: None,
  23           };
  24           insert_into(federation_allowlist::table)
  25             .values(form)
  26             .get_result::<Self>(conn)?;
  27         }
  28         Ok(())
  29       } else {
  30         Ok(())
  31       }
  32     })
  33   }
  34
  35   pub fn clear(conn: &mut PgConnection) -> Result<usize, Error> {
  36     diesel::delete(federation_allowlist::table).execute(conn)
  37   }
  38 }
  39 #[cfg(test)]
  40 mod tests {
  41   use crate::{
  42     source::{federation_allowlist::FederationAllowList, instance::Instance},
  43     utils::establish_unpooled_connection,
  44   };
  45   use serial_test::serial;
  46
  47   #[test]
  48   #[serial]
  49   fn test_allowlist_insert_and_clear() {
  50     let conn = &mut establish_unpooled_connection();
  51     let allowed = Some(vec![
  52       "tld1.xyz".to_string(),
  53       "tld2.xyz".to_string(),
  54       "tld3.xyz".to_string(),
  55     ]);
  56
  57     FederationAllowList::replace(conn, allowed).unwrap();
  58
  59     let allows = Instance::allowlist(conn).unwrap();
  60
  61     assert_eq!(3, allows.len());
  62     assert_eq!(
  63       vec![
  64         "tld1.xyz".to_string(),
  65         "tld2.xyz".to_string(),
  66         "tld3.xyz".to_string()
  67       ],
  68       allows
  69     );
  70
  71     // Now test clearing them via Some(empty vec)
  72     let clear_allows = Some(Vec::new());
  73
  74     FederationAllowList::replace(conn, clear_allows).unwrap();
  75     let allows = Instance::allowlist(conn).unwrap();
  76
  77     assert_eq!(0, allows.len());
  78
  79     Instance::delete_all(conn).unwrap();
  80   }
  81 }