Add ephemeral message table and delete them on startup

Signed-off-by: trivernis <trivernis@protonmail.com>

Cargo.lock

@@ -193,7 +193,7 @@ dependencies = [
 
 [[package]]
 name = "bot-database"
-version = "0.4.0"
+version = "0.5.0"
 dependencies = [
  "chrono",
  "diesel",

bot-database/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "bot-database"
-version = "0.4.0"
+version = "0.5.0"
 authors = ["trivernis <trivernis@protonmail.com>"]
 edition = "2018"
 

bot-database/migrations/2021-04-16-064403_create_ephemeral_messages/down.sql

@@ -0,0 +1,2 @@
+-- This file should undo anything in `up.sql`
+DROP TABLE ephemeral_messages;

bot-database/migrations/2021-04-16-064403_create_ephemeral_messages/up.sql

@@ -0,0 +1,7 @@
+-- Your SQL goes here
+CREATE TABLE ephemeral_messages (
+    channel_id BIGINT NOT NULL,
+    message_id BIGINT NOT NULL,
+    timeout TIMESTAMP NOT NULL,
+    PRIMARY KEY (channel_id, message_id)
+)

bot-database/src/database/ephemeral_messages.rs

@@ -0,0 +1,58 @@
+use std::time::SystemTime;
+
+use diesel::prelude::*;
+use diesel::{delete, insert_into};
+use tokio_diesel::*;
+
+use crate::error::DatabaseResult;
+use crate::models::*;
+use crate::schema::*;
+use crate::Database;
+
+impl Database {
+    /// Adds a command statistic to the database
+    pub async fn add_ephemeral_message(
+        &self,
+        channel_id: u64,
+        message_id: u64,
+        timeout: SystemTime,
+    ) -> DatabaseResult<()> {
+        use ephemeral_messages::dsl;
+        insert_into(dsl::ephemeral_messages)
+            .values(EphemeralMessageInsert {
+                channel_id: channel_id as i64,
+                message_id: message_id as i64,
+                timeout,
+            })
+            .execute_async(&self.pool)
+            .await?;
+
+        Ok(())
+    }
+
+    /// Returns a vec of all ephemeral messages
+    pub async fn get_ephemeral_messages(&self) -> DatabaseResult<Vec<EphemeralMessage>> {
+        use ephemeral_messages::dsl;
+        let messages: Vec<EphemeralMessage> = dsl::ephemeral_messages
+            .load_async::<EphemeralMessage>(&self.pool)
+            .await?;
+
+        Ok(messages)
+    }
+
+    /// Deletes a single ephemeral message
+    pub async fn delete_ephemeral_message(
+        &self,
+        channel_id: i64,
+        message_id: i64,
+    ) -> DatabaseResult<()> {
+        use ephemeral_messages::dsl;
+        delete(dsl::ephemeral_messages)
+            .filter(dsl::channel_id.eq(channel_id))
+            .filter(dsl::message_id.eq(message_id))
+            .execute_async(&self.pool)
+            .await?;
+
+        Ok(())
+    }
+}

bot-database/src/database/mod.rs

@@ -1,3 +1,4 @@
+pub use ephemeral_messages::*;
 pub use gifs::*;
 pub use guild_playlists::*;
 pub use guild_playlists::*;
@@ -6,6 +7,7 @@ pub use youtube_songs::*;
 
 use crate::PoolConnection;
 
+mod ephemeral_messages;
 mod gifs;
 mod guild_playlists;
 mod guild_settings;

bot-database/src/models.rs

@@ -77,3 +77,18 @@ pub struct YoutubeSongInsert {
     pub album: String,
     pub url: String,
 }
+
+#[derive(Queryable, Debug, Clone)]
+pub struct EphemeralMessage {
+    pub channel_id: i64,
+    pub message_id: i64,
+    pub timeout: SystemTime,
+}
+
+#[derive(Insertable, Debug)]
+#[table_name = "ephemeral_messages"]
+pub struct EphemeralMessageInsert {
+    pub channel_id: i64,
+    pub message_id: i64,
+    pub timeout: SystemTime,
+}

bot-database/src/schema.rs

@@ -1,3 +1,11 @@
+table! {
+    ephemeral_messages (channel_id, message_id) {
+        channel_id -> Int8,
+        message_id -> Int8,
+        timeout -> Timestamp,
+    }
+}
+
 table! {
     gifs (id) {
         id -> Int8,
@@ -47,6 +55,7 @@ table! {
 }
 
 allow_tables_to_appear_in_same_query!(
+    ephemeral_messages,
     gifs,
     guild_playlists,
     guild_settings,

src/handler.rs

@@ -9,6 +9,7 @@ use serenity::model::voice::VoiceState;
 use serenity::prelude::*;
 
 use crate::commands::music::get_queue_for_guild;
+use crate::utils::delete_messages_from_database;
 use bot_serenityutils::menu::{
     handle_message_delete, handle_message_delete_bulk, handle_reaction_add, handle_reaction_remove,
     start_update_loop,
@@ -21,6 +22,9 @@ impl EventHandler for Handler {
     async fn cache_ready(&self, ctx: Context, _: Vec<GuildId>) {
         log::info!("Cache Ready");
         start_update_loop(&ctx).await;
+        if let Err(e) = delete_messages_from_database(&ctx).await {
+            log::error!("Failed to delete expired messages {:?}", e);
+        }
     }
 
     /// Fired when a message was deleted

src/messages/mod.rs

@@ -1,3 +1,24 @@
+use crate::utils::context_data::get_database_from_context;
+use crate::utils::error::BotResult;
+use bot_serenityutils::core::MessageHandle;
+use serenity::client::Context;
+use std::time::{Duration, SystemTime};
+
 pub mod gifs;
 pub mod music;
 pub mod sauce;
+
+/// Adds an ephemeral message to the database
+pub async fn add_ephemeral_handle_to_database(
+    ctx: &Context,
+    handle: MessageHandle,
+    timeout: Duration,
+) -> BotResult<()> {
+    let timeout = SystemTime::now() + timeout;
+    let database = get_database_from_context(ctx).await;
+    database
+        .add_ephemeral_message(handle.channel_id, handle.message_id, timeout)
+        .await?;
+
+    Ok(())
+}

src/messages/music/now_playing.rs

@@ -6,6 +6,7 @@ use serenity::model::prelude::ChannelId;
 use songbird::input::Metadata;
 
 use crate::commands::music::{get_queue_for_guild, get_voice_manager, is_dj};
+use crate::messages::add_ephemeral_handle_to_database;
 use crate::providers::music::add_youtube_song_to_database;
 use crate::providers::music::queue::MusicQueue;
 use crate::utils::context_data::{DatabaseContainer, Store};
@@ -68,6 +69,8 @@ pub async fn create_now_playing_msg(
         .build(ctx, channel_id)
         .await?;
 
+    add_ephemeral_handle_to_database(ctx, *handle.read().await, Duration::from_secs(0)).await?;
+
     Ok(handle)
 }
 

src/messages/sauce.rs

@@ -40,7 +40,7 @@ pub async fn show_sauce_menu(
             .add_pages(pages)
             .build(ctx, msg.channel_id)
             .await?;
-    };
+    }
 
     Ok(())
 }

src/utils/mod.rs

@@ -1,7 +1,12 @@
 use serenity::client::Context;
 use serenity::model::channel::Message;
 
+use crate::utils::context_data::get_database_from_context;
 use crate::utils::error::BotResult;
+use std::ops::Add;
+use std::sync::Arc;
+use std::time::SystemTime;
+use tokio::time::Instant;
 
 pub(crate) mod context_data;
 pub(crate) mod error;
@@ -24,3 +29,45 @@ pub async fn get_previous_message_or_reply(
 
     Ok(referenced)
 }
+
+/// Deletes all expired ephemeral messages that are stored in the database
+pub async fn delete_messages_from_database(ctx: &Context) -> BotResult<()> {
+    let database = get_database_from_context(ctx).await;
+    let messages = database.get_ephemeral_messages().await?;
+
+    for message in messages {
+        if message.timeout <= SystemTime::now() {
+            log::debug!("Deleting message {:?}", message);
+            let _ = ctx
+                .http
+                .delete_message(message.channel_id as u64, message.message_id as u64)
+                .await;
+            database
+                .delete_ephemeral_message(message.channel_id, message.message_id)
+                .await?;
+        } else {
+            let http = Arc::clone(&ctx.http);
+            let database = database.clone();
+            log::debug!(
+                "Creating future to delete ephemeral message {:?} later",
+                message
+            );
+
+            tokio::spawn(async move {
+                tokio::time::sleep_until(
+                    Instant::now().add(message.timeout.duration_since(SystemTime::now()).unwrap()),
+                )
+                .await;
+                log::debug!("Deleting message {:?}", message);
+                let _ = http
+                    .delete_message(message.channel_id as u64, message.message_id as u64)
+                    .await;
+                let _ = database
+                    .delete_ephemeral_message(message.channel_id, message.message_id)
+                    .await;
+            });
+        }
+    }
+
+    Ok(())
+}