#[macro_use] extern crate lazy_static; mod commands; #[cfg(not(test))] mod component_models; mod consts; #[cfg(not(test))] mod event_handlers; #[cfg(not(test))] mod hooks; mod interval_parser; mod metrics; #[cfg(not(test))] mod models; mod postman; #[cfg(test)] mod test; mod time_parser; mod utils; mod web; use std::{ collections::HashMap, env, error::Error as StdError, fmt::{Debug, Display, Formatter}, path::Path, }; use chrono_tz::Tz; use log::warn; use poise::serenity_prelude::{ model::{ gateway::GatewayIntents, id::{GuildId, UserId}, }, ClientBuilder, }; use serenity::all::ActivityData; use sqlx::{MySql, Pool}; use tokio::sync::{broadcast, broadcast::Sender, RwLock}; use crate::metrics::init_metrics; #[cfg(test)] use crate::test::TestContext; #[cfg(not(test))] use crate::{ commands::{ allowed_dm, clock, clock_context_menu::clock_context_menu, command_macro, dashboard, delete, donate, help, info, look, multiline, nudge, offset, pause, remind, settings, timer, timezone, todo, webhook, }, consts::THEME_COLOR, event_handlers::listener, hooks::all_checks, models::command_macro::CommandMacro, }; type Database = MySql; type Error = Box; #[cfg(test)] type Context<'a> = TestContext<'a>; #[cfg(not(test))] type Context<'a> = poise::Context<'a, Data, Error>; type ApplicationContext<'a> = poise::ApplicationContext<'a, Data, Error>; pub struct Data { database: Pool, #[cfg(not(test))] recording_macros: RwLock>, popular_timezones: Vec, _broadcast: Sender<()>, } struct Ended; impl Debug for Ended { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { f.write_str("Process ended.") } } impl Display for Ended { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { f.write_str("Process ended.") } } impl StdError for Ended {} #[tokio::main(flavor = "multi_thread")] #[cfg(not(test))] async fn main() -> Result<(), Box> { let (tx, mut rx) = broadcast::channel(16); tokio::select! { output = _main(tx) => output, _ = rx.recv() => Err(Box::new(Ended) as Box) } } #[cfg(not(test))] async fn _main(tx: Sender<()>) -> Result<(), Box> { env_logger::init(); if Path::new("/etc/reminder-rs/config.env").exists() { dotenv::from_path("/etc/reminder-rs/config.env")?; } else { let _ = dotenv::dotenv(); } let args = env::args().collect::>(); let cmd_word = args.last().map(|w| w.as_str()); let database = Pool::connect(&env::var("DATABASE_URL").expect("No database URL provided")).await.unwrap(); match cmd_word { Some("clean") => { let sent_clean_age = env::var("SENT_CLEAN_AGE").expect("No SENT_CLEAN_AGE provided"); if sent_clean_age.is_empty() { panic!("SENT_CLEAN_AGE empty") } sqlx::query!( " DELETE FROM reminders WHERE `utc_time` < NOW() - INTERVAL ? DAY AND status != 'pending' ORDER BY `utc_time` LIMIT 1000 ", sent_clean_age ) .execute(&database) .await?; let total_clean_age = env::var("TOTAL_CLEAN_AGE"); if let Ok(total_clean_age) = total_clean_age { sqlx::query!( " DELETE FROM reminders WHERE `utc_time` < NOW() - INTERVAL ? DAY ORDER BY `utc_time` LIMIT 1000 ", total_clean_age ) .execute(&database) .await?; } Ok(()) } _ => { let discord_token = env::var("DISCORD_TOKEN").expect("Missing DISCORD_TOKEN from environment"); let options = poise::FrameworkOptions { commands: vec![ help::command(), info::command(), clock::command(), donate::command(), clock_context_menu(), dashboard::command(), timezone::command(), poise::Command { subcommands: vec![allowed_dm::set::set(), allowed_dm::unset::unset()], ..allowed_dm::allowed_dm() }, poise::Command { subcommands: vec![poise::Command { subcommands: vec![ settings::ephemeral_confirmations::set::set(), settings::ephemeral_confirmations::unset::unset(), ], ..settings::ephemeral_confirmations::ephemeral_confirmations() }], ..settings::settings() }, webhook::command(), poise::Command { subcommands: vec![ command_macro::delete_macro::delete_macro(), command_macro::finish_macro::finish_macro(), command_macro::list_macro::list_macro(), command_macro::record_macro::record_macro(), command_macro::run_macro::run_macro(), ], ..command_macro::command_macro() }, pause::command(), offset::command(), nudge::command(), look::command(), delete::command(), poise::Command { subcommands: vec![ timer::list::list(), timer::start::start(), timer::delete::delete(), ], ..timer::timer() }, multiline::command(), remind::command(), poise::Command { subcommands: vec![ poise::Command { subcommands: vec![ todo::guild::add::add(), todo::guild::view::view(), ], ..todo::guild::guild() }, poise::Command { subcommands: vec![ todo::channel::add::add(), todo::channel::view::view(), ], ..todo::channel::channel() }, poise::Command { subcommands: vec![todo::user::add::add(), todo::user::view::view()], ..todo::user::user() }, ], ..todo::todo() }, ], allowed_mentions: None, command_check: Some(|ctx| Box::pin(all_checks(ctx))), event_handler: |ctx, event, _framework, data| Box::pin(listener(ctx, event, data)), on_error: |error| { Box::pin(async move { match error { poise::FrameworkError::CommandCheckFailed { .. } => { // suppress error } error => { if let Err(e) = poise::builtins::on_error(error).await { log::error!("Error while handling error: {}", e); } } } }) }, ..Default::default() }; // Start metrics init_metrics(); sqlx::migrate!().run(&database).await?; let popular_timezones = sqlx::query!( " SELECT IFNULL(timezone, 'UTC') AS timezone FROM users WHERE timezone IS NOT NULL GROUP BY timezone ORDER BY COUNT(timezone) DESC LIMIT 21 " ) .fetch_all(&database) .await .unwrap() .iter() .map(|t| t.timezone.parse::().unwrap()) .collect::>(); let framework = poise::Framework::builder() .setup(move |ctx, _bot, framework| { Box::pin(async move { poise::builtins::register_globally(ctx, &framework.options().commands) .await?; let kill_tx = tx.clone(); let kill_recv = tx.subscribe(); let ctx1 = ctx.clone(); let ctx2 = ctx.clone(); let pool1 = database.clone(); let pool2 = database.clone(); let run_settings = env::var("DONTRUN").unwrap_or_else(|_| "".to_string()); if !run_settings.contains("postman") { tokio::spawn(async move { match postman::initialize(kill_recv, ctx1, &pool1).await { Ok(_) => {} Err(e) => { panic!("postman exiting: {}", e); } }; }); } else { warn!("Not running postman"); } if !run_settings.contains("web") { tokio::spawn(async move { web::initialize(kill_tx, ctx2, pool2).await.unwrap(); }); } else { warn!("Not running web"); } Ok(Data { database, popular_timezones, recording_macros: Default::default(), _broadcast: tx, }) }) }) .options(options) .build(); let mut client = ClientBuilder::new(&discord_token, GatewayIntents::GUILDS) .framework(framework) .activity(ActivityData::watching("for /remind")) .await?; client.start_autosharded().await?; Ok(()) } } }