feat: /reconnect command and Reconnect button on exports (#19)

Co-authored-by: Raphael Darley <raphael@raphaeldarley.com>

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "surreal_bot"
-version = "0.3.2"
+version = "0.3.3"
 edition = "2021"
 authors = [
     "Raphael Darley <raphael.darley@surrealdb.com>",

src/commands/mod.rs

@@ -11,6 +11,7 @@ pub mod export;
 pub mod load;
 pub mod q;
 pub mod query;
+pub mod reconnect;
 pub mod share;
 
 use serenity::builder::CreateApplicationCommands;
@@ -29,6 +30,7 @@ pub fn register_all(commands: &mut CreateApplicationCommands) -> &mut CreateAppl
         .create_application_command(|command| configure_channel::register(command))
         .create_application_command(|command| query::register(command))
         .create_application_command(|command| q::register(command))
+        .create_application_command(|command| reconnect::register(command))
         .create_application_command(|command| connect::register(command))
         .create_application_command(|command| export::register(command))
 }

src/commands/query.rs

@@ -33,7 +33,7 @@ pub async fn run(
             conn.query(
                 &ctx,
                 &command.channel_id,
-                Some(&command),
+                Some(command),
                 &command.user,
                 query,
                 None,

src/commands/reconnect.rs

@@ -0,0 +1,175 @@
+use std::sync::Arc;
+
+use anyhow::anyhow;
+use serenity::{
+    builder::CreateApplicationCommand,
+    futures::StreamExt,
+    model::prelude::{application_command, Attachment, ChannelId},
+    prelude::Context,
+};
+use tracing::Instrument;
+
+use crate::{
+    components::configurable_session::show,
+    config::Config,
+    utils::{
+        create_db_instance, ephemeral_interaction, ephemeral_interaction_edit, register_db,
+        CmdError, ToInteraction,
+    },
+    DB, DBCONNS,
+};
+
+pub async fn run(
+    command: &application_command::ApplicationCommandInteraction,
+    ctx: Context,
+) -> Result<(), anyhow::Error> {
+    match command.guild_id {
+        Some(guild_id) => {
+            if let Some(_) = DBCONNS.lock().await.get_mut(command.channel_id.as_u64()) {
+                CmdError::ExpectedNoSession.reply(&ctx, command).await
+            } else {
+                let result: Result<Option<Config>, surrealdb::Error> =
+                    DB.select(("guild_config", guild_id.to_string())).await;
+
+                let config = match result {
+                    Ok(response) => match response {
+                        Some(c) => c,
+                        None => return CmdError::NoConfig.reply(&ctx, command).await,
+                    },
+                    Err(e) => return CmdError::GetConfig(e).reply(&ctx, command).await,
+                };
+
+                ephemeral_interaction(
+                    &ctx,
+                    command,
+                    "Looking for last export",
+                    "I will now check the last 20 messages for .surql attachments.",
+                    None,
+                )
+                .await?;
+                tokio::spawn(
+                    reconnect(
+                        ctx,
+                        Arc::new(command.clone()),
+                        command.channel_id,
+                        config,
+                        20,
+                    )
+                    .in_current_span(),
+                );
+                Ok(())
+            }
+        }
+        None => CmdError::NoGuild.reply(&ctx, command).await,
+    }
+}
+
+async fn reconnect(
+    ctx: Context,
+    i: impl ToInteraction,
+    channel_id: ChannelId,
+    config: Config,
+    limit: usize,
+) -> Result<(), anyhow::Error> {
+    match find_attachment(&ctx, &i, channel_id, limit).await {
+        Some(att) => new_db_from_attachment(ctx.clone(), i, channel_id, config, att).await,
+        None => {
+            ephemeral_interaction_edit(
+                &ctx,
+                i,
+                "No export found!",
+                "Bot could not find any .surql attachments in the last 20 messages.",
+                Some(false),
+            )
+            .await
+        }
+    }
+}
+
+pub async fn new_db_from_attachment(
+    ctx: Context,
+    i: impl ToInteraction,
+    channel_id: ChannelId,
+    config: Config,
+    att: Attachment,
+) -> Result<(), anyhow::Error> {
+    let channel = channel_id
+        .to_channel(&ctx)
+        .await?
+        .guild()
+        .ok_or(anyhow!("Not in a guild"))?;
+
+    match create_db_instance(&config).await {
+        Ok(db) => {
+            match register_db(
+                ctx.clone(),
+                db.clone(),
+                channel.clone(),
+                config.clone(),
+                crate::ConnType::ConnectedChannel,
+                true,
+            )
+            .await
+            {
+                Ok(conn) => {
+                    ephemeral_interaction_edit(&ctx, i.clone(), "Session loading!", "Successfully created a new session, registered it with this channel and now loading your export.", None).await?;
+                    if let Err(err) = conn.import_from_attachment(&ctx, i.clone(), &att).await {
+                        error!(error = %err, "Error importing from attachment")
+                    }
+                    show(&ctx, &channel, conn.conn_type, &config).await
+                }
+                Err(e) => CmdError::RegisterDB(e).edit(&ctx, i).await,
+            }
+        }
+        Err(err) => {
+            error!(error = %err, "Error creating DB instance");
+            CmdError::CreateDB(err).edit(&ctx, i).await
+        }
+    }
+}
+
+#[tracing::instrument(skip(ctx, i), fields(channel_id = %channel_id, limit = %limit))]
+async fn find_attachment(
+    ctx: &Context,
+    i: &impl ToInteraction,
+    channel_id: ChannelId,
+    limit: usize,
+) -> Option<Attachment> {
+    let mut messages = channel_id.messages_iter(ctx).boxed();
+    let mut total = 0;
+    loop {
+        total += 1;
+        if total > limit {
+            break None;
+        }
+        if let Some(message_result) = messages.next().await {
+            match message_result {
+                Ok(message) => match message.attachments.first() {
+                    Some(att) => break Some(att.clone()),
+                    None => continue,
+                },
+                Err(error) => {
+                    error!(error = %error, "Error getting message");
+                    ephemeral_interaction_edit(
+                        ctx,
+                        i.clone(),
+                        "Failed to get message",
+                        format!("Couldn't load a message:\n```rust\n{error}\n```"),
+                        Some(false),
+                    )
+                    .await
+                    .unwrap();
+                    break None;
+                }
+            }
+        } else {
+            break None;
+        }
+    }
+}
+
+pub fn register(command: &mut CreateApplicationCommand) -> &mut CreateApplicationCommand {
+    command
+        .name("reconnect")
+        .description("Recreates a SurrealDB instance using most recent export and associates it with the current channel")
+}

src/components/configurable_session.rs

@@ -6,7 +6,7 @@ use crate::{
         clean_channel, ephemeral_interaction, user_interaction, CmdError, BOT_VERSION,
         SURREALDB_VERSION,
     },
-    ConnType, BIG_QUERY_SENT_KEY, BIG_QUERY_VARS_KEY, DBCONNS,
+    ConnType, BIG_QUERY_SENT_KEY, BIG_QUERY_VARS_KEY, DB, DBCONNS,
 };
 
 use anyhow::Result;
@@ -257,7 +257,46 @@ pub async fn handle_component(
                 .await?;
         }
         ("reconnect", false) => {
-            // TODO: feature creep but maybe a button to re-create a session after it's been deleted
+            let result: Result<Option<Config>, surrealdb::Error> = DB
+                .select((
+                    "guild_config",
+                    &event.guild_id.unwrap_or_default().to_string(),
+                ))
+                .await;
+            match (event.message.attachments.first(), result) {
+                (Some(att), Ok(Some(config))) => {
+                    ephemeral_interaction(
+                        &ctx,
+                        event,
+                        "Creating new DB",
+                        "Your new DB is being created, this may take a while.",
+                        None,
+                    )
+                    .await?;
+                    tokio::spawn(
+                        crate::commands::reconnect::new_db_from_attachment(
+                            ctx.clone(),
+                            Arc::new(event.clone()),
+                            channel.clone(),
+                            config,
+                            att.clone(),
+                        )
+                        .in_current_span(),
+                    );
+                }
+                (Some(_), Err(err)) => {
+                    CmdError::GetConfig(err).reply(ctx, event).await?;
+                }
+                (Some(_), Ok(None)) => {
+                    CmdError::NoConfig.reply(ctx, event).await?;
+                }
+                (None, _) => {
+                    CmdError::ExpectedAttachment.reply(ctx, event).await?;
+                }
+            }
+        }
+        ("reconnect", true) => {
+            CmdError::ExpectedNoSession.reply(&ctx, event).await?;
         }
         ("rename_thread", _) => {
             let channel_name = channel

src/handler.rs

@@ -102,6 +102,7 @@ impl EventHandler for Handler {
                         }
                         "query" => commands::query::run(&command, ctx.clone()).await,
                         "q" => commands::q::run(&command, ctx.clone()).await,
+                        "reconnect" => commands::reconnect::run(&command, ctx.clone()).await,
                         "connect" => commands::connect::run(&command, ctx.clone()).await,
                         "export" => commands::export::run(&command, ctx.clone()).await,
                         _ => {

src/lib.rs

@@ -15,7 +15,7 @@ use serenity::{
     model::{
         prelude::{
             application_command::ApplicationCommandInteraction, component::ButtonStyle::Primary,
-            AttachmentType, ChannelId,
+            Attachment, AttachmentType, ChannelId,
         },
         user::User,
     },
@@ -24,7 +24,7 @@ use serenity::{
 use surrealdb::{opt::IntoQuery, sql::Value, Error, Response};
 use tokio::sync::Mutex;
 use tokio::time::{Duration, Instant};
-use utils::MAX_FILE_SIZE;
+use utils::{ephemeral_interaction_edit, CmdError, ToInteraction, MAX_FILE_SIZE};
 
 #[macro_use]
 extern crate tracing;
@@ -61,6 +61,47 @@ pub enum ConnType {
 }
 
 impl Conn {
+    pub async fn import_from_attachment(
+        &self,
+        http: impl AsRef<Http>,
+        i: impl ToInteraction,
+        attachment: &Attachment,
+    ) -> Result<(), anyhow::Error> {
+        ephemeral_interaction_edit(
+            &http,
+            i.clone(),
+            "Downloading attachment",
+            format!("Now downloading `{}`, please wait.", attachment.filename),
+            None,
+        )
+        .await?;
+        match attachment.download().await {
+            Ok(bytes) => {
+                ephemeral_interaction_edit(&http, i.clone(), "Downloaded, now importing...", "Your data is currently being loaded, soon you'll be able to query your dataset! \n_Please wait for a confirmation that the dataset is loaded!_", None).await?;
+                match self
+                    .db
+                    .query(String::from_utf8_lossy(&bytes).into_owned())
+                    .await
+                {
+                    Ok(_) => {
+                        ephemeral_interaction_edit(http, i, "Imported successfully!", "Your data has been imported successfully!\nYou can now query your dataset.", Some(true)).await?;
+                        Ok(())
+                    }
+                    Err(why) => {
+                        CmdError::BadQuery(why).edit(http, i).await?;
+                        Ok(())
+                    }
+                }
+            }
+            Err(err) => {
+                CmdError::AttachmentDownload(err.into())
+                    .edit(http, i)
+                    .await?;
+                Ok(())
+            }
+        }
+    }
+
     #[must_use]
     pub async fn export_to_attachment(&self) -> Result<Option<AttachmentType>, anyhow::Error> {
         let mut acc = Vec::new();
@@ -77,7 +118,7 @@ impl Conn {
 
         let reply_attachment = AttachmentType::Bytes {
             data: std::borrow::Cow::Owned(acc),
-            filename: format!("export.surql"),
+            filename: "export.surql".to_string(),
         };
         Ok(Some(reply_attachment))
     }
@@ -258,7 +299,7 @@ pub async fn shutdown(http: impl AsRef<Http>) -> Result<(), anyhow::Error> {
                 let res = channel.send_message(&http, |m| {
                     m.embed(|e| {
                         e.title("Pre-shutdown DB Exported successfully").description("Sorry! The bot had to go offline for maintenance, your session has been exported. You can find the .surql file attached.\nYou can either use `/reconnect` and load a new session with it when the bot is back online, or use it locally with `surreal import` CLI.").color(0x00ff00)
-                    }).add_file(attchment)
+                    }).add_file(attchment).components(|c| c.create_action_row(|r| r.create_button(|b| b.label("Reconnect").custom_id("configurable_session:reconnect").style(Primary).emoji('📦'))))
                 }).await;
                 if let Err(why) = res {
                     errors.push(why.to_string())