feat: /reconnect command

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,119 @@
+use std::sync::Arc;
+
+use serenity::{
+    builder::CreateApplicationCommand, futures::StreamExt, model::prelude::application_command,
+    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,
+    },
+    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?;
+
+                let command = Arc::new(command.clone());
+                let channel = command.channel_id.to_channel(&ctx).await?.guild().unwrap();
+                tokio::spawn(async move {
+                    let mut messages = command.channel_id.messages_iter(&ctx).boxed();
+                    let mut total = 0;
+                    let att = loop {
+                        total += 1;
+                        if total > 20 {
+                            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, &command, "Failed to get message", format!("Couldn't load a message:\n```rust\n{error}\n```"), Some(false)).await.unwrap();
+                                    break None;
+                                }
+                            }
+                        } else {
+                            break None;
+                        }
+                    };
+                    if let Some(att) = att {
+                        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, &command, "Session loading!", "Bot has successfully created a new session, registered it with this channel and is now loading your export.", None).await.unwrap();
+                                        if let Err(err) = conn.import_from_attachment(&ctx, &command, &att).await {
+                                            error!(error = %err, "Error importing from attachment")
+                                        }
+                                        show(&ctx, &channel, conn.conn_type, &config).await.unwrap();
+                                    },
+                                    Err(e) => {
+                                        CmdError::RegisterDB(e).edit(&ctx, &command).await.unwrap();
+                                    }
+                                }
+                            },
+                            Err(err) => {
+                                error!(error = %err, "Error creating DB instance");
+                                CmdError::CreateDB(err).edit(&ctx, &command).await.unwrap();
+                            },
+                        }
+                    } else {
+                        ephemeral_interaction_edit(&ctx, &command, "No export found!", "Bot could not find any .surql attachments in the last 20 messages.", Some(false)).await.unwrap();
+                    }
+                }.in_current_span());
+
+                Ok(())
+            }
+        }
+        None => CmdError::NoGuild.reply(&ctx, command).await,
+    }
+}
+
+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/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,
+        ctx: &Context,
+        i: impl ToInteraction,
+        attachment: &Attachment,
+    ) -> Result<(), anyhow::Error> {
+        ephemeral_interaction_edit(
+            ctx,
+            i,
+            "Downloading attachment",
+            format!("Now downloading `{}`, please wait.", attachment.filename),
+            None,
+        )
+        .await?;
+        match attachment.download().await {
+            Ok(bytes) => {
+                ephemeral_interaction_edit(ctx, i, "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(ctx, 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(ctx, i).await?;
+                        Ok(())
+                    }
+                }
+            }
+            Err(err) => {
+                CmdError::AttachmentDownload(err.into())
+                    .edit(ctx, 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))
     }

src/utils.rs

@@ -1,4 +1,4 @@
-use std::{borrow::Cow, cmp::Ordering};
+use std::{borrow::Cow, cmp::Ordering, sync::Arc};
 
 use once_cell::sync::Lazy;
 use serenity::{
@@ -52,6 +52,8 @@ pub enum CmdError {
     ExportTooLarge,
     BadQuery(surrealdb::Error),
     AttachmentDownload(anyhow::Error),
+    CreateDB(anyhow::Error),
+    RegisterDB(anyhow::Error),
 }
 
 impl CmdError {
@@ -136,6 +138,14 @@ impl CmdError {
                 "Attachment download failed".into(),
                 format!("There was an error while loading the attachment:\n```rust\n{e}\n```").into(),
             ),
+            CmdError::CreateDB(e) => (
+                "Database creation failed".into(),
+                format!("There was an error while creating the database:\n```rust\n{e}\n```").into(),
+            ),
+            CmdError::RegisterDB(e) => (
+                "Database registration failed".into(),
+                format!("There was an error while registering the database:\n```rust\n{e}\n```").into(),
+            )
         }
     }
 
@@ -159,7 +169,7 @@ impl CmdError {
 }
 
 /// ToInteraction is a trait that allows for easy conversion of different interaction types to a tuple of the interaction id and token.
-pub trait ToInteraction {
+pub trait ToInteraction: std::marker::Copy {
     fn to_interaction(&self) -> (&InteractionId, &str);
 }
 
@@ -169,6 +179,12 @@ impl ToInteraction for &ApplicationCommandInteraction {
     }
 }
 
+impl ToInteraction for &Arc<ApplicationCommandInteraction> {
+    fn to_interaction(&self) -> (&InteractionId, &str) {
+        (&self.id, &self.token)
+    }
+}
+
 impl ToInteraction for &MessageComponentInteraction {
     fn to_interaction(&self) -> (&InteractionId, &str) {
         (&self.id, &self.token)
@@ -429,20 +445,21 @@ pub async fn register_db(
     config: Config,
     conn_type: ConnType,
     require_query: bool,
-) -> Result<(), anyhow::Error> {
+) -> Result<Conn, anyhow::Error> {
     info!("Registering a new database");
-    DBCONNS.lock().await.insert(
-        *channel.id.as_u64(),
-        crate::Conn {
-            db,
-            last_used: Instant::now(),
-            conn_type,
-            ttl: config.ttl,
-            pretty: config.pretty,
-            json: config.json,
-            require_query,
-        },
-    );
+    let conn = crate::Conn {
+        db,
+        last_used: Instant::now(),
+        conn_type,
+        ttl: config.ttl,
+        pretty: config.pretty,
+        json: config.json,
+        require_query,
+    };
+    DBCONNS
+        .lock()
+        .await
+        .insert(*channel.id.as_u64(), conn.clone());
 
     tokio::spawn(
         async move {
@@ -467,7 +484,7 @@ pub async fn register_db(
         }
         .in_current_span(),
     );
-    Ok(())
+    Ok(conn)
 }
 
 pub async fn respond(