freechat/chan-like/src/main.rs

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 <D> 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<u64> = 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(())
}