diff options
Diffstat (limited to 'src/db/mod.rs')
| -rw-r--r-- | src/db/mod.rs | 338 |
1 files changed, 212 insertions, 126 deletions
diff --git a/src/db/mod.rs b/src/db/mod.rs index e8e6596..5f19d54 100644 --- a/src/db/mod.rs +++ b/src/db/mod.rs @@ -1,9 +1,6 @@ -use std::{ - convert::AsRef, - env, - str::FromStr, -}; - +pub use self::models::*; +use self::schema::*; +use crate::db::schema::memes::title; use anyhow::{ anyhow, Error, @@ -22,7 +19,12 @@ use deadpool_postgres::{ }; use diesel::{ prelude::*, + row::NamedRow, + BoolExpressionMethods, + ExpressionMethods, NotFound, + PgTextExpressionMethods, + QueryDsl, }; use diesel_async::{ pooled_connection::{ @@ -35,11 +37,19 @@ use diesel_async::{ RunQueryDsl, }; use grate::tracing; +use rand::Rng; use serenity::FutureExt; +use std::{ + convert::{ + AsRef, + From, + }, + env, + str::FromStr, +}; +use tokio_postgres::types::FromSql; -pub use self::models::*; -use self::schema::*; - +pub mod manual_migrate; mod models; mod schema; @@ -72,7 +82,7 @@ lazy_static::lazy_static! { #[inline] pub async fn connection() --> Result<diesel_async::pooled_connection::deadpool::Object<AsyncPgConnection>> { + -> Result<diesel_async::pooled_connection::deadpool::Object<AsyncPgConnection>> { POOL.get() .then(|mut conn| async move { if let Ok(ref mut conn) = conn { @@ -105,16 +115,30 @@ async fn raw_connection() -> Result<deadpool_postgres::Object> { RAW_CONN_MGR.get().await.map_err(Error::from) } -pub async fn find_meme<T: AsRef<str>>(conn: &mut AsyncPgConnection, search: T) -> Result<Meme> { +pub async fn find_meme<T: AsRef<str>>( + conn: &mut AsyncPgConnection, + search: T, + guild_id: u64, +) -> Result<Meme> { let search = search.as_ref(); + let guild_match = memes::guild.eq(guild_id as i64); - let mut meme = memes::table.filter(memes::title.eq(search)).limit(1).first::<Meme>(conn).await; + let mut meme = memes::table + .filter(memes::title.eq(search).and(guild_match)) + .limit(1) + .first::<Meme>(conn) + .await; if let Err(NotFound) = meme { let format_search = format!("%{}%", search); meme = memes::table - .filter(memes::title.ilike(&format_search).or(memes::content.ilike(&format_search))) + .filter( + memes::title + .ilike(&format_search) + .or(memes::content.ilike(&format_search)) + .and(guild_match), + ) .limit(1) .first::<Meme>(conn) .await; @@ -127,9 +151,11 @@ pub async fn query_meme<T: AsRef<str>>( search: T, user_id: Option<u64>, age_desc: bool, + guild: u64, ) -> Result<Vec<(Meme, Metadata)>> { let raw_conn = raw_connection().await?; + let guild = guild as i64; let search = format!("%{}%", search.as_ref()); let rows = raw_conn @@ -140,7 +166,7 @@ pub async fn query_meme<T: AsRef<str>>( FROM memes INNER JOIN metadata ON memes.metadata_id = metadata.id WHERE (memes.title ILIKE $1 OR memes.content ILIKE $1) - AND (metadata.created_by = $2 OR $3) + AND (metadata.created_by = $2 OR $3) AND memes.guild = $4 ORDER BY metadata.created {} LIMIT 100 "#, @@ -150,7 +176,7 @@ pub async fn query_meme<T: AsRef<str>>( "ASC" }, ), - &[&search, &(user_id.unwrap_or(0) as i64), &user_id.is_none()], + &[&search, &(user_id.unwrap_or(0) as i64), &user_id.is_none(), &guild], ) .await?; @@ -158,17 +184,18 @@ pub async fn query_meme<T: AsRef<str>>( .iter() .map(|row| { let meme = Meme { - id: row.get(0), - title: row.get(1), - content: row.get(2), - image_id: row.get(3), - audio_id: row.get(4), + id: row.get(0), + guild, + title: row.get(1), + content: row.get(2), + image_id: row.get(3), + audio_id: row.get(4), metadata_id: row.get(5), }; let metadata = Metadata { - id: row.get(5), - created: row.get(6), + id: row.get(5), + created: row.get(6), created_by: row.get(7), }; @@ -196,7 +223,7 @@ pub async fn del_memes_by_userid( meme_deletes AS ( DELETE FROM memes WHERE metadata_id IN (SELECT id FROM metadata WHERE created_by = $1) - RETURNING id, metadata_id + RETURNING id, guild, metadata_id ) , image_deletes AS ( DELETE FROM images @@ -209,8 +236,8 @@ pub async fn del_memes_by_userid( RETURNING id, metadata_id ) , tombstone_inserts AS ( - INSERT INTO tombstones (meme_id, metadata_id, deleted_at, deleted_by) - SELECT id, metadata_id, NOW(), $2 FROM meme_deletes + INSERT INTO tombstones (meme_id, guild, metadata_id, deleted_at, deleted_by) + SELECT id, guild, metadata_id, NOW(), $2 FROM meme_deletes RETURNING meme_id, metadata_id ) @@ -229,12 +256,16 @@ pub async fn delete_meme<T: AsRef<str>>( conn: &mut AsyncPgConnection, search: T, deleted_by: u64, + guild: u64, ) -> Result<()> { let search = search.as_ref(); conn.transaction::<(), Error, _>(|tx| { async move { - let deleted = memes::table.filter(memes::title.eq(search)).first::<Meme>(tx).await?; + let deleted = memes::table + .filter(memes::title.eq(search).and(memes::guild.eq(guild as i64))) + .first::<Meme>(tx) + .await?; diesel::delete(memes::table).filter(memes::id.eq(deleted.id)).execute(tx).await?; @@ -263,23 +294,29 @@ pub async fn delete_meme<T: AsRef<str>>( } let tombstone = NewTombstone { - deleted_by: deleted_by as i64, + guild: guild as i64, + deleted_by: deleted_by as i64, metadata_id: deleted.metadata_id, - meme_id: deleted.id, + meme_id: deleted.id, }; let _ = diesel::insert_into(tombstones::table).values(&tombstone).execute(tx).await?; Ok(()) } - .scope_boxed() + .scope_boxed() }) - .await + .await } -pub async fn rare_meme(conn: &mut AsyncPgConnection, audio: bool) -> Result<Meme> { +pub async fn rare_meme( + conn: &mut AsyncPgConnection, + audio: bool, + guild: u64, +) -> Result<Option<Meme>> { use rand::prelude::*; + let guild = guild as i64; let raw_conn = raw_connection().await?; let rows = raw_conn @@ -291,6 +328,7 @@ pub async fn rare_meme(conn: &mut AsyncPgConnection, audio: bool) -> Result<Meme meme_id, COUNT(*) AS ct FROM invocation_records + WHERE guild = $3 GROUP BY meme_id ), aggregate AS ( @@ -301,7 +339,7 @@ pub async fn rare_meme(conn: &mut AsyncPgConnection, audio: bool) -> Result<Meme FROM meme_count RIGHT JOIN memes ON memes.id = meme_count.meme_id INNER JOIN metadata ON metadata.id = memes.metadata_id - WHERE (memes.audio_id IS NULL) = $1 OR $2 + WHERE memes.guild = $3 AND ((memes.audio_id IS NULL) = $1 OR $2) ), least_used AS ( SELECT @@ -315,7 +353,7 @@ pub async fn rare_meme(conn: &mut AsyncPgConnection, audio: bool) -> Result<Meme FROM least_used LIMIT 100; "#, - &[&!audio, &audio], + &[&!audio, &audio, &guild], ) .await?; @@ -325,7 +363,7 @@ pub async fn rare_meme(conn: &mut AsyncPgConnection, audio: bool) -> Result<Meme .collect::<Vec<_>>(); if elems.is_empty() { - return Err(anyhow!("no rare memes found")); + return Ok(None); } let target_prob = { @@ -339,10 +377,14 @@ pub async fn rare_meme(conn: &mut AsyncPgConnection, audio: bool) -> Result<Meme .ok_or_else(|| anyhow!("couldn't locate meme satisfying target probability"))? .0; - Meme::find(conn, meme_id).await + Meme::find(conn, meme_id).await.map(Some) } -pub async fn rand_meme(conn: &mut AsyncPgConnection, audio: bool) -> Result<Meme> { +pub async fn rand_meme( + conn: &mut AsyncPgConnection, + audio: bool, + guild: u64, +) -> Result<Option<Meme>> { use rand::{ seq::SliceRandom, thread_rng, @@ -352,10 +394,12 @@ pub async fn rand_meme(conn: &mut AsyncPgConnection, audio: bool) -> Result<Meme memes::table .select(memes::id) .filter( - memes::content - .is_not_null() - .or(memes::image_id.is_not_null()) - .or(memes::audio_id.is_not_null()), + memes::guild.eq(guild as i64).and( + memes::content + .is_not_null() + .or(memes::image_id.is_not_null()) + .or(memes::audio_id.is_not_null()), + ), ) .load(conn) .await @@ -363,18 +407,26 @@ pub async fn rand_meme(conn: &mut AsyncPgConnection, audio: bool) -> Result<Meme } else { memes::table .select(memes::id) - .filter(memes::content.is_not_null().or(memes::image_id.is_not_null())) + .filter( + memes::guild + .eq(guild as i64) + .and(memes::content.is_not_null().or(memes::image_id.is_not_null())), + ) .load(conn) .await .map_err(Error::from)? }; + if ids.is_empty() { + return Ok(None); + } + let id = ids.choose(&mut thread_rng()).ok_or_else(|| anyhow!("couldn't load meme"))?; - memes::table.find(id).first::<Meme>(conn).await.map_err(Error::from) + memes::table.find(id).first::<Meme>(conn).await.map_err(Error::from).map(Some) } -pub async fn rand_audio_meme(conn: &mut AsyncPgConnection) -> Result<Meme> { +pub async fn rand_audio_meme(conn: &mut AsyncPgConnection, guild: u64) -> Result<Option<Meme>> { use rand::{ seq::SliceRandom, thread_rng, @@ -382,17 +434,21 @@ pub async fn rand_audio_meme(conn: &mut AsyncPgConnection) -> Result<Meme> { let ids: Vec<i32> = memes::table .select(memes::id) - .filter(memes::audio_id.is_not_null()) + .filter(memes::audio_id.is_not_null().and(memes::guild.eq(guild as i64))) .load(conn) .await .map_err(Error::from)?; + if ids.is_empty() { + return Ok(None); + } + let id = ids.choose(&mut thread_rng()).ok_or_else(|| anyhow!("couldn't load audio meme"))?; - memes::table.find(id).first::<Meme>(conn).await.map_err(Error::from) + memes::table.find(id).first::<Meme>(conn).await.map_err(Error::from).map(Some) } -pub async fn rand_silent_meme(conn: &mut AsyncPgConnection) -> Result<Meme> { +pub async fn rand_silent_meme(conn: &mut AsyncPgConnection, guild: u64) -> Result<Option<Meme>> { use rand::{ seq::SliceRandom, thread_rng, @@ -400,218 +456,246 @@ pub async fn rand_silent_meme(conn: &mut AsyncPgConnection) -> Result<Meme> { let ids: Vec<i32> = memes::table .select(memes::id) - .filter(memes::audio_id.is_null()) + .filter(memes::audio_id.is_null().and(memes::guild.eq(guild as i64))) .load(conn) .await .map_err(Error::from)?; + if ids.is_empty() { + return Ok(None); + } + let id = ids.choose(&mut thread_rng()).ok_or_else(|| anyhow!("couldn't load audio meme"))?; - memes::table.find(id).first::<Meme>(conn).await.map_err(Error::from) + memes::table.find(id).first::<Meme>(conn).await.map_err(Error::from).map(Some) } #[derive(Debug, Clone)] pub struct Stats { - pub memes_overall: usize, - pub audio_memes: usize, - pub image_memes: usize, - pub started_recording: DateTime<Utc>, - pub total_meme_invocations: usize, - pub audio_meme_invocations: usize, + pub memes_overall: usize, + pub audio_memes: usize, + pub image_memes: usize, + pub started_recording: Option<DateTime<Utc>>, + pub total_meme_invocations: usize, + pub audio_meme_invocations: usize, pub random_meme_invocations: usize, - pub most_active_day: NaiveDate, + pub most_active_day: Option<NaiveDate>, pub most_active_day_count: usize, - pub most_audio_active_day: NaiveDate, + pub most_audio_active_day: Option<NaiveDate>, pub most_audio_active_count: usize, - pub most_random_meme_user: u64, - pub most_random_meme_user_count: usize, - pub most_directly_named_meme_user: u64, + pub most_random_meme_user: Option<u64>, + pub most_random_meme_user_count: usize, + pub most_directly_named_meme_user: Option<u64>, pub most_directly_named_meme_count: usize, - pub most_popular_named_meme: String, + pub most_popular_named_meme: Option<String>, pub most_popular_named_meme_count: usize, - pub most_popular_random_meme: String, + pub most_popular_random_meme: Option<String>, pub most_popular_random_meme_count: usize, - pub most_popular_meme_overall: String, + pub most_popular_meme_overall: Option<String>, pub most_popular_meme_overall_count: usize, } -pub async fn stats(conn: &mut AsyncPgConnection) -> Result<Stats> { +pub async fn stats(conn: &mut AsyncPgConnection, guild: u64) -> Result<Stats> { use chrono::NaiveDateTime; use diesel::dsl::{ count, count_star, }; + let guild = guild as i64; + + let guild_match = memes::guild.eq(guild); + let invoke_guild_match = invocation_records::guild.eq(guild); + #[inline] fn to_utc(ndt: NaiveDateTime) -> DateTime<Utc> { Utc.from_utc_datetime(&ndt) } - let total_count: i64 = - memes::table.select(count_star()).first(conn).await.map_err(Error::from)?; + let total_count: i64 = memes::table + .select(count_star()) + .filter(guild_match) + .first(conn) + .await + .map_err(Error::from)?; let image_count: i64 = memes::table .select(count(memes::image_id)) - .filter(memes::image_id.is_not_null()) + .filter(memes::image_id.is_not_null().and(guild_match)) .first(conn) .await .map_err(Error::from)?; let audio_count: i64 = memes::table .select(count(memes::audio_id)) - .filter(memes::audio_id.is_not_null()) + .filter(memes::audio_id.is_not_null().and(guild_match)) .first(conn) .await .map_err(Error::from)?; - let started_recording: NaiveDateTime = invocation_records::table + let started_recording = invocation_records::table .select(invocation_records::time) + .filter(invoke_guild_match) .order(invocation_records::time) .first(conn) .await - .map_err(Error::from)?; + .optional(); + + tracing::debug!(?started_recording); - let started_recording = to_utc(started_recording); + let started_recording: Option<NaiveDateTime> = started_recording.map_err(Error::from)?; + let started_recording = started_recording.map(to_utc); - let total_meme_invocations: i64 = - invocation_records::table.select(count_star()).first(conn).await.map_err(Error::from)?; + let total_meme_invocations: i64 = invocation_records::table + .select(count_star()) + .filter(invoke_guild_match) + .first(conn) + .await + .map_err(Error::from)?; let audio_meme_invocations: i64 = invocation_records::table .inner_join(memes::table) .select(count_star()) - .filter(memes::audio_id.is_not_null()) + .filter(memes::audio_id.is_not_null().and(guild_match)) .first(conn) .await .map_err(Error::from)?; let random_meme_invocations: i64 = invocation_records::table .select(count_star()) - .filter(invocation_records::random.eq(true)) + .filter(invocation_records::random.eq(true).and(invoke_guild_match)) .first(conn) .await .map_err(Error::from)?; let raw_conn = raw_connection().await?; + #[inline] + fn option_first_count<T>(rows: Vec<tokio_postgres::Row>) -> (Option<T>, i64) + where + for<'a> T: FromSql<'a>, + { + let Some(row) = <[_]>::first(&rows) else { + return (None, 0); + }; + + (row.get(0), row.get(1)) + } + let row = raw_conn - .query_one( + .query( r#" SELECT DATE(time) as dt, COUNT(*) FROM invocation_records + WHERE guild = $1 GROUP BY dt ORDER BY COUNT(*) DESC LIMIT 1; "#, - &[], + &[&guild], ) .await?; - let most_active_day = row.get(0); - let most_active_day_count: i64 = row.get(1); + let (most_active_day, most_active_day_count) = option_first_count(row); let row = raw_conn - .query_one( + .query( r#" SELECT DATE(time) as dt, COUNT(*) FROM invocation_records - INNER JOIN memes ON invocation_records.meme_id = memes.id - WHERE memes.audio_id IS NOT NULL + INNER JOIN memes ON (invocation_records.meme_id = memes.id AND invocation_records.guild = memes.guild) + WHERE memes.audio_id IS NOT NULL AND memes.guild = $1 GROUP BY dt ORDER BY COUNT(*) DESC LIMIT 1; "#, - &[], + &[&guild], ) .await?; - let most_active_audio_day = row.get(0); - let most_active_audio_day_count: i64 = row.get(1); + let (most_active_audio_day, most_active_audio_day_count) = option_first_count(row); let row = raw_conn - .query_one( + .query( r#" SELECT user_id, COUNT(*) FROM invocation_records - WHERE random IS TRUE + WHERE random IS TRUE AND guild = $1 GROUP BY user_id ORDER BY COUNT(*) DESC LIMIT 1; "#, - &[], + &[&guild], ) .await?; - let most_random_invoker: i64 = row.get(0); - let most_random_invoker_count: i64 = row.get(1); + let (most_random_invoker, most_random_invoker_count) = option_first_count::<i64>(row); let row = raw_conn - .query_one( + .query( r#" SELECT user_id, COUNT(*) FROM invocation_records - WHERE random IS FALSE + WHERE random IS FALSE AND guild = $1 GROUP BY user_id ORDER BY COUNT(*) DESC LIMIT 1; "#, - &[], + &[&guild], ) .await?; - let most_specific_invoker: i64 = row.get(0); - let most_specific_invoker_count: i64 = row.get(1); + let (most_specific_invoker, most_specific_invoker_count) = option_first_count::<i64>(row); let row = raw_conn - .query_one( + .query( r#" SELECT memes.title, COUNT(*) FROM invocation_records - INNER JOIN memes ON meme_id = memes.id - WHERE random IS FALSE + INNER JOIN memes ON (meme_id = memes.id AND invocation_records.guild = memes.guild) + WHERE random IS FALSE AND memes.guild = $1 GROUP BY memes.title ORDER BY COUNT(*) DESC LIMIT 1; "#, - &[], + &[&guild], ) .await?; - let most_requested_meme = row.get(0); - let most_requested_meme_count: i64 = row.get(1); + let (most_requested_meme, most_requested_meme_count) = option_first_count(row); let row = raw_conn - .query_one( + .query( r#" SELECT memes.title, COUNT(*) FROM invocation_records - INNER JOIN memes ON meme_id = memes.id - WHERE random IS TRUE + INNER JOIN memes ON (meme_id = memes.id AND invocation_records.guild = memes.guild) + WHERE random IS TRUE AND memes.guild = $1 GROUP BY memes.title ORDER BY COUNT(*) DESC LIMIT 1; "#, - &[], + &[&guild], ) .await?; - let most_random_meme = row.get(0); - let most_random_meme_count: i64 = row.get(1); + let (most_random_meme, most_random_meme_count) = option_first_count(row); let row = raw_conn - .query_one( + .query( r#" SELECT memes.title, COUNT(*) FROM invocation_records - INNER JOIN memes ON meme_id = memes.id + INNER JOIN memes ON (meme_id = memes.id AND invocation_records.guild = memes.guild) + WHERE memes.guild = $1 GROUP BY memes.title ORDER BY COUNT(*) DESC LIMIT 1; "#, - &[], + &[&guild], ) .await?; - let most_invoked_meme = row.get(0); - let most_invoked_meme_count: i64 = row.get(1); + let (most_invoked_meme, most_invoked_meme_count) = option_first_count(row); Ok(Stats { memes_overall: total_count as usize, @@ -624,12 +708,13 @@ pub async fn stats(conn: &mut AsyncPgConnection) -> Result<Stats> { most_active_day, most_active_day_count: most_active_day_count as usize, + most_audio_active_day: most_active_audio_day, most_audio_active_count: most_active_audio_day_count as usize, - most_random_meme_user: most_random_invoker as u64, + most_random_meme_user: most_random_invoker.map(|x| x as u64), most_random_meme_user_count: most_random_invoker_count as usize, - most_directly_named_meme_user: most_specific_invoker as u64, + most_directly_named_meme_user: most_specific_invoker.map(|x| x as u64), most_directly_named_meme_count: most_specific_invoker_count as usize, most_popular_named_meme: most_requested_meme, @@ -645,33 +730,34 @@ pub async fn stats(conn: &mut AsyncPgConnection) -> Result<Stats> { #[derive(Clone, Debug, Hash, PartialEq, Eq, Default)] pub struct MemerInfo { - pub user_id: u64, - pub random_memes: usize, - pub specific_memes: usize, - pub most_used_meme: String, + pub user_id: u64, + pub random_memes: usize, + pub specific_memes: usize, + pub most_used_meme: String, pub most_used_meme_count: usize, } -pub async fn memers() -> Result<Vec<MemerInfo>> { +pub async fn memers(guild: u64) -> Result<Vec<MemerInfo>> { + let guild = guild as i64; let raw_conn = raw_connection().await?; let rows = raw_conn.query(r#" WITH random_count AS ( SELECT user_id, COUNT(*) as count FROM invocation_records - WHERE random = TRUE + WHERE random = TRUE AND guild = $1 GROUP BY user_id ), specific_count AS ( SELECT user_id, COUNT(*) as count FROM invocation_records - WHERE random = FALSE + WHERE random = FALSE AND guild = $1 GROUP BY user_id ), user_meme_counts AS ( SELECT user_id, meme_id, COUNT(meme_id) as meme_count FROM invocation_records - WHERE EXISTS (SELECT id FROM memes WHERE id = invocation_records.meme_id) + WHERE EXISTS (SELECT id FROM memes WHERE id = invocation_records.meme_id) AND guild = $1 GROUP BY user_id, meme_id ORDER BY user_id, meme_count DESC ), @@ -692,7 +778,7 @@ pub async fn memers() -> Result<Vec<MemerInfo>> { INNER JOIN specific_count ON specific_count.user_id = random_count.user_id INNER JOIN memes ON memes.id = most_memed.meme_id ORDER BY (random_count.count + specific_count.count) DESC - "#, &[]).await?; + "#, &[&guild]).await?; let result = rows .iter() @@ -704,10 +790,10 @@ pub async fn memers() -> Result<Vec<MemerInfo>> { let most_memed_count: i64 = row.get(4); MemerInfo { - user_id: user_id as u64, - random_memes: random_count as usize, - specific_memes: specific_count as usize, - most_used_meme: most_memed_meme, + user_id: user_id as u64, + random_memes: random_count as usize, + specific_memes: specific_count as usize, + most_used_meme: most_memed_meme, most_used_meme_count: most_memed_count as usize, } }) |
