use tokio; use mysql_async::{Pool, params}; use mysql_async::prelude::*; use chrono::{Utc, Duration}; use std::env::var; use clap::{App, Arg}; const MAX_DAYS_HELP: &'static str = "Purge messages older than days old If the value is 0 then messages are not deleted by their age "; const MSG_LIMIT_HELP: &'static str = "Sets the max number of messages per channel. NOTE: Deletes based on time so oldest messages are removed first until the max number of messages has been reached. Example: #channel-a contains 1500 messages. Limit is set to 1000. Oldest 500 messages are removed "; async fn remove_old(pool: &Pool, age_limit: i64) -> Result<(), mysql_async::Error> { // really stupid guard but whatever if age_limit == 0 { Ok(()) } else { // Straight removal of messages by age let now = Utc::now(); let oldest = (now - Duration::days(age_limit)).timestamp(); let query = "DELETE FROM messages WHERE time < :age"; let p = params!{ "age" => oldest }; let mut conn = pool.get_conn().await?; conn.exec_drop(query, p).await?; Ok(()) } } async fn remove_maxed(pool: &Pool, max_msg: u64) -> Result<(), mysql_async::Error> { // For N channels we're going to do n+1 queries making this really expensive // TODO: reduce this to a single query for reasons below: tags[network, sql, slow] /* * Multiple network hits blow(moreso if the the db is located away from the * main server binary but i digress. * The first hit is no the expensive part if I'm being honest, * It's everything after that leverages the channel id's as the messages * table is likely to be filled with a ton of shit. * * The trick is going to be in some really nasty group by/limit thing but * frankly this works for now so into the backlog it goes */ let mut conn = pool.get_conn().await?; let channel_ids: Vec = conn.exec_map( "SELECT id FROM channels", (), | row | { row } ).await?; // reverse the order and apply the limit to get the N newest items for cid in channel_ids { let query = "DELETE msg FROM messages msg join (SELECT id FROM messages imsg WHERE imsg.channel_id = :cid ORDER BY id DESC LIMIT 999 OFFSET :lim ) imsg on imsg.id = msg.id;"; let param = params!{"cid" => cid, "lim" => max_msg}; if let Err(e) = conn.exec_drop(query , param).await { eprintln!("Couldn't remove messages from channel: {}", cid); eprintln!("Mysql error: {}", e); break; } } Ok(()) } #[tokio::main] async fn main() -> Result<(), String>{ let args = App::new("chan-chron") .version("1.0") .author("shockrah - with <3") .about("Removes \"old\" messages periodically") .arg(Arg::with_name("max_days") .short("d") .long("max-days") .value_name("D") .takes_value(true) .long_help(MAX_DAYS_HELP)) .arg(Arg::with_name("config_file") .short("f") .long("file") .required(true) .takes_value(true)) .arg(Arg::with_name("msg_limit") .short("l") .long("msg-limit") .required(false) .takes_value(true) .help(MSG_LIMIT_HELP)) .get_matches(); //if let Some(db_url) = args.value_of("db-url") { let db_url = if let Some(config) = args.value_of("config_file") { match dotenv::from_filename(config) { // Not bothering with fucked up configs thats an admin problem not mine Ok(_) => var("DATABASE_URL").unwrap(), _ => panic!("Config not found!") } } else { panic!("No config file provided!"); }; // could fail but again, thats a problem with configs sucking let pool = Pool::from_url(db_url).unwrap(); let mut exec = 0; if let Some(max_days) = args.value_of("max_days") { let max: i64 = match max_days.parse() { Ok(val) => val, _ => panic!("Couldn't parse max_days value") }; // Panicing in hopes that we avoid running later on if let Err(e) = remove_old(&pool, max).await { panic!("{}", e); } else { exec = 1; } } if let Some(limit) = args.value_of("msg_limit") { let limit: u64 = match limit.parse() { Ok(val) => val, _ => panic!("couldn't parse msg_limit value") }; if let Err(e) = remove_maxed(&pool, limit).await { panic!("{}", e); } else { exec |= 2; } } if exec == 0 { println!("Nothing todo"); } else { match exec { 1 => println!("Age Based Purge succeed"), 2 => println!("Max message based purge succeed"), 3 => println!("All succeed"), _ => {} }; } Ok(()) }