From d6bea61fa917d257219a43386c88a4f04cf82408 Mon Sep 17 00:00:00 2001 From: Nathan Perry Date: Fri, 6 Apr 2018 19:22:55 -0400 Subject: database alteration in-flight --- src/commands/meme.rs | 121 ++++++++++++++++++++++++++++++++++++++++++++++++++- src/db/mod.rs | 25 +++++++++-- src/db/models.rs | 108 +++++++++++++++++++++++++++++++++++++++++++-- src/db/schema.rs | 2 + src/main.rs | 4 ++ 5 files changed, 251 insertions(+), 9 deletions(-) (limited to 'src') diff --git a/src/commands/meme.rs b/src/commands/meme.rs index 89a979d..277b5f8 100644 --- a/src/commands/meme.rs +++ b/src/commands/meme.rs @@ -1,7 +1,20 @@ +use std::time::Duration; + use rand::{thread_rng, distributions::{Weighted, WeightedChoice, Distribution}}; use serenity::http::AttachmentType; use serenity::builder::CreateMessage; use diesel::PgConnection; +use reqwest::{ + Client, + header::{ + Headers, + ContentLength, + UserAgent, + AcceptEncoding, + Encoding, + qitem, + } +}; use super::*; use super::playback::CtxExt; @@ -28,10 +41,116 @@ static mut TTS_WEIGHTS: [Weighted; 2] = [ ]; command!(meme(ctx, msg, args) { - if args.len() == 0 { + if args.len_quoted() == 0 { rand_meme(ctx, msg)?; return Ok(()); } + + macro_rules! next { () => { args.single_quoted::()?.to_lowercase() }; } + + match next!().as_ref() { + "add" => { // e.g.: !thulani meme add title [image IMAGE] [audio|sound AUDIO] [text TEXT...] + let mut new_meme = NewMeme { + title: next!(), + content: None, + image_id: None, + audio_id: None, + metadata_id: 0, + }; + + let mut headers = Headers::new(); + headers.set(AcceptEncoding(vec!(qitem(Encoding::Gzip)))); + headers.set(UserAgent::new("Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:59.0) Gecko/20100101 Firefox/59.0)")); + + let client = Client::builder() + .default_headers(headers) + .timeout(Duration::from_secs(5)) + .build()?; + + let conn = connection()?; + + while args.len() > 0 { + match next!().as_ref() { + "text" => new_meme.content = Some(args.full().to_owned()), + "image" => { + let url = args.single_quoted::()?; + let resp = client.head(&url).send()?; + + if !resp.status().is_success() { + send(msg.channel_id, "pick a better url next time thanks", msg.tts)?; + return Ok(()); + } + + let len = resp.headers().get::() + .map(|ct_len| **ct_len) + .unwrap_or(0); + + if len > 20_000_000 { + send(msg.channel_id, "are you trying to bankrupt my disk space", msg.tts)?; + return Ok(()); + } + + let mut resp = client.get(&url).send()?; + + if !resp.status().is_success() { + send(msg.channel_id, "pick a better url next time thanks", msg.tts)?; + return Ok(()); + } + + let len = resp.headers().get::() + .map(|ct_len| **ct_len) + .unwrap_or(0); + + if len > 20_000_000 { + send(msg.channel_id, "are you fucking serious", msg.tts)?; + return Ok(()); + } + + if !resp.status().is_success() { + send(msg.channel_id, "bad link reeeeee", msg.tts)?; + return Ok(()); + } + + let mut data = Vec::with_capacity(len as usize); + ::std::io::copy(&mut resp, &mut data)?; + + let image_id = Image::create(&conn, data, msg.author.id.0)?; + new_meme.image_id = Some(image_id); + }, + "audio" | "sound" => { + let _url = args.single_quoted::()?; + }, + _ => { + send(msg.channel_id, "hueh?", msg.tts)?; + return Ok(()); + } + } + } + + if new_meme.content.is_none() && new_meme.image_id.is_none() && new_meme.audio_id.is_none() { + send(msg.channel_id, "haha it's empty lol xdddd", msg.tts)?; + return Ok(()); + } + + new_meme.save(&conn, msg.author.id.0)?; + send(msg.channel_id, "i hate my job", msg.tts)?; + }, + "delete" | "remove" => { + send(msg.channel_id, "hwaet", msg.tts)?; + }, + search => { + let conn = connection()?; + let mem = match find_meme(&conn, search) { + Ok(x) => x, + Err(_) => { + send(msg.channel_id, "what in ryan's name", msg.tts)?; + return Ok(()); + }, + }; + + send_meme(ctx, &mem, &conn, msg)?; + } + } }); fn rand_meme(ctx: &Context, message: &Message) -> Result<()> { diff --git a/src/db/mod.rs b/src/db/mod.rs index 258d5a5..0862883 100644 --- a/src/db/mod.rs +++ b/src/db/mod.rs @@ -1,4 +1,5 @@ use std::env; +use std::convert::AsRef; use diesel::prelude::*; use diesel::r2d2::{ConnectionManager, ManageConnection}; @@ -19,8 +20,10 @@ pub fn connection() -> Result { CONN_MGR.connect().map_err(Error::from) } -pub fn find_text(conn: &PgConnection, search: String) -> Result { +pub fn find_meme>(conn: &PgConnection, search: T) -> Result { use diesel::dsl::sql; + + let search = search.as_ref(); let format_search = format!("%{}%", search); memes::table @@ -30,9 +33,23 @@ pub fn find_text(conn: &PgConnection, search: String) -> Result { .map_err(Error::from) } -pub fn find_audio(conn: &PgConnection, search: String) -> Result { +pub fn find_text>(conn: &PgConnection, search: T) -> Result { + use diesel::dsl::sql; + + let search = search.as_ref(); let format_search = format!("%{}%", search); + memes::table + .filter((memes::title.ilike(&format_search).or(sql(&format!("content ILIKE %{}%", search)))) + .and(memes::content.is_not_null())) + .limit(1) + .first::(conn) + .map_err(Error::from) +} + +pub fn find_audio>(conn: &PgConnection, search: T) -> Result { + let format_search = format!("%{}%", search.as_ref()); + memes::table .filter(memes::title.ilike(format_search).and(memes::audio_id.is_not_null())) .limit(1) @@ -40,8 +57,8 @@ pub fn find_audio(conn: &PgConnection, search: String) -> Result { .map_err(Error::from) } -pub fn find_image(conn: &PgConnection, search: String) -> Result { - let format_search = format!("%{}%", search); +pub fn find_image>(conn: &PgConnection, search: T) -> Result { + let format_search = format!("%{}%", search.as_ref()); memes::table .filter(memes::title.ilike(format_search).and(memes::image_id.is_not_null())) diff --git a/src/db/models.rs b/src/db/models.rs index 7479fff..32ccc4b 100644 --- a/src/db/models.rs +++ b/src/db/models.rs @@ -17,11 +17,11 @@ pub struct Meme { impl Meme { pub fn image(&self, conn: &PgConnection) -> Option> { - self.image_id.map(|x: i32| images::table.find(x).first(conn).map_err(Error::from)) + self.image_id.map(|x: i32| images::table.filter(images::id.eq(x)).first(conn).map_err(Error::from)) } pub fn audio(&self, conn: &PgConnection) -> Option> { - self.audio_id.map(|x: i32| audio::table.find(x).first(conn).map_err(Error::from)) + self.audio_id.map(|x: i32| audio::table.filter(audio::id.eq(x)).first(conn).map_err(Error::from)) } } @@ -35,19 +35,65 @@ pub struct NewMeme { pub metadata_id: i32, } +impl NewMeme { + pub fn save(mut self, conn: &PgConnection, by_user: u64) -> Result { + let metadata = Metadata::create(conn, by_user)?; + + self.metadata_id = metadata.id; + + ::diesel::insert_into(memes::table) + .values(&self) + .get_result::(conn) + .map_err(Error::from) + } +} + #[derive(Queryable, Identifiable, PartialEq, Debug)] #[table_name="audio"] pub struct Audio { pub id: i32, pub data: Vec, + pub data_hash: Vec, pub metadata_id: i32, } +impl Audio { + pub fn create(conn: &PgConnection, data: Vec, by_user: u64) -> Result { + let mut data_hash = ::sha1::Sha1::new(); + data_hash.update(&data); + let data_hash = data_hash.digest().bytes().to_vec(); + + let id = audio::table + .select(audio::id) + .filter(audio::data_hash.eq(&data_hash)) + .get_results::(conn)?; + + if let Some(id) = id.first() { + return Ok(*id); + } + + let metadata = Metadata::create(conn, by_user)?; + + let new_audio = NewAudio { + data, + data_hash, + metadata_id: metadata.id, + }; + + ::diesel::insert_into(audio::table) + .values(&new_audio) + .returning(audio::id) + .get_result(conn) + .map_err(Error::from) + } +} + #[derive(Insertable, PartialEq, Debug)] #[table_name="audio"] pub struct NewAudio { pub data: Vec, + pub data_hash: Vec, pub metadata_id: i32, } @@ -57,13 +103,46 @@ pub struct NewAudio { pub struct Image { pub id: i32, pub data: Vec, + pub data_hash: Vec, pub metadata_id: i32, } +impl Image { + pub fn create(conn: &PgConnection, data: Vec, by_user: u64) -> Result { + let mut data_hash = ::sha1::Sha1::new(); + data_hash.update(&data); + let data_hash = data_hash.digest().bytes().to_vec(); + + let id = images::table + .select(images::id) + .filter(images::data_hash.eq(&data_hash)) + .get_results::(conn)?; + + if let Some(id) = id.first() { + return Ok(*id); + } + + let metadata = Metadata::create(conn, by_user)?; + + let new_image = NewImage { + data, + data_hash, + metadata_id: metadata.id, + }; + + ::diesel::insert_into(images::table) + .values(&new_image) + .returning(images::id) + .get_result(conn) + .map_err(Error::from) + } +} + #[derive(Insertable, PartialEq, Debug)] #[table_name="images"] pub struct NewImage { pub data: Vec, + pub data_hash: Vec, pub metadata_id: i32, } @@ -76,10 +155,20 @@ pub struct Metadata { pub created_by: i64, } +impl Metadata { + pub fn create(conn: &PgConnection, by_user: u64) -> Result { + ::diesel::insert_into(metadata::table) + .values(&NewMetadata { + created_by: by_user as i64, + }) + .get_result::(conn) + .map_err(Error::from) + } +} + #[derive(Insertable, PartialEq, Debug)] #[table_name="metadata"] pub struct NewMetadata { - pub created: NaiveDateTime, pub created_by: i64, } @@ -93,10 +182,21 @@ pub struct AuditRecord { pub metadata_id: i32, } +impl AuditRecord { + pub fn create(conn: &PgConnection, metadata: i32, by_user: u64) -> Result { + ::diesel::insert_into(audit_records::table) + .values(&NewAuditRecord { + updated_by: by_user as i64, + metadata_id: metadata, + }) + .get_result::(conn) + .map_err(Error::from) + } +} + #[derive(Insertable, PartialEq, Debug)] #[table_name="audit_records"] pub struct NewAuditRecord { - pub updated: NaiveDateTime, pub updated_by: i64, pub metadata_id: i32, } diff --git a/src/db/schema.rs b/src/db/schema.rs index cf15dcd..d4ba0fa 100644 --- a/src/db/schema.rs +++ b/src/db/schema.rs @@ -3,6 +3,7 @@ table! { id -> Int4, data -> Bytea, metadata_id -> Int4, + data_hash -> Bytea, } } @@ -20,6 +21,7 @@ table! { id -> Int4, data -> Bytea, metadata_id -> Int4, + data_hash -> Bytea, } } diff --git a/src/main.rs b/src/main.rs index 551449b..80ed991 100644 --- a/src/main.rs +++ b/src/main.rs @@ -14,6 +14,8 @@ extern crate typemap; extern crate url; extern crate rand; extern crate either; +extern crate reqwest; +extern crate sha1; use commands::register_commands; use dotenv::dotenv; @@ -127,6 +129,8 @@ fn main() { const MIN_RUN_DURATION: Duration = Duration::from_secs(120); + info!("starting"); + dotenv().ok(); use fern::colors::{Color, ColoredLevelConfig}; -- cgit v1.3.1