From 66e72aeb54c1aa200e15b42ffab7b4c4719a0e1d Mon Sep 17 00:00:00 2001 From: oliver Date: Mon, 29 Jul 2024 16:08:04 -0600 Subject: [PATCH] surreal testing --- src/command.rs | 246 ++++++++++++++++++++++++++++++++++++++++--------- src/main.rs | 3 +- 2 files changed, 202 insertions(+), 47 deletions(-) diff --git a/src/command.rs b/src/command.rs index 8a0925a..9df1ef0 100644 --- a/src/command.rs +++ b/src/command.rs @@ -2,7 +2,12 @@ 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 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; @@ -28,6 +33,10 @@ impl Channel { messages: Vec::new(), } } + + fn is_category(&self) -> bool { + !self.children.is_empty() && self.messages.is_empty() + } } #[derive(Serialize, Deserialize)] @@ -45,12 +54,13 @@ struct Server { 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) { 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); } } @@ -72,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); } } @@ -87,16 +96,15 @@ impl Display for Server { } impl Server { - fn new(name: impl Into) -> 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, find: &ChannelId) -> Option<&'a mut Channel> { for child in target { if child.this.id == *find { @@ -105,9 +113,9 @@ impl Server { 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 @@ -117,12 +125,12 @@ impl Server { 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)); @@ -131,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 { @@ -165,20 +175,28 @@ impl Server { // 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); @@ -186,13 +204,16 @@ impl Server { if let Some(l) = channel.messages.last() { last_id = Some(l.id); } - } - }, + } + } Err(e) => { - error!("While reading messages in \"{}\" before `{}` - {e}", channel.this.name, last); + 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 @@ -209,7 +230,7 @@ impl Server { for channel in this { total += channel.messages.len(); total += walk(&channel.children); - }; + } total } walk(&self.channels) @@ -222,21 +243,140 @@ impl Server { db.signin(Root { username: "root", password: "root", - }).await?; + }) + .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, + topic: String, + } + let chan = &cat.this; + let dpi = if let Some(val) = chan.parent_id {Some(val.get())} else { None }; + let new_channel: Vec = 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, + attachments: Vec, + reactions: Vec, + pinned: bool, + } + + let _: Vec = 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 = db - .create("category") - .content(&cat.this) - .await?; + let new_category: Vec = 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 @@ -244,7 +384,11 @@ impl Server { import_channel(&cat.children, &new_category[0], &db).await?; } - async fn import_channel(channels: &Vec, parent: &Thing, db: &Surreal) -> surrealdb::Result<()> { + async fn import_channel( + channels: &Vec, + parent: &Thing, + db: &Surreal, + ) -> surrealdb::Result<()> { for channel in channels.iter() { trace!("Importing channel \"{}\"", channel.this.name); #[derive(Serialize)] @@ -267,14 +411,17 @@ impl Server { Ok(()) } - async fn import_messages(msgs: &Vec, parent: &Thing, db: &Surreal) -> surrealdb::Result<()> { + async fn import_messages( + msgs: &Vec, + parent: &Thing, + db: &Surreal, + ) -> 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 + surreal_parent: &'b Thing, } let created: Vec = db @@ -283,7 +430,7 @@ impl Server { message: &msg, surreal_parent: &parent, }) - .await?; + .await?; trace!("Imported message {:?}", created); } @@ -295,12 +442,16 @@ impl Server { Ok(()) } - } #[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 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 { @@ -344,18 +495,17 @@ pub async fn scrape_all(ctx: Context<'_>, pretty_print: bool) -> Result<(), Erro &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 fn from_json() { +pub async fn from_json() { let data = fs::read_to_string("server.json").unwrap(); let server: Server = serde_json::from_str(&data).unwrap(); - - black_box(server); + server.to_surreal().await.unwrap(); } /// Get server's topology (and runs clean) @@ -375,7 +525,12 @@ async fn index(map: HashMap, name: impl Into) - // 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 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); @@ -387,9 +542,8 @@ pub async fn index_cmd(ctx: Context<'_>) -> Result<(), Error> { Ok(ok) => { let server = index(ok, guild.name).await; let _ = ctx.reply(server.to_string()).await; - }, + } Err(_) => todo!(), } Ok(()) } - diff --git a/src/main.rs b/src/main.rs index c3f4109..cacad10 100644 --- a/src/main.rs +++ b/src/main.rs @@ -15,7 +15,8 @@ static ENV: Lazy = Lazy::new(|| { #[tokio::main] async fn main() { - command::from_json(); + command::from_json().await; + return; // Start the tracing subscriber let filter = EnvFilter::builder()