crates/db_schema/src/utils.rs

   1 use crate::newtypes::DbUrl;
   2 use activitypub_federation::{core::object_id::ObjectId, traits::ApubObject};
   3 use chrono::NaiveDateTime;
   4 use diesel::{
   5   backend::Backend,
   6   deserialize::FromSql,
   7   serialize::{Output, ToSql},
   8   sql_types::Text,
   9   Connection,
  10   PgConnection,
  11 };
  12 use lemmy_utils::error::LemmyError;
  13 use once_cell::sync::Lazy;
  14 use regex::Regex;
  15 use std::{env, env::VarError, io::Write};
  16 use url::Url;
  17
  18 pub type DbPool = diesel::r2d2::Pool<diesel::r2d2::ConnectionManager<diesel::PgConnection>>;
  19
  20 pub fn get_database_url_from_env() -> Result<String, VarError> {
  21   env::var("LEMMY_DATABASE_URL")
  22 }
  23
  24 pub fn fuzzy_search(q: &str) -> String {
  25   let replaced = q.replace('%', "\\%").replace('_', "\\_").replace(' ', "%");
  26   format!("%{}%", replaced)
  27 }
  28
  29 pub fn limit_and_offset(page: Option<i64>, limit: Option<i64>) -> (i64, i64) {
  30   let page = page.unwrap_or(1);
  31   let limit = limit.unwrap_or(10);
  32   let offset = limit * (page - 1);
  33   (limit, offset)
  34 }
  35
  36 pub fn is_email_regex(test: &str) -> bool {
  37   EMAIL_REGEX.is_match(test)
  38 }
  39
  40 pub fn diesel_option_overwrite(opt: &Option<String>) -> Option<Option<String>> {
  41   match opt {
  42     // An empty string is an erase
  43     Some(unwrapped) => {
  44       if !unwrapped.eq("") {
  45         Some(Some(unwrapped.to_owned()))
  46       } else {
  47         Some(None)
  48       }
  49     }
  50     None => None,
  51   }
  52 }
  53
  54 pub fn diesel_option_overwrite_to_url(
  55   opt: &Option<String>,
  56 ) -> Result<Option<Option<DbUrl>>, LemmyError> {
  57   match opt.as_ref().map(|s| s.as_str()) {
  58     // An empty string is an erase
  59     Some("") => Ok(Some(None)),
  60     Some(str_url) => match Url::parse(str_url) {
  61       Ok(url) => Ok(Some(Some(url.into()))),
  62       Err(e) => Err(LemmyError::from_error_message(e, "invalid_url")),
  63     },
  64     None => Ok(None),
  65   }
  66 }
  67
  68 embed_migrations!();
  69
  70 pub fn establish_unpooled_connection() -> PgConnection {
  71   let db_url = match get_database_url_from_env() {
  72     Ok(url) => url,
  73     Err(e) => panic!(
  74       "Failed to read database URL from env var LEMMY_DATABASE_URL: {}",
  75       e
  76     ),
  77   };
  78   let conn =
  79     PgConnection::establish(&db_url).unwrap_or_else(|_| panic!("Error connecting to {}", db_url));
  80   embedded_migrations::run(&conn).expect("load migrations");
  81   conn
  82 }
  83
  84 pub fn naive_now() -> NaiveDateTime {
  85   chrono::prelude::Utc::now().naive_utc()
  86 }
  87
  88 static EMAIL_REGEX: Lazy<Regex> = Lazy::new(|| {
  89   Regex::new(r"^[a-zA-Z0-9.!#$%&’*+/=?^_`{|}~-]+@[a-zA-Z0-9-]+(?:\.[a-zA-Z0-9-]+)*$")
  90     .expect("compile email regex")
  91 });
  92
  93 pub mod functions {
  94   use diesel::sql_types::*;
  95
  96   sql_function! {
  97     fn hot_rank(score: BigInt, time: Timestamp) -> Integer;
  98   }
  99
 100   sql_function!(fn lower(x: Text) -> Text);
 101 }
 102
 103 impl<DB: Backend> ToSql<Text, DB> for DbUrl
 104 where
 105   String: ToSql<Text, DB>,
 106 {
 107   fn to_sql<W: Write>(&self, out: &mut Output<W, DB>) -> diesel::serialize::Result {
 108     self.0.to_string().to_sql(out)
 109   }
 110 }
 111
 112 impl<DB: Backend> FromSql<Text, DB> for DbUrl
 113 where
 114   String: FromSql<Text, DB>,
 115 {
 116   fn from_sql(bytes: Option<&DB::RawValue>) -> diesel::deserialize::Result<Self> {
 117     let str = String::from_sql(bytes)?;
 118     Ok(DbUrl(Url::parse(&str)?))
 119   }
 120 }
 121
 122 impl<Kind> From<ObjectId<Kind>> for DbUrl
 123 where
 124   Kind: ApubObject + Send + 'static,
 125   for<'de2> <Kind as ApubObject>::ApubType: serde::Deserialize<'de2>,
 126 {
 127   fn from(id: ObjectId<Kind>) -> Self {
 128     DbUrl(id.into())
 129   }
 130 }
 131
 132 #[cfg(test)]
 133 mod tests {
 134   use super::{fuzzy_search, *};
 135   use crate::utils::is_email_regex;
 136
 137   #[test]
 138   fn test_fuzzy_search() {
 139     let test = "This %is% _a_ fuzzy search";
 140     assert_eq!(
 141       fuzzy_search(test),
 142       "%This%\\%is\\%%\\_a\\_%fuzzy%search%".to_string()
 143     );
 144   }
 145
 146   #[test]
 147   fn test_email() {
 148     assert!(is_email_regex("gush@gmail.com"));
 149     assert!(!is_email_regex("nada_neutho"));
 150   }
 151
 152   #[test]
 153   fn test_diesel_option_overwrite() {
 154     assert_eq!(diesel_option_overwrite(&None), None);
 155     assert_eq!(diesel_option_overwrite(&Some("".to_string())), Some(None));
 156     assert_eq!(
 157       diesel_option_overwrite(&Some("test".to_string())),
 158       Some(Some("test".to_string()))
 159     );
 160   }
 161
 162   #[test]
 163   fn test_diesel_option_overwrite_to_url() {
 164     assert!(matches!(diesel_option_overwrite_to_url(&None), Ok(None)));
 165     assert!(matches!(
 166       diesel_option_overwrite_to_url(&Some("".to_string())),
 167       Ok(Some(None))
 168     ));
 169     assert!(matches!(
 170       diesel_option_overwrite_to_url(&Some("invalid_url".to_string())),
 171       Err(_)
 172     ));
 173     let example_url = "https://example.com";
 174     assert!(matches!(
 175       diesel_option_overwrite_to_url(&Some(example_url.to_string())),
 176       Ok(Some(Some(url))) if url == Url::parse(example_url).unwrap().into()
 177     ));
 178   }
 179 }