surreal testing

I don't want normies tyring to scrape the whole server

.gitea/workflows/rust.yaml

@@ -0,0 +1,20 @@
+name: Test Rust project
+on: [push]
+jobs:
+  test:
+    runs-on: ${{ matrix.os }}
+    strategy:
+      matrix:
+        os: [ubuntu-latest]
+        rust: [stable]
+
+    steps:
+    - uses: https://git.oliveratkinson.net/Oliver/setup-rust-action@master
+      with:
+        rust-version: ${{ matrix.rust }}
+    - uses: actions/checkout@master
+    - name: Run tests
+      run: cargo test --verbose
+    - name: Clippy
+      run: cargo clippy
+

.gitignore

@@ -1,3 +1,4 @@
 /target
 /.vscode
 .env
+server.json

Cargo.toml

@@ -10,10 +10,12 @@ edition = "2021"
 
 [dependencies]
 tokio = { version = "1.21.2", features = ["macros", "rt-multi-thread"] }
-# songbird = { version = "0.3.2", features = ["yt-dlp"] }
-poise = { version = "0.6", features = ["cache"] }
+poise = { features = ["cache"], git="https://github.com/Rushmore75/poise.git", rev="6afaf44a791cd3ff590c89fda26c94b4411b3266" }
 dotenv = "0.15.0"
 anyhow = "1.0.75"
 once_cell = "1.19.0"
 tracing = "0.1.40"
 tracing-subscriber = { version = "0.3.18", features = ["env-filter"] }
+serde_json = "1.0.120"
+serde = { version = "1.0.204", features = ["derive"] }
+surrealdb = "1.5.4"

Dockerfile

@@ -1,15 +0,0 @@
-# Build the execuitible 
-FROM rustlang/rust:nightly as builder
-ENV RUSTFLAGS=""
-
-WORKDIR /bot
-COPY . .
-RUN cargo build --release
-
-# Now make the runtime container
-FROM debian:bookworm-slim
-
-COPY --from=builder /bot/target/release/discord_egress /usr/local/bin/discord_egress
-COPY Cargo.lock /
-
-CMD ["/usr/local/bin/discord_egress"]

README.md

@@ -2,6 +2,26 @@
 
 egress a discord server
 
+### NOTE:
+Using custom version of poise (the discord library) because I need to use a specific version of serenity (the discord library for poise). Here's how to set it up:
+
+1) Clone Poise next to this repo:
+
+```shell
+git clone https://github.com/serenity-rs/poise.git
+```
+
+2) Change the version of serenity that poise uses:
+
+```diff
+[dependencies.serenity]
+- version = "x.xx.x"
++ git = "https://github.com/serenity-rs/serenity.git"
++ rev = "e34f4491ee3a0b20ea8dc30cbc77b257f402f692"
+```
+
+> Just make sure that the rev is a sha1 commit hash sometime after 7/25/24
+
 ## Getting started
 * goto: [Discord applicatoins](https://discord.com/developers/applications) and
 create an application.

docker-compose.yaml

@@ -1,30 +0,0 @@
-version: '3.1'
-name: discord-egress
-services:
-  bot:
-    container_name: discord-egress_bot
-    build: https://git.oliveratkinson.net/Oliver/discord-egress.git
-    restart: always
-    environment:
-      # These will read from the .env file
-      DISCORD_INTENTS: ${DISCORD_INTENTS}
-      DISCORD_TOKEN: ${DISCORD_TOKEN}
-      DISCORD_ID: ${DISCORD_ID}
-    networks:
-      - external
-      - internal
-  redis:
-    container_name: discord-egress_redis
-    image: redis
-    restart: always
-    ports:
-      - 6379:6379
-    networks:
-      - internal
-
-networks:
-  internal:
-    driver: bridge
-    internal: true
-  external:
-    driver: bridge

src/command.rs

@@ -1,19 +1,26 @@
-use std::{collections::HashMap, fmt::Display, sync::Arc};
+use std::{collections::HashMap, fmt::Display, fs, hint::black_box, sync::Arc};
 
 use crate::Context;
 use anyhow::Error;
-use poise::{serenity_prelude::{Cache, CacheHttp, ChannelId, ChannelType, GetMessages, GuildChannel, Http, Message}, CreateReply};
-use tracing::{debug, error, trace};
-
-struct Server {
-    channels: Vec<Channel>,
-    orphanage: Vec<GuildChannel>,
-    needs_clean: bool,
-}
+use poise::{
+    serenity_prelude::{
+        Cache, CacheHttp, ChannelId, ChannelType, GetMessages, GuildChannel, Http, Message,
+    },
+    CreateReply,
+};
+use serde::{Deserialize, Serialize};
+use surrealdb::{engine::remote::ws::Ws, opt::auth::Root, sql::Thing, Surreal};
+use tokio::time::Instant;
+use tracing::{debug, error, info, trace};
 
+#[derive(Serialize, Deserialize)]
 struct Channel {
     this: GuildChannel,
+    #[serde(skip_serializing_if = "Vec::is_empty")]
+    #[serde(default)]
     children: Vec<Channel>,
+    #[serde(skip_serializing_if = "Vec::is_empty")]
+    #[serde(default)]
     messages: Vec<Message>,
 }
 
@@ -26,16 +33,34 @@ impl Channel {
             messages: Vec::new(),
         }
     }
