Compare commits

...

4 Commits

Author SHA1 Message Date
66e72aeb54 surreal testing
All checks were successful
Test Rust project / test (ubuntu-latest, stable) (push) Successful in 4m29s
2024-07-29 16:08:04 -06:00
a701716958 surreal v1 2024-07-29 15:00:38 -06:00
8e53f2a234 add deserializing 2024-07-29 15:00:10 -06:00
9492f378d3 re-arrange 2024-07-29 13:38:51 -06:00
4 changed files with 2353 additions and 80 deletions

2051
Cargo.lock generated

File diff suppressed because it is too large Load Diff

View File

@ -18,3 +18,4 @@ tracing = "0.1.40"
tracing-subscriber = { version = "0.3.18", features = ["env-filter"] } tracing-subscriber = { version = "0.3.18", features = ["env-filter"] }
serde_json = "1.0.120" serde_json = "1.0.120"
serde = { version = "1.0.204", features = ["derive"] } serde = { version = "1.0.204", features = ["derive"] }
surrealdb = "1.5.4"

View File

@ -1,27 +1,26 @@
use std::{collections::HashMap, fmt::Display, fs, sync::Arc}; use std::{collections::HashMap, fmt::Display, fs, hint::black_box, sync::Arc};
use crate::Context; use crate::Context;
use anyhow::Error; use anyhow::Error;
use poise::{serenity_prelude::{Cache, CacheHttp, ChannelId, ChannelType, GetMessages, GuildChannel, Http, Message}, CreateReply}; use poise::{
use serde::Serialize; 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 tokio::time::Instant;
use tracing::{debug, error, info, trace}; use tracing::{debug, error, info, trace};
#[derive(Serialize)] #[derive(Serialize, Deserialize)]
struct Server {
channels: Vec<Channel>,
#[serde(skip_serializing_if = "Vec::is_empty")]
orphanage: Vec<GuildChannel>,
#[serde(skip_serializing)]
needs_clean: bool,
}
#[derive(Serialize)]
struct Channel { struct Channel {
this: GuildChannel, this: GuildChannel,
#[serde(skip_serializing_if = "Vec::is_empty")] #[serde(skip_serializing_if = "Vec::is_empty")]
#[serde(default)]
children: Vec<Channel>, children: Vec<Channel>,
#[serde(skip_serializing_if = "Vec::is_empty")] #[serde(skip_serializing_if = "Vec::is_empty")]
#[serde(default)]
messages: Vec<Message>, messages: Vec<Message>,
} }
@ -34,16 +33,34 @@ impl Channel {
messages: Vec::new(), 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 { impl Display for Server {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
fn print(f: &mut std::fmt::Formatter<'_>, tab: usize, channel: &Vec<Channel>) { fn print(f: &mut std::fmt::Formatter<'_>, tab: usize, channel: &Vec<Channel>) {
for i in channel { for i in channel {
for _ in 0..tab { let _ = write!(f, "\t"); } for _ in 0..tab {
let _ = writeln!(f, "{} {}", prefix(i.this.kind),i.this.name); let _ = write!(f, "\t");
print(f, tab+1, &i.children); }
let _ = writeln!(f, "{} {}", prefix(i.this.kind), i.this.name);
print(f, tab + 1, &i.children);
} }
} }
@ -65,13 +82,12 @@ impl Display for Server {
} }
} }
print(f, 0, &self.channels); print(f, 0, &self.channels);
if self.needs_clean { if self.needs_clean {
let _ = writeln!(f, "Orphans: (please clean() before displaying...)"); let _ = writeln!(f, "Orphans: (please clean() before displaying...)");
for i in &self.orphanage { for i in &self.orphanage {
let _ = write!(f, "{} {},", prefix(i.kind),i.name); let _ = write!(f, "{} {},", prefix(i.kind), i.name);
} }
} }
@ -80,15 +96,15 @@ impl Display for Server {
} }
impl Server { impl Server {
fn new(name: impl Into<String>) -> Self {
fn new() -> Self {
Self { Self {
name: name.into(),
channels: Vec::new(), channels: Vec::new(),
orphanage: Vec::new(), orphanage: Vec::new(),
needs_clean: false needs_clean: false,
} }
} }
fn search_by_id<'a>(target: &'a mut Vec<Channel>, find: &ChannelId) -> Option<&'a mut Channel> { fn search_by_id<'a>(target: &'a mut Vec<Channel>, find: &ChannelId) -> Option<&'a mut Channel> {
for child in target { for child in target {
if child.this.id == *find { if child.this.id == *find {
@ -97,9 +113,9 @@ impl Server {
if let Some(x) = Self::search_by_id(&mut child.children, find) { if let Some(x) = Self::search_by_id(&mut child.children, find) {
return Some(x); return Some(x);
} }
} }
None None
} }
fn add(&mut self, insert: GuildChannel) { fn add(&mut self, insert: GuildChannel) {
// make sure the new item wants a parent // make sure the new item wants a parent
@ -109,12 +125,12 @@ impl Server {
match Self::search_by_id(&mut self.channels, parent_id) { match Self::search_by_id(&mut self.channels, parent_id) {
Some(parent_node) => { Some(parent_node) => {
parent_node.children.push(Channel::new(insert)); parent_node.children.push(Channel::new(insert));
}, }
None => { None => {
// couldn't find parent, store somewhere else until it's parent is added... // couldn't find parent, store somewhere else until it's parent is added...
self.orphanage.push(insert); self.orphanage.push(insert);
self.needs_clean = true; self.needs_clean = true;
}, }
} }
} else { } else {
self.channels.push(Channel::new(insert)); self.channels.push(Channel::new(insert));
@ -123,7 +139,9 @@ impl Server {
/// Cleans out the orphan channels, finding them parents. You'll want to use this before displaying anything. /// Cleans out the orphan channels, finding them parents. You'll want to use this before displaying anything.
fn clean(&mut self) { fn clean(&mut self) {
if !self.needs_clean {return;} if !self.needs_clean {
return;
}
// Look thru the orphanage and try to find parents // Look thru the orphanage and try to find parents
for orphan in &self.orphanage { for orphan in &self.orphanage {
@ -157,20 +175,28 @@ impl Server {
// Clone *should* be cheap - it's Arc under the hood // Clone *should* be cheap - it's Arc under the hood
// get the messages // get the messages
get_messages(channel, cache.clone(), settings).await; get_messages(channel, cache.clone(), settings).await;
} }
/// Get all messages for 1 channel and children /// 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. // Loop thru all the messages in the channel in batches.
// Adding each batch to the current channel's messages the whole time. // Adding each batch to the current channel's messages the whole time.
let mut last_id = channel.this.last_message_id; let mut last_id = channel.this.last_message_id;
while let Some(last) = last_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) => { Ok(mut ok) => {
if ok.is_empty() { if ok.is_empty() {
debug!("Reached the beginning of {}", channel.this.name); debug!("Reached the beginning of {}", channel.this.name);
// Stop the loop if there are no more messages. // Stop the loop if there are no more messages.
last_id = None; last_id = None;
} else { } else {
trace!("Adding {} messages to \"{}\"", ok.len(), channel.this.name); trace!("Adding {} messages to \"{}\"", ok.len(), channel.this.name);
channel.messages.append(&mut ok); channel.messages.append(&mut ok);
@ -178,13 +204,16 @@ impl Server {
if let Some(l) = channel.messages.last() { if let Some(l) = channel.messages.last() {
last_id = Some(l.id); last_id = Some(l.id);
} }
} }
}, }
Err(e) => { 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. // Stop reading this channel on an error.
last_id = None; last_id = None;
}, }
} }
} }
// Then recurse into children channels // Then recurse into children channels
@ -201,17 +230,228 @@ impl Server {
for channel in this { for channel in this {
total += channel.messages.len(); total += channel.messages.len();
total += walk(&channel.children); total += walk(&channel.children);
}; }
total total
} }
walk(&self.channels) 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, owners_only)] #[poise::command(slash_command, rename = "scrape_all", guild_only, owners_only)]
pub async fn scrape_all(ctx: Context<'_>, pretty_print: bool) -> Result<(), Error> { 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(); let invoker = ctx.author().name.clone();
if let Some(nickname) = ctx.author().nick_in(ctx.http(), guild.id).await { if let Some(nickname) = ctx.author().nick_in(ctx.http(), guild.id).await {
@ -221,7 +461,7 @@ pub async fn scrape_all(ctx: Context<'_>, pretty_print: bool) -> Result<(), Erro
} }
if let Ok(map) = guild.channels(ctx.http()).await { 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 { match ctx.reply("Starting scrape...").await {
Ok(ok) => { Ok(ok) => {
let start = Instant::now(); let start = Instant::now();
@ -229,38 +469,48 @@ pub async fn scrape_all(ctx: Context<'_>, pretty_print: bool) -> Result<(), Erro
let end = start.elapsed().as_millis(); let end = start.elapsed().as_millis();
let msg_count = server.message_count(); let msg_count = server.message_count();
let print = if pretty_print { if let Err(e) = server.to_surreal().await {
serde_json::to_string_pretty(&server) error!("{e}");
} else {
serde_json::to_string(&server)
}; };
match print { // let print = if pretty_print {
Ok(ok) => { // serde_json::to_string_pretty(&server)
if let Err(e) = fs::write("server.json", ok) { // } else {
error!("Problem writing server to disk: {e}"); // serde_json::to_string(&server)
} // };
},
Err(err) => { // match print {
error!("Trying to serialize server: {err}"); // 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. // Done. Print stats.
let _ = ok.edit(ctx, CreateReply::default().content( let _ = ok.edit(ctx, CreateReply::default().content(
&format!("Done. Stats: \n```toml\nMessages saved: {msg_count}\nElapsed time: {end}ms\n```") &format!("Done. Stats: \n```toml\nMessages saved: {msg_count}\nElapsed time: {end}ms\n```")
)).await; )).await;
debug!("Scraped server in {}ms", end); debug!("Scraped server in {}ms", end);
}, }
Err(e) => error!("{e} - While trying to reply to scrape command"), Err(e) => error!("{e} - While trying to reply to scrape command"),
} }
} }
Ok(()) 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) /// Get server's topology (and runs clean)
async fn index(map: HashMap<ChannelId, GuildChannel>) -> Server { async fn index(map: HashMap<ChannelId, GuildChannel>, name: impl Into<String>) -> Server {
let mut server = Server::new(); let mut server = Server::new(name);
// iterate thru all channels // iterate thru all channels
map.into_iter().for_each(|(_id, current)| { map.into_iter().for_each(|(_id, current)| {
// println!("{} {} {:?}", current.name, current.id, current.parent_id); // println!("{} {} {:?}", current.name, current.id, current.parent_id);
@ -275,7 +525,12 @@ async fn index(map: HashMap<ChannelId, GuildChannel>) -> Server {
// NOTE!!! Make sure these names in quotes are lowercase! // NOTE!!! Make sure these names in quotes are lowercase!
#[poise::command(slash_command, rename = "index", guild_only)] #[poise::command(slash_command, rename = "index", guild_only)]
pub async fn index_cmd(ctx: Context<'_>) -> Result<(), Error> { 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(); let invoker = ctx.author().name.clone();
if let Some(nickname) = ctx.author().nick_in(ctx.http(), guild.id).await { if let Some(nickname) = ctx.author().nick_in(ctx.http(), guild.id).await {
info!("{invoker} ({nickname}) is indexing {}", guild.name); info!("{invoker} ({nickname}) is indexing {}", guild.name);
@ -285,9 +540,9 @@ pub async fn index_cmd(ctx: Context<'_>) -> Result<(), Error> {
match guild.channels(ctx.http()).await { match guild.channels(ctx.http()).await {
Ok(ok) => { Ok(ok) => {
let server = index(ok).await; let server = index(ok, guild.name).await;
let _ = ctx.reply(server.to_string()).await; let _ = ctx.reply(server.to_string()).await;
}, }
Err(_) => todo!(), Err(_) => todo!(),
} }
Ok(()) Ok(())

View File

@ -15,6 +15,8 @@ static ENV: Lazy<BotEnv> = Lazy::new(|| {
#[tokio::main] #[tokio::main]
async fn main() { async fn main() {
command::from_json().await;
return;
// Start the tracing subscriber // Start the tracing subscriber
let filter = EnvFilter::builder() let filter = EnvFilter::builder()