riftenlabs-indexer/src/db/init.rs

// Copyright (C) 2024-2026 Whiterun LLC
//
// This software is licensed under the GNU Affero General Public License (AGPL), version 3.0 or later.
// A copy of the license can be found in the LICENSE file or at https://www.gnu.org/licenses/agpl-3.0.html

use anyhow::Result;
use log::info;
use rusqlite::OpenFlags;
use std::path::Path;
use std::sync::Arc;
use std::thread::sleep;
use std::time::{Duration, Instant};

use super::{DBPool, DB};
use crate::db::bcmr::prepare_tables as bcmr_prepare_tables;
use crate::db::cauldron::config::check_db_version;
use crate::db::cauldron::prepare_tables as cauldron_prepare_tables;
use crate::db::crc20::prepare_tables as crc20_prepare_tables;
use crate::db::oracle::prepare_tables as oracle_prepare_tables;

/// Initialize a database connection with proper settings
pub fn db_initialize(c: &mut rusqlite::Connection) -> Result<(), rusqlite::Error> {
    // Set busy timeout for normal queries.
    c.busy_timeout(Duration::from_secs(30))?;
    c.execute("PRAGMA foreign_keys=1;", [])?;

    // PRAGMA journal_mode may not honor busy_timeout, so retry manually.
    let start = Instant::now();
    let timeout = Duration::from_secs(30);
    loop {
        match c.execute_batch("PRAGMA journal_mode=WAL;") {
            Ok(_) => break,
            Err(rusqlite::Error::SqliteFailure(err, _))
                if err.code == rusqlite::ErrorCode::DatabaseBusy =>
            {
                if start.elapsed() >= timeout {
                    return Err(rusqlite::Error::SqliteFailure(
                        err,
                        Some("database locked after retries".into()),
                    ));
                }
                sleep(Duration::from_millis(100));
            }
            Err(e) => return Err(e),
        }
    }
    Ok(())
}

/// Verify that database settings are correct
pub fn db_sanity_check(c: &rusqlite::Connection) -> rusqlite::Result<()> {
    // Verify that foreign_keys are enabled
    let mut stmt = c.prepare("PRAGMA foreign_keys;")?;
    let foreign_keys_enabled: i32 = stmt.query_row([], |row| row.get(0))?;
    if foreign_keys_enabled != 1 {
        panic!(
            "Foreign key enforcement is not enabled for the connection. (result: {foreign_keys_enabled})"
        );
    }
    Ok(())
}

/// Create read and write database pools for a given database path
fn create_db_pool(db_path: &str, read_max_size: u32) -> (bool, DBPool, DBPool) {
    let db_exists = Path::new(db_path).exists();
    info!("Initializing connection to {db_path} (read_slots={read_max_size})");

    let write_manager = r2d2_sqlite::SqliteConnectionManager::file(db_path)
        .with_flags(OpenFlags::SQLITE_OPEN_READ_WRITE | OpenFlags::SQLITE_OPEN_CREATE)
        .with_init(db_initialize);

    let write_pool = Arc::new(
        r2d2::Pool::builder()
            .max_size(2) // Write pools need fewer connections
            .build(write_manager)
            .expect("Failed to initialize database"),
    );

    let read_manager = r2d2_sqlite::SqliteConnectionManager::file(db_path)
        .with_flags(OpenFlags::SQLITE_OPEN_READ_ONLY)
        .with_init(db_initialize);

    let read_pool = Arc::new(
        r2d2::Pool::builder()
            .max_size(read_max_size)
            .build(read_manager)
            .expect("Failed to initialize database"),
    );

    db_sanity_check(&read_pool.get().unwrap()).expect("db sanity check failed");
    db_sanity_check(&write_pool.get().unwrap()).expect("db sanity check failed");

    (db_exists, write_pool, read_pool)
}

/// Get the database directory path based on network
pub fn get_db_directory(network: &str) -> &'static str {
    if network.to_lowercase() == "chipnet" {
        "chipnet"
    } else {
        "."
    }
}

/// Create network-specific directory if needed
pub fn ensure_db_directory(db_dir: &str) -> Result<()> {
    if db_dir != "." {
        std::fs::create_dir_all(db_dir).map_err(|e| {
            anyhow::anyhow!("Failed to create database directory '{}': {}", db_dir, e)
        })?;
    }
    Ok(())
}

/// Get the full database path for a filename based on the database directory
fn db_path(db_dir: &str, filename: &str) -> String {
    if db_dir == "." {
        filename.to_string()
    } else {
        format!("{}/{}", db_dir, filename)
    }
}

/// Per-database read pool sizes
pub struct ReadSlots {
    pub cauldron: u32,
    pub bcmr: u32,
    pub crc20: u32,
    pub oracle: u32,
}

/// Initialize all databases and return a DB struct
pub fn initialize_databases(network: &str, read_slots: ReadSlots) -> Result<DB> {
    let db_dir = get_db_directory(network);
    ensure_db_directory(db_dir)?;

    // Initialize cauldron database
    let (db_exists, cauldron_db_write, cauldron_db_read) =
        create_db_pool(&db_path(db_dir, "cauldron.db"), read_slots.cauldron);
    if !db_exists {
        cauldron_prepare_tables(
            &cauldron_db_write
                .get()
                .expect("failed to get sqlite connection"),
        );
    } else {
        // Check database version for existing database
        check_db_version(&cauldron_db_read.get().unwrap()).expect("Database version check failed");
    }

    // Initialize BCMR database
    let (db_exists, bcmr_db_write, bcmr_db_read) =
        create_db_pool(&db_path(db_dir, "bcmr.db"), read_slots.bcmr);
    if !db_exists {
        bcmr_prepare_tables(
            &bcmr_db_write
                .get()
                .expect("failed to create sql connection"),
        );
    }

    // Initialize CRC20 database
    let (db_exists, crc20_db_write, crc20_db_read) =
        create_db_pool(&db_path(db_dir, "crc20.db"), read_slots.crc20);
    if !db_exists {
        crc20_prepare_tables(
            &crc20_db_write
                .get()
                .expect("failed to create sqlite crc20 connection"),
        );
    }

    // Initialize oracle database
    let (db_exists, oracle_db_write, oracle_db_read) =
        create_db_pool(&db_path(db_dir, "oracle.db"), read_slots.oracle);
    if !db_exists {
        oracle_prepare_tables(
            &oracle_db_write
                .get()
                .expect("failed to create sqlite oracle connection"),
        );
    }

    Ok(DB {
        cauldron_w: cauldron_db_write,
        cauldron_r: cauldron_db_read,
        bcmr_w: bcmr_db_write,
        bcmr_r: bcmr_db_read,
        crc20_w: crc20_db_write,
        crc20_r: crc20_db_read,
        oracle_w: oracle_db_write,
        oracle_r: oracle_db_read,
    })
}