crates/db_schema/src/utils.rs

   1 use crate::{newtypes::DbUrl, CommentSortType, SortType};
   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 pub fn post_to_comment_sort_type(sort: SortType) -> CommentSortType {
 122   match sort {
 123     SortType::Active | SortType::Hot => CommentSortType::Hot,
 124     SortType::New | SortType::NewComments | SortType::MostComments => CommentSortType::New,
 125     SortType::Old => CommentSortType::Old,
 126     SortType::TopDay
 127     | SortType::TopAll
 128     | SortType::TopWeek
 129     | SortType::TopYear
 130     | SortType::TopMonth => CommentSortType::Top,
 131   }
 132 }
 133
 134 static EMAIL_REGEX: Lazy<Regex> = Lazy::new(|| {
 135   Regex::new(r"^[a-zA-Z0-9.!#$%&’*+/=?^_`{|}~-]+@[a-zA-Z0-9-]+(?:\.[a-zA-Z0-9-]+)*$")
 136     .expect("compile email regex")
 137 });
 138
 139 pub mod functions {
 140   use diesel::sql_types::*;
 141
 142   sql_function! {
 143     fn hot_rank(score: BigInt, time: Timestamp) -> Integer;
 144   }
 145
 146   sql_function!(fn lower(x: Text) -> Text);
 147 }
 148
 149 impl<DB: Backend> ToSql<Text, DB> for DbUrl
 150 where
 151   String: ToSql<Text, DB>,
 152 {
 153   fn to_sql<W: Write>(&self, out: &mut Output<W, DB>) -> diesel::serialize::Result {
 154     self.0.to_string().to_sql(out)
 155   }
 156 }
 157
 158 impl<DB: Backend> FromSql<Text, DB> for DbUrl
 159 where
 160   String: FromSql<Text, DB>,
 161 {
 162   fn from_sql(bytes: Option<&DB::RawValue>) -> diesel::deserialize::Result<Self> {
 163     let str = String::from_sql(bytes)?;
 164     Ok(DbUrl(Url::parse(&str)?))
 165   }
 166 }
 167
 168 impl<Kind> From<ObjectId<Kind>> for DbUrl
 169 where
 170   Kind: ApubObject + Send + 'static,
 171   for<'de2> <Kind as ApubObject>::ApubType: serde::Deserialize<'de2>,
 172 {
 173   fn from(id: ObjectId<Kind>) -> Self {
 174     DbUrl(id.into())
 175   }
 176 }
 177
 178 #[cfg(test)]
 179 mod tests {
 180   use super::{fuzzy_search, *};
 181   use crate::utils::is_email_regex;
 182
 183   #[test]
 184   fn test_fuzzy_search() {
 185     let test = "This %is% _a_ fuzzy search";
 186     assert_eq!(
 187       fuzzy_search(test),
 188       "%This%\\%is\\%%\\_a\\_%fuzzy%search%".to_string()
 189     );
 190   }
 191
 192   #[test]
 193   fn test_email() {
 194     assert!(is_email_regex("gush@gmail.com"));
 195     assert!(!is_email_regex("nada_neutho"));
 196   }
 197
 198   #[test]
 199   fn test_diesel_option_overwrite() {
 200     assert_eq!(diesel_option_overwrite(&None), None);
 201     assert_eq!(diesel_option_overwrite(&Some("".to_string())), Some(None));
 202     assert_eq!(
 203       diesel_option_overwrite(&Some("test".to_string())),
 204       Some(Some("test".to_string()))
 205     );
 206   }
 207
 208   #[test]
 209   fn test_diesel_option_overwrite_to_url() {
 210     assert!(matches!(diesel_option_overwrite_to_url(&None), Ok(None)));
 211     assert!(matches!(
 212       diesel_option_overwrite_to_url(&Some("".to_string())),
 213       Ok(Some(None))
 214     ));
 215     assert!(matches!(
 216       diesel_option_overwrite_to_url(&Some("invalid_url".to_string())),
 217       Err(_)
 218     ));
 219     let example_url = "https://example.com";
 220     assert!(matches!(
 221       diesel_option_overwrite_to_url(&Some(example_url.to_string())),
 222       Ok(Some(Some(url))) if url == Url::parse(example_url).unwrap().into()
 223     ));
 224   }
 225 }