+
+    fn is_category(&self) -> bool {
+        !self.children.is_empty() && self.messages.is_empty()
+    }
+}
+
+#[derive(Serialize, Deserialize)]
+struct Server {
+    #[serde(default)]
+    name: String,
+    channels: Vec<Channel>,
+    #[serde(skip_serializing_if = "Vec::is_empty")]
+    #[serde(default)]
+    orphanage: Vec<GuildChannel>,
+    #[serde(skip_serializing)]
+    #[serde(default)]
+    needs_clean: bool,
 }
 
 impl Display for Server {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-
         fn print(f: &mut std::fmt::Formatter<'_>, tab: usize, channel: &Vec<Channel>) {
             for i in channel {
-                for _ in 0..tab { let _ = write!(f, "\t"); }
-                let _ = writeln!(f, "{} {}", prefix(i.this.kind),i.this.name);
-                print(f, tab+1, &i.children);
+                for _ in 0..tab {
+                    let _ = write!(f, "\t");
+                }
+                let _ = writeln!(f, "{} {}", prefix(i.this.kind), i.this.name);
+                print(f, tab + 1, &i.children);
             }
         }
 
@@ -57,13 +82,12 @@ impl Display for Server {
             }
         }
 
-        print(f, 0, &self.channels); 
-
+        print(f, 0, &self.channels);
 
         if self.needs_clean {
             let _ = writeln!(f, "Orphans: (please clean() before displaying...)");
             for i in &self.orphanage {
-                let _ = write!(f, "{} {},", prefix(i.kind),i.name);
+                let _ = write!(f, "{} {},", prefix(i.kind), i.name);
             }
         }
 
@@ -72,42 +96,41 @@ impl Display for Server {
 }
 
 impl Server {
-
-    fn new() -> Self {
+    fn new(name: impl Into<String>) -> Self {
         Self {
+            name: name.into(),
             channels: Vec::new(),
-            orphanage: Vec::new(), 
-            needs_clean: false
+            orphanage: Vec::new(),
+            needs_clean: false,
         }
     }
-    
+
     fn search_by_id<'a>(target: &'a mut Vec<Channel>, find: &ChannelId) -> Option<&'a mut Channel> {
         for child in target {
             if child.this.id == *find {
                 return Some(child);
             }
-            match Self::search_by_id(&mut child.children, find) {
-                Some(x) => return Some(x),
-                None => {},
+            if let Some(x) = Self::search_by_id(&mut child.children, find) {
+                return Some(x);
             }
         }
         None
-   }
+    }
 
     fn add(&mut self, insert: GuildChannel) {
         // make sure the new item wants a parent
         if let Some(parent_id) = &insert.parent_id {
             // find the parent (needs to go thru all nodes)
 
-            match Self::search_by_id(&mut self.channels, &parent_id) {
+            match Self::search_by_id(&mut self.channels, parent_id) {
                 Some(parent_node) => {
                     parent_node.children.push(Channel::new(insert));
-                },
+                }
                 None => {
                     // couldn't find parent, store somewhere else until it's parent is added...
                     self.orphanage.push(insert);
                     self.needs_clean = true;
-                },
+                }
             }
         } else {
             self.channels.push(Channel::new(insert));
@@ -116,7 +139,9 @@ impl Server {
 
     /// Cleans out the orphan channels, finding them parents. You'll want to use this before displaying anything.
     fn clean(&mut self) {
-        if !self.needs_clean {return;}
+        if !self.needs_clean {
+            return;
+        }
 
         // Look thru the orphanage and try to find parents
         for orphan in &self.orphanage {
@@ -136,30 +161,42 @@ impl Server {
 
     /// Scrapes messages for all the channels in `self`.
     async fn scrape_all(&mut self) {
+        // TODO add status from this command, just because it's cool
         let cache: (&Arc<Cache>, &Http) = (&Arc::new(Cache::new()), &Http::new(&crate::ENV.token));
         walk_channels(&mut self.channels, cache).await;
 
         /// Recursive walk thru the channels
         async fn walk_channels(all: &mut Vec<Channel>, cache: impl CacheHttp + Clone) {
-            let settings = GetMessages::default().limit(5);
+            // Qty of messages to take at a time, max=100
+            let batch_size = 100;
+
+            let settings = GetMessages::default().limit(batch_size);
             for channel in all {
                 // Clone *should* be cheap - it's Arc under the hood
                 // get the messages
                 get_messages(channel, cache.clone(), settings).await;
-            }           
+            }
 
             /// Get all messages for 1 channel and children
-            async fn get_messages(channel: &mut Channel, cache: impl CacheHttp + Clone, settings: GetMessages) {
+            async fn get_messages(
+                channel: &mut Channel,
+                cache: impl CacheHttp + Clone,
+                settings: GetMessages,
+            ) {
                 // Loop thru all the messages in the channel in batches.
                 // Adding each batch to the current channel's messages the whole time.
                 let mut last_id = channel.this.last_message_id;
                 while let Some(last) = last_id {
-                    match channel.this.messages(cache.clone(), settings.before(last)).await {
+                    match channel
+                        .this
+                        .messages(cache.clone(), settings.before(last))
+                        .await
+                    {
                         Ok(mut ok) => {
                             if ok.is_empty() {
                                 debug!("Reached the beginning of {}", channel.this.name);
                                 // Stop the loop if there are no more messages.
-                                last_id = None;  
+                                last_id = None;
                             } else {
                                 trace!("Adding {} messages to \"{}\"", ok.len(), channel.this.name);
                                 channel.messages.append(&mut ok);
@@ -167,9 +204,16 @@ impl Server {
                                 if let Some(l) = channel.messages.last() {
                                     last_id = Some(l.id);
                                 }
-                            }                        
-                        },
-                        Err(e) => error!("Error while trying to get messages - {e}"),
+                            }
+                        }
+                        Err(e) => {
+                            error!(
+                                "While reading messages in \"{}\" before `{}` - {e}",
+                                channel.this.name, last
+                            );
+                            // Stop reading this channel on an error.
+                            last_id = None;
+                        }
                     }
                 }
                 // Then recurse into children channels
@@ -186,32 +230,287 @@ impl Server {
             for channel in this {
                 total += channel.messages.len();
                 total += walk(&channel.children);
-            };
+            }
             total
         }
         walk(&self.channels)
     }
+
+    async fn to_surreal(&self) -> surrealdb::Result<()> {
+        trace!("Connecting to surrealdb...");
+        // Connect to the server
+        let db = Surreal::new::<Ws>("127.0.0.1:8000").await?;
+        db.signin(Root {
+            username: "root",
+            password: "root",
+        })
+        .await?;
+
+        db.use_ns("egress").use_db(self.name.clone()).await?;
+
+        // =========================================================
+        // Ingress data
+
+        // Data will only be in three layers
+        // Layer 1: Categories (no parent)
+        // Layer 2: Channels (might have parent)
+        // Layer 3: Messages (has parent)
+
+        trace!("Starting ingress...");
+        for cat in self.channels.iter() {
+
+            match cat.this.kind {
+                ChannelType::Text => {
+                    // This is a text channel
+
+                    #[derive(Serialize)]
+                    struct ChannelWrapper {
+                        name: String,
+                        nsfw: bool,
+                        discord_id: u64,
+                        discord_parent_id: Option<u64>,
+                        topic: String,
+                    }
+                    let chan = &cat.this;
+                    let dpi = if let Some(val) = chan.parent_id {Some(val.get())} else { None }; 
+                    let new_channel: Vec<Thing> = db
+                        .create("channel")
+                        .content(ChannelWrapper {
+                            name: chan.name.to_owned(),
+                            nsfw: chan.nsfw,
+                            discord_id: chan.id.get(),
+                            discord_parent_id: dpi,
+                            topic: chan.topic.to_owned().unwrap_or(String::new()),
+                        })
+                        .await?;
+
+                    for msg in cat.messages.iter() {
+
+                        #[derive(Serialize)]
+                        struct Author {
+                            nickname: String,
+                            username: String,
+                            /// B64 encoded string of image (for now)
+                            avatar: String,
+                            id: u64,
+                        }
+
+                        #[derive(Serialize)]
+                        struct Attachment {
+                            content_type: String,
+                            filename: String,
+                            url: String,
+                        }
+
+                        #[derive(Serialize)]
+                        struct Reaction {
+                            count: u64,
+                            emoji: String,
+                        }
+
+                        #[derive(Serialize)]
+                        struct MessageWrapper {
+                            // FIXME learn how to do references
+                            parent: Thing,
+                            author: Author,
+                            content: String,
+                            utc_timestamp: String,
+                            mentions: Vec<u64>,
+                            attachments: Vec<Attachment>,
+                            reactions: Vec<Reaction>,
+                            pinned: bool,
+                        }
+
+                        let _: Vec<Thing> = db
+                            .create("message")
+                            .content(MessageWrapper {
+                                parent: new_channel[0].clone(),
+                                author: Author {
+                                    id: msg.author.id.get(),
+                                    nickname: msg.author.name.to_owned(),
+                                    username: msg.author.global_name.clone().unwrap_or(String::new()),
+                                    avatar: {
+                                        match msg.author.avatar {
+                                            Some(hash) => {
+                                                format!(
+                                                    "https://cdn.discordapp.com/avatars/{}/{}.webp",
+                                                    msg.author.id,
+                                                    hash,
+                                                )
+                                            },
+                                            None => String::new(),
+                                        }
+                                    }
+                                },
+                                content: msg.content.clone(),
+                                utc_timestamp: msg.timestamp.to_utc().to_string(),
+                                mentions: msg.mentions.iter().map(|f| f.id.get()).collect(),
+                                attachments: msg.attachments.iter().map(|f| Attachment {
+                                    content_type: f.content_type.clone().unwrap_or(String::new()),
+                                    filename: f.filename.to_owned(),
+                                    url: f.url.to_owned(),
+                                }).collect(),
+                                reactions: msg.reactions.iter().map(|f| Reaction {
+                                    count: f.count,
+                                    emoji: f.reaction_type.as_data(),
+                                }).collect(),
+                                pinned: msg.pinned,
+                            })
+                            .await?;
+                    }
+                }
+                ChannelType::Private => todo!(),
+                ChannelType::Voice => todo!(),
+                ChannelType::GroupDm => todo!(),
+                ChannelType::Category => todo!(),
+                ChannelType::News => todo!(),
+                ChannelType::NewsThread => todo!(),
+                ChannelType::PublicThread => todo!(),
+                ChannelType::PrivateThread => todo!(),
+                ChannelType::Stage => todo!(),
+                ChannelType::Directory => todo!(),
+                ChannelType::Forum => todo!(),
+                ChannelType::Unknown(_) => todo!(),
+                _ => todo!(),
+            }
+
+            // TODO learn why this is a vec
+            // Do the first iteration of channels a bit different, so as to name it "category".
+            let new_category: Vec<Thing> = db.create("category").content(&cat.this).await?;
+            import_messages(&cat.messages, &new_category[0], &db).await?;
+
+            // Ok, now automatically recurse the rest of the structure and auto import as channels
+            // and messages.
+            import_channel(&cat.children, &new_category[0], &db).await?;
+        }
+
+        async fn import_channel(
+            channels: &Vec<Channel>,
+            parent: &Thing,
+            db: &Surreal<surrealdb::engine::remote::ws::Client>,
+        ) -> surrealdb::Result<()> {
+            for channel in channels.iter() {
+                trace!("Importing channel \"{}\"", channel.this.name);
+                #[derive(Serialize)]
+                struct ChannelWrapper<'a, 'b> {
+                    channel: &'a GuildChannel,
+                    surreal_parent: &'b Thing,
+                }
+                let new_channel: Vec<Thing> = db
+                    .create("channel")
+                    .content(ChannelWrapper {
+                        channel: &channel.this,
+                        surreal_parent: &parent,
+                    })
+                    .await?;
+
+                import_messages(&channel.messages, &new_channel[0], &db).await?;
+                // async recursion - thus box
+                Box::pin(import_channel(&channel.children, &new_channel[0], &db)).await?;
+            }
+            Ok(())
+        }
+
+        async fn import_messages(
+            msgs: &Vec<Message>,
+            parent: &Thing,
+            db: &Surreal<surrealdb::engine::remote::ws::Client>,
+        ) -> surrealdb::Result<()> {
+            trace!("Importing {} messages...", msgs.len());
+            for msg in msgs.iter() {
+                #[derive(Serialize)]
+                struct MessageWrapper<'a, 'b> {
+                    message: &'a Message,
+                    surreal_parent: &'b Thing,
+                }
+
+                let created: Vec<Thing> = db
+                    .create("message")
+                    .content(MessageWrapper {
+                        message: &msg,
+                        surreal_parent: &parent,
+                    })
+                    .await?;
+
+                trace!("Imported message {:?}", created);
+            }
+            Ok(())
+        }
+
+        // Data is all in
+        // =========================================================
+
+        Ok(())
+    }
 }
 
-#[poise::command(slash_command, rename = "scrape_all", guild_only)]
-pub async fn scrape_all(ctx: Context<'_>) -> Result<(), Error> {
-    let guild = ctx.guild_id().unwrap().to_partial_guild(ctx.serenity_context()).await.unwrap();
+#[poise::command(slash_command, rename = "scrape_all", guild_only, owners_only)]
+pub async fn scrape_all(ctx: Context<'_>, pretty_print: bool) -> Result<(), Error> {
+    let guild = ctx
+        .guild_id()
+        .unwrap()
+        .to_partial_guild(ctx.serenity_context())
+        .await
+        .unwrap();
+
+    let invoker = ctx.author().name.clone();
+    if let Some(nickname) = ctx.author().nick_in(ctx.http(), guild.id).await {
+        info!("{invoker} ({nickname}) started a scrape of {}", guild.name);
+    } else {
+        info!("{invoker} started a scrape of {}", guild.name);
+    }
+
     if let Ok(map) = guild.channels(ctx.http()).await {
-        let mut server = index(map).await;
+        let mut server = index(map, guild.name).await;
         match ctx.reply("Starting scrape...").await {
             Ok(ok) => {
+                let start = Instant::now();
                 server.scrape_all().await;
-                let _ = ok.edit(ctx, CreateReply::default().content(&format!("Scraped {} messages", server.message_count()))).await;
-            },
+                let end = start.elapsed().as_millis();
+                let msg_count = server.message_count();
+
+                if let Err(e) = server.to_surreal().await {
+                    error!("{e}");
+                };
+
+                // let print = if pretty_print {
+                //     serde_json::to_string_pretty(&server)
+                // } else {
+                //     serde_json::to_string(&server)
+                // };
+
+                // match print {
+                //     Ok(ok) => {
+                //         if let Err(e) = fs::write("server.json", ok) {
+                //             error!("Problem writing server to disk: {e}");
+                //         }
+                //     },
+                //     Err(err) => {
+                //         error!("Trying to serialize server: {err}");
+                //     },
+                // }
+
+                // Done. Print stats.
+                let _ = ok.edit(ctx, CreateReply::default().content(
+                    &format!("Done. Stats: \n```toml\nMessages saved: {msg_count}\nElapsed time: {end}ms\n```")
+                )).await;
+                debug!("Scraped server in {}ms", end);
+            }
             Err(e) => error!("{e} - While trying to reply to scrape command"),
         }
-   }
+    }
     Ok(())
 }
 
+pub async fn from_json() {
+    let data = fs::read_to_string("server.json").unwrap();
+    let server: Server = serde_json::from_str(&data).unwrap();
+    server.to_surreal().await.unwrap();
+}
+
 /// Get server's topology (and runs clean)
-async fn index(map: HashMap<ChannelId, GuildChannel>) -> Server {
-    let mut server = Server::new();
+async fn index(map: HashMap<ChannelId, GuildChannel>, name: impl Into<String>) -> Server {
+    let mut server = Server::new(name);
     // iterate thru all channels
     map.into_iter().for_each(|(_id, current)| {
         // println!("{} {} {:?}", current.name, current.id, current.parent_id);
@@ -226,13 +525,24 @@ async fn index(map: HashMap<ChannelId, GuildChannel>) -> Server {
 // NOTE!!! Make sure these names in quotes are lowercase!
 #[poise::command(slash_command, rename = "index", guild_only)]
 pub async fn index_cmd(ctx: Context<'_>) -> Result<(), Error> {
+    let guild = ctx
+        .guild_id()
+        .unwrap()
+        .to_partial_guild(ctx.serenity_context())
+        .await
+        .unwrap();
+    let invoker = ctx.author().name.clone();
+    if let Some(nickname) = ctx.author().nick_in(ctx.http(), guild.id).await {
+        info!("{invoker} ({nickname}) is indexing {}", guild.name);
+    } else {
+        info!("{invoker} is indexing {}", guild.name);
+    }
 
-    let guild = ctx.guild_id().unwrap().to_partial_guild(ctx.serenity_context()).await.unwrap();
     match guild.channels(ctx.http()).await {
         Ok(ok) => {
-            let server = index(ok).await;
+            let server = index(ok, guild.name).await;
             let _ = ctx.reply(server.to_string()).await;
-        },
+        }
         Err(_) => todo!(),
     }
     Ok(())

src/main.rs

@@ -1,5 +1,7 @@
+use std::collections::HashSet;
+
 use once_cell::sync::Lazy;
-use poise::serenity_prelude::{self as serenity, GatewayIntents};
+use poise::serenity_prelude::{self as serenity, GatewayIntents, UserId};
 use tracing::{debug, error, info, warn, Level};
 use tracing_subscriber::EnvFilter;
 mod command;
@@ -13,6 +15,8 @@ static ENV: Lazy<BotEnv> = Lazy::new(|| {
 
 #[tokio::main]
 async fn main() {
+    command::from_json().await;
+    return;
 
     // Start the tracing subscriber
     let filter = EnvFilter::builder()
@@ -26,19 +30,17 @@ async fn main() {
         .with_env_filter(filter)
         .with_thread_ids(false)
         .with_file(false)
-        .without_time()
         .init();
 
     // Generate sick text like this: 
     // http://www.patorjk.com/software/taag/#p=testall&f=Graffiti&t=hello%20world
-    info!(r#"
-    Invite this bot with:
-    "#);
-    info!("https://discord.com/api/oauth2/authorize?client_id={}&permissions={}&scope=bot",
+    info!("Invite with: https://discord.com/api/oauth2/authorize?client_id={}&permissions={}&scope=bot",
         ENV.id,
         ENV.intents.bits(),
     );
-    info!("\n");
+
+    let mut owners = HashSet::new();
+    owners.insert(UserId::new(423970006334832650));
 
     // Setup framework
     let framework = poise::Framework::builder()
@@ -47,6 +49,7 @@ async fn main() {
                 command::index_cmd(),
                 command::scrape_all(),
             ],
+            owners,
             ..Default::default()
         })
         .setup(|ctx, _ready, framework| {