database alteration in-flight

5 files changed, 251 insertions, 9 deletions

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<bool>; 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::<String>()?.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::<String>()?;
+                        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::<ContentLength>()
+                            .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::<ContentLength>()
+                            .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::<String>()?;
+                    },
+                    _ => {
+                        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<PgConnection> {
     CONN_MGR.connect().map_err(Error::from)

 }

 

-pub fn find_text(conn: &PgConnection, search: String) -> Result<Meme> {

+pub fn find_meme<T: AsRef<str>>(conn: &PgConnection, search: T) -> Result<Meme> {

     use diesel::dsl::sql;

+

+    let search = search.as_ref();

     let format_search = format!("%{}%", search);

 

     memes::table

@@ -30,18 +33,32 @@ pub fn find_text(conn: &PgConnection, search: String) -> Result<Meme> {
         .map_err(Error::from)

 }

 

-pub fn find_audio(conn: &PgConnection, search: String) -> Result<Meme> {

+pub fn find_text<T: AsRef<str>>(conn: &PgConnection, search: T) -> Result<Meme> {

+    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::<Meme>(conn)

+        .map_err(Error::from)

+}

+

+pub fn find_audio<T: AsRef<str>>(conn: &PgConnection, search: T) -> Result<Meme> {

+    let format_search = format!("%{}%", search.as_ref());

+

+    memes::table

         .filter(memes::title.ilike(format_search).and(memes::audio_id.is_not_null()))

         .limit(1)

         .first::<Meme>(conn)

         .map_err(Error::from)

 }

 

-pub fn find_image(conn: &PgConnection, search: String) -> Result<Meme> {

-    let format_search = format!("%{}%", search);

+pub fn find_image<T: AsRef<str>>(conn: &PgConnection, search: T) -> Result<Meme> {

+    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<Result<Image>> {

-        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<Result<Audio>> {

-        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<Meme> {

+        let metadata = Metadata::create(conn, by_user)?;

+

+        self.metadata_id = metadata.id;

+

+        ::diesel::insert_into(memes::table)

+            .values(&self)

+            .get_result::<Meme>(conn)

+            .map_err(Error::from)

+    }

+}

+

 

 #[derive(Queryable, Identifiable, PartialEq, Debug)]

 #[table_name="audio"]

 pub struct Audio {

     pub id: i32,

     pub data: Vec<u8>,

+    pub data_hash: Vec<u8>,

     pub metadata_id: i32,

 }

 

+impl Audio {

+    pub fn create(conn: &PgConnection, data: Vec<u8>, by_user: u64) -> Result<i32> {

+        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::<i32>(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<u8>,

+    pub data_hash: Vec<u8>,

     pub metadata_id: i32,

 }

 

@@ -57,13 +103,46 @@ pub struct NewAudio {
 pub struct Image {

     pub id: i32,

     pub data: Vec<u8>,

+    pub data_hash: Vec<u8>,

     pub metadata_id: i32,

 }

 

+impl Image {

+    pub fn create(conn: &PgConnection, data: Vec<u8>, by_user: u64) -> Result<i32> {

+        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::<i32>(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<u8>,

+    pub data_hash: Vec<u8>,

     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<Metadata> {

+        ::diesel::insert_into(metadata::table)

+            .values(&NewMetadata {

+                created_by: by_user as i64,

+            })

+            .get_result::<Metadata>(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<AuditRecord> {

+        ::diesel::insert_into(audit_records::table)

+            .values(&NewAuditRecord {

+                updated_by: by_user as i64,

+                metadata_id: metadata,

+            })

+            .get_result::<AuditRecord>(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};