reminder-bot/src/main.rs

#[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<dyn std::error::Error + Send + Sync>;

#[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<Database>,
    #[cfg(not(test))]
    recording_macros: RwLock<HashMap<(GuildId, UserId), CommandMacro>>,
    popular_timezones: Vec<Tz>,
    _broadcast: Sender<()>,
}

impl Debug for Data {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(f, "Data {{ .. }}")
    }
}

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<dyn StdError + Send + Sync>> {
    let (tx, mut rx) = broadcast::channel(16);

    tokio::select! {
        output = _main(tx) => output,
        _ = rx.recv() => Err(Box::new(Ended) as Box<dyn StdError + Send + Sync>)
    }
}

#[cfg(not(test))]
async fn _main(tx: Sender<()>) -> Result<(), Box<dyn StdError + Send + Sync>> {
    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::<Vec<_>>();
    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")?;
            if sent_clean_age.is_empty() {
                panic!("No SENT_CLEAN_AGE")
            }
            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::<Tz>().unwrap())
            .collect::<Vec<Tz>>();

            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(())
        }
    }
}