Modrinth-plus/packages/app-lib/src/state/db.rs

use crate::state::DirectoryInfo;
use sqlx::sqlite::{
    SqliteConnectOptions, SqliteJournalMode, SqlitePoolOptions,
};
use sqlx::{Pool, Sqlite};
use std::str::FromStr;
use std::time::Duration;

static MIGRATOR: sqlx::migrate::Migrator = sqlx::migrate!();
const HISTORICAL_INIT_MIGRATION_VERSION: i64 = 20240711194701;

pub(crate) async fn connect(
    app_identifier: &str,
) -> crate::Result<Pool<Sqlite>> {
    let settings_dir = DirectoryInfo::initial_settings_dir_path(app_identifier)
        .ok_or(crate::ErrorKind::FSError(
            "Could not find valid config dir".to_string(),
        ))?;

    if !settings_dir.exists() {
        crate::util::io::create_dir_all(&settings_dir).await?;
    }

    let uri = format!("sqlite:{}", settings_dir.join("app.db").display());

    let conn_options = SqliteConnectOptions::from_str(&uri)?
        .busy_timeout(Duration::from_secs(30))
        .journal_mode(SqliteJournalMode::Wal)
        .optimize_on_close(true, None)
        .create_if_missing(true);

    let pool = SqlitePoolOptions::new()
        .max_connections(100)
        .connect_with(conn_options)
        .await?;

    reconcile_historical_migration_checksums(&pool).await?;
    MIGRATOR.run(&pool).await?;

    if let Err(err) = stale_data_cleanup(&pool).await {
        tracing::warn!(
            "Failed to clean up stale data from state database: {err}"
        );
    }

    Ok(pool)
}

async fn reconcile_historical_migration_checksums(
    pool: &Pool<Sqlite>,
) -> crate::Result<()> {
    let Some(init_migration) = MIGRATOR.iter().find(|migration| {
        migration.version == HISTORICAL_INIT_MIGRATION_VERSION
    })
    else {
        return Ok(());
    };

    let has_migration_table: i64 = sqlx::query_scalar(
        "SELECT COUNT(*) FROM sqlite_master WHERE type = 'table' AND name = '_sqlx_migrations'",
    )
    .fetch_one(pool)
    .await?;

    if has_migration_table == 0 {
        return Ok(());
    }

    let stored_checksum: Option<Vec<u8>> = sqlx::query_scalar(
        "SELECT checksum FROM _sqlx_migrations WHERE version = ? AND success = TRUE",
    )
    .bind(HISTORICAL_INIT_MIGRATION_VERSION)
    .fetch_optional(pool)
    .await?;

    if let Some(stored_checksum) = stored_checksum {
        let bundled_checksum: &[u8] = init_migration.checksum.as_ref();

        if stored_checksum.as_slice() != bundled_checksum {
            tracing::warn!(
                "Reconciling checksum for historical Modrinth init migration {}",
                HISTORICAL_INIT_MIGRATION_VERSION
            );

            sqlx::query(
                "UPDATE _sqlx_migrations SET checksum = ? WHERE version = ? AND success = TRUE",
            )
            .bind(bundled_checksum)
            .bind(HISTORICAL_INIT_MIGRATION_VERSION)
            .execute(pool)
            .await?;
        }
    }

    Ok(())
}

/// Cleans up data from the database that is no longer referenced, but must be
/// kept around for a little while to allow users to recover from accidental
/// deletions.
async fn stale_data_cleanup(pool: &Pool<Sqlite>) -> crate::Result<()> {
    let mut tx = pool.begin().await?;

    sqlx::query!(
        "DELETE FROM default_minecraft_capes WHERE minecraft_user_uuid NOT IN (SELECT uuid FROM minecraft_users)"
    )
    .execute(&mut *tx)
    .await?;
    sqlx::query!(
        "DELETE FROM custom_minecraft_skins WHERE minecraft_user_uuid NOT IN (SELECT uuid FROM minecraft_users)"
    )
    .execute(&mut *tx)
    .await?;

    tx.commit().await?;

    Ok(())
}