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