use std::collections::HashSet; use std::{fs::File, io::Read, collections::HashMap, path::Path}; use poise::serenity_prelude::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), } 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 )", 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) -> Result<()> { 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()?; 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, submissions) in serde_json::from_str::>(&contents) .unwrap() .iter() .map(|legacy| (legacy, legacy.parse().unwrap())) { let mut already_updated = false; for submission in submissions { conn.execute( "INSERT INTO Submission(author_id, timestamp, image, challenge) VALUES (?1, ?2, ?3, ?4)", params![ &submission.author_id, &submission.timestamp, &submission.image, n, ] )?; let id = submission.author_id; if !self.has_submitted(id)? { println!("Fetching user {id}..."); let previously_archived = archived_users.contains_key(&id); // Parse archived user out of legacy and insert into HashMap let mut archive = || { if already_updated { return; } if previously_archived { println!("Updating archived data for user {id}"); } else { println!("Adding archived data for user {id}"); } let (name, discriminator) = { let mut iter = legacy.username.split('#'); let name = iter.next().unwrap().to_owned(); let discriminator = iter .next() .map(|str| str.parse().unwrap()) .unwrap_or(0); (name, discriminator) }; archived_users.insert(id, User { id, name, discriminator, avatar: None, }); already_updated = true; }; if previously_archived { // If it already contains the archived user, // overwrite write their data since they may have updated // their username/discriminator since their previous submission archive(); } else { match User::fetch(http, submission.author_id).await { Ok(user) => { conn.execute( "INSERT INTO User(id, name, discriminator, avatar) VALUES (?1, ?2, ?3, ?4)", params![user.id, user.name, user.discriminator, user.avatar] )?; }, Err(error) => { println!("Failed to fetch user {}, may update archived data: {error}", submission.author_id); archive(); }, }; } } } } } 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)?, }) })? .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 .filter_map(|id| -> Option> { match conn.prepare("SELECT name, discriminator, avatar FROM User WHERE id = ?1") { Ok(mut statement) => Some(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)?, }))).map_err(DatabaseError::Rusqlite)), Err(error) => Some(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!() } }