FrostByte/server/src/state.rs

use std::collections::BTreeMap;
use std::sync::Arc;
use std::sync::Mutex;

use log::error;
use log::info;
use sqlx::postgres::PgPoolOptions;
use sqlx::PgPool;

#[derive(Clone)]
pub struct CaptchaState {
    pub capthca_db: Arc<Mutex<BTreeMap<i32, String>>>,
}

impl CaptchaState {
    pub fn new() -> Self {
        Self {
            capthca_db: Arc::new(Mutex::new(BTreeMap::new())),
        }
    }
}

#[derive(Clone)]
pub struct ServerState {
    pub pool: PgPool,
}

impl ServerState {
    pub async fn new() -> Self {
        // This is almost certainly bad practice for more reasons than I can count
        dotenvy::dotenv().ok();
        let db_url = dotenvy::var("DATABASE_URL").unwrap_or_else(|_| {
            error!("DATABASE_URL not set in environment!");
            std::process::exit(1);
        });

        info!("Using db_url: {}", &db_url);

        let pool = PgPoolOptions::new()
            .max_connections(5)
            .connect(&db_url)
            .await
            .unwrap();

        sqlx::migrate!("./migrations_pg").run(&pool).await.unwrap();

        match crate::db::db_new_user("imbus".to_string(), "kartellen1234".to_string(), &pool).await
        {
            Some(u) => info!("Created default user {}", u.username),
            None => error!("Failed to create default user..."),
        }

        #[cfg(debug_assertions)]
        debug_setup(&pool).await.unwrap();

        Self { pool }
    }
}

// Inserts a bunch of dummy data into the database
// Mostly useful for debugging new posts, as we need to satisfy foreign key constraints.
#[cfg(debug_assertions)]
async fn debug_setup(pool: &PgPool) -> Result<(), sqlx::Error> {
    use lipsum::lipsum;
    use sqlx::query;

    use crate::db::db_new_user;

    db_new_user("user".to_string(), "pass".to_string(), pool).await;

    // Check if the demo post already exists
    let posted = query!("SELECT * FROM posts WHERE id = 1",)
        .fetch_one(pool)
        .await
        .ok();

    // If the demo user already has a post, don't insert another one
    if !posted.is_some() {
        // This requires that the user with id 1 exists in the user table
        query!("INSERT INTO posts (user_id, content) VALUES (1, 'Hello world! The demo username is user and the password is pass.')",)
        .execute(pool)
        .await?;

        for _ in 0..10 {
            query!(
                "INSERT INTO posts (user_id, content) VALUES (1, $1)",
                lipsum(50)
            )
            .execute(pool)
            .await?;
        }
    }

    Ok(())
}