use std::collections::HashSet; use std::{fs::File, io::Read, collections::HashMap, path::Path}; use poise::serenity_prelude::{SerenityError, Http}; use r2d2::{Pool, PooledConnection}; use r2d2_sqlite::SqliteConnectionManager; use r2d2_sqlite::rusqlite::{self, params}; use derive_more::From; use crate::{utils::get_challenge_number, models::User}; use super::{LegacySubmission, Submission}; pub struct Database { // Must be Arc because Connection contains RefCell, // which cannot be shared between threads safely connection_pool: Pool, } const DATABASE_FILENAME: &str = "database.db"; #[derive(From, Debug)] pub enum DatabaseError { Rusqlite(rusqlite::Error), Pool(r2d2::Error), } #[derive(From, Debug)] pub enum LoadLegacyError { Database(DatabaseError), Serenity(SerenityError), } type Result = std::result::Result; impl Database { pub fn file_exists() -> bool { Path::new(DATABASE_FILENAME).exists() } pub fn new( testing: bool, ) -> Result { let connection_manager = if testing { SqliteConnectionManager::memory() } else { SqliteConnectionManager::file(DATABASE_FILENAME) }; let connection_pool = Pool::new(connection_manager)?; let conn = connection_pool.get()?; conn.execute( "CREATE TABLE IF NOT EXISTS Submission ( id INTEGER PRIMARY KEY, author_id INTEGER NOT NULL, timestamp DATETIME DEFAULT CURRENT_TIMESTAMP, image TEXT NOT NULL, challenge INTEGER NOT NULL, FOREIGN KEY (author_id) REFERENCES User(id) )", params![], )?; conn.execute( "CREATE TABLE IF NOT EXISTS User ( id INTEGER PRIMARY KEY, name TEXT NOT NULL, discriminator INTEGER NOT NULL, avatar TEXT, deleted INTEGER NOT NULL )", params![], )?; Ok(Self { connection_pool }) } fn conn(&self) -> std::result::Result, r2d2::Error> { self.connection_pool.get() } pub fn has_submitted(&self, user_id: u64) -> Result { Ok(self.conn()? .prepare("SELECT 1 FROM User WHERE id = ?1 LIMIT 1")? .query_row(params![user_id], |_row| Ok(())) .is_ok()) } pub async fn load_legacy(&self, http: &Http) -> std::result::Result<(), LoadLegacyError> { let latest_challenge = get_challenge_number(); // HashMap of archived users that are no longer sharing a server with 字ちゃん // Their historical usernames and discriminators will be used let mut archived_users = HashMap::new(); let conn = self.conn().map_err(DatabaseError::Pool)?; for n in 1..=latest_challenge { println!("Loading legacy challenge {n}/{latest_challenge}..."); let mut file = File::open(format!("data/challenges/{n}.json")).unwrap(); let mut contents = String::new(); file.read_to_string(&mut contents).unwrap(); for legacy in serde_json::from_str::>(&contents).unwrap() { let id = legacy.id; if !self.has_submitted(id)? { println!("Fetching user {id}..."); // If it already contains the archived user, // overwrite write their data since they may have updated // their username/discriminator since their previous submission match archived_users.get(&id) { Some(User { deleted, .. }) => { archived_users.insert(id, User { id, avatar: None, deleted: *deleted, ..User::from_username(&legacy.username) }); }, None => match User::fetch(http, id).await { Ok(User { deleted: true, .. }) => { archived_users.insert(id, User { id, avatar: None, deleted: true, ..User::from_username(&legacy.username) }); }, Ok(user) => { conn.execute( "INSERT INTO User(id, name, discriminator, avatar, deleted) VALUES (?1, ?2, ?3, ?4, ?5)", params![user.id, user.name, user.discriminator, user.avatar, user.deleted] ).map_err(DatabaseError::Rusqlite)?; }, Err(error) if error.to_string().eq("Unknown User") => { // This will also be called in the case of an invalid user ID println!("Failed to fetch user {id}, adding to archive"); archived_users.insert(id, User { id, avatar: None, deleted: false, ..User::from_username(&legacy.username) }); }, Err(error) => return Err(LoadLegacyError::Serenity(error)), }, }; } for submission in legacy.parse(n as u32).unwrap() { conn.execute( "INSERT INTO Submission(author_id, timestamp, image, challenge) VALUES (?1, ?2, ?3, ?4)", params![ &submission.author_id, &submission.timestamp, &submission.image, n, ] ).map_err(DatabaseError::Rusqlite)?; } } } for (_id, user) in archived_users { conn.execute( "INSERT INTO USER (id, name, discriminator, avatar, deleted) VALUES (?1, ?2, ?3, ?4, ?5)", params![user.id, user.name, user.discriminator, user.avatar, user.deleted] ).map_err(DatabaseError::Rusqlite)?; } Ok(()) } pub fn get_challenge_user_data(&self, challenge: u32) -> Result<(Vec, HashMap)> { let submissions = self.get_submissions(challenge)?; let users = self.get_users({ let mut user_ids = HashSet::new(); for submission in &submissions { user_ids.insert(submission.author_id); } user_ids })?; Ok((submissions, users)) } pub fn get_submissions(&self, challenge: u32) -> Result> { Ok(self.conn()? .prepare("SELECT author_id, timestamp, image FROM Submission WHERE challenge = ?1")? .query_map(params![challenge], |row| { Ok(Submission { author_id: row.get(0)?, timestamp: row.get(1)?, image: row.get(2)?, challenge, }) })? .collect::, rusqlite::Error>>()?) } pub fn get_user_by_name(&self, name: &str) -> Result> { Ok(self.conn()? .prepare("SELECT id, discriminator, avatar, deleted FROM User where name = ?1")? .query_row(params![name], |row| { Ok(Some(User { id: row.get(0)?, name: name.to_owned(), discriminator: row.get(1)?, avatar: row.get(2)?, deleted: row.get(3)?, })) })?) } pub fn get_submissions_by_user_name(&self, name: &str) -> Result> { Ok(self.conn()? .prepare(" SELECT author_id, timestamp, image, challenge FROM Submission s JOIN User u ON s.author_id = u.id WHERE u.name = ?1", )? .query_map(params![name], |row| { Ok(Submission { author_id: row.get(0)?, timestamp: row.get(1)?, image: row.get(2)?, challenge: row.get(3)?, }) })? .collect::, rusqlite::Error>>()?) } fn get_users(&self, users: HashSet) -> Result> { let conn = self.conn()?; // Not sure why derive_more::From is unable to convert these errors users .iter() // u64 must be converted to String for templates .map(|id| -> Result<(String, User)> { match conn.prepare("SELECT name, discriminator, avatar, deleted FROM User WHERE id = ?1") { Ok(mut statement) => statement.query_row(params![id], |row| { Ok((id.to_string(), User { id: *id, name: row.get(0)?, discriminator: row.get(1)?, avatar: row.get(2)?, deleted: row.get(3)?, })) }).map_err(DatabaseError::Rusqlite), Err(error) => Err(DatabaseError::Rusqlite(error)), } }) .collect() } #[allow(dead_code)] pub fn refresh_users(&self) -> rusqlite::Result<()> { // Periodically refresh all changable user data (name, discriminator, avatar) // Ideally this should run periodically. todo!() } #[allow(dead_code, unused_variables)] pub fn insert_submission(&self, submission: &Submission) -> rusqlite::Result<()> { // For new submissions only todo!() } }