use poise::serenity_prelude as sere; use std::env; use tracing::{debug, error, info, trace}; use tracing_subscriber::EnvFilter; pub mod commands; pub mod event_handler; pub mod message_archive; pub mod messaging; pub mod web; pub struct Data { pub pool: sqlx::PgPool, } pub type Error = Box; pub type Context<'a> = poise::Context<'a, Data, Error>; #[tokio::main] async fn main() { dotenv::dotenv().ok(); let log_level = env::var("TG_BOT_LOG").unwrap_or_else(|_| "info".to_string()); let filter = EnvFilter::new(format!("tg_archive_bot={}", log_level)); tracing_subscriber::fmt() .with_env_filter(filter) .with_file(true) .with_thread_ids(true) .with_thread_names(true) .with_line_number(true) .with_target(true) .init(); trace!("Logging system started"); info!("Log Level: {}", log_level); trace!("Env loaded"); trace!("Loading database url"); let database_url = env::var("TG_BOT_DATABASE_URL").expect("Expected a database url in the environment"); trace!("Database url loaded"); trace!("Loading token"); let token = env::var("TG_BOT_DISCORD_TOKEN").expect("Expected a token in the environment"); trace!("Token loaded"); trace!("Loading web address"); let web_address = env::var("TG_BOT_WEB_ADDRESS").unwrap_or_else(|_| "0.0.0.0:3000".to_string()); trace!("Web address loaded"); let discord_client_id = env::var("TG_BOT_DISCORD_CLIENT_ID") .expect("Expected a Discord client ID in the environment"); let discord_client_secret = env::var("TG_BOT_DISCORD_CLIENT_SECRET") .expect("Expected a Discord client secret in the environment"); let discord_redirect_uri = env::var("TG_BOT_DISCORD_REDIRECT_URI") .expect("Expected a Discord redirect URI in the environment"); info!("Starting bot..."); debug!("Initalizing database"); let pool = connect_database(&database_url) .await .expect("Failed to connect to database"); debug!("Database connected"); trace!("Applying migrations"); sqlx::migrate!() .run(&pool) .await .expect("Failed to apply database migrations"); debug!("Migrations applied"); debug!("Starting web server on {}", web_address); let web_listener = tokio::net::TcpListener::bind(&web_address) .await .expect("Failed to bind web server"); tokio::spawn(web::run( web_listener, pool.clone(), token.clone(), discord_client_id, discord_client_secret, discord_redirect_uri, )); debug!("Web server started"); trace!("Loading intents"); let intents = sere::GatewayIntents::non_privileged() | sere::GatewayIntents::MESSAGE_CONTENT | sere::GatewayIntents::GUILD_MEMBERS; trace!("Intents loaded"); trace!("Initalizing framework"); let framework = poise::Framework::builder() .options(poise::FrameworkOptions { event_handler: |ctx, event, framework, user_data| { Box::pin(event_handler::event_handler( ctx, event, framework, user_data, )) }, on_error: |error| { Box::pin(async move { match error { poise::FrameworkError::Command { error, ctx, .. } => { error!("Error in command `{}`: {:?}", ctx.command().name, error); } error => { if let Err(e) = poise::builtins::on_error(error).await { error!("Error while handling error: {}", e); } } } }) }, commands: vec![ commands::ping(), commands::version(), commands::git(), commands::stats(), ], ..Default::default() }) .setup(|ctx, _ready, framework| { Box::pin(async move { if let Ok(guild_id) = env::var("TG_BOT_GUILD_ID") { let guild_id = sere::GuildId::new(guild_id.parse()?); poise::builtins::register_in_guild( ctx, &framework.options().commands, guild_id, ) .await?; } else { poise::builtins::register_globally(ctx, &framework.options().commands).await?; } Ok(Data { pool: pool.clone() }) }) }) .build(); trace!("Framework initialized"); debug!("Connecting to Discord"); let mut client = sere::ClientBuilder::new(&token, intents) .framework(framework) .await .expect("Err creating client"); debug!("Connected to Discord"); info!("Bot running"); if let Err(why) = client.start().await { error!("Client error: {why:?}"); } } async fn connect_database(database_url: &str) -> Result { let options: sqlx::postgres::PgConnectOptions = database_url.parse()?; match sqlx::postgres::PgPoolOptions::new() .connect_with(options.clone()) .await { Ok(pool) => Ok(pool), Err(sqlx::Error::Database(error)) if error.code().as_deref() == Some("3D000") => { let database_name = options.get_database().ok_or_else(|| { std::io::Error::other("Database URL does not contain a database name") })?; let database_name = quote_identifier(database_name); info!("Creating database {}", database_name); let maintenance_pool = sqlx::postgres::PgPoolOptions::new() .connect_with(options.database("postgres")) .await?; let create_database = format!("CREATE DATABASE {database_name};"); sqlx::query(sqlx::AssertSqlSafe(create_database)) .execute(&maintenance_pool) .await?; maintenance_pool.close().await; sqlx::postgres::PgPoolOptions::new() .connect(database_url) .await .map_err(Into::into) } Err(error) => Err(error.into()), } } fn quote_identifier(identifier: &str) -> String { let mut escaped = String::with_capacity(identifier.len()); for character in identifier.chars() { escaped.push(character); if character == '"' { escaped.push(character); } } format!("\"{}\"", escaped) }