wsb 0.258.6441

use crate::persist::global::config::{DatabaseConfig, Schema};
use crate::prelude::*;
use r2d2::{CustomizeConnection, Pool};
use r2d2_sqlite::SqliteConnectionManager;
use rusqlite::OpenFlags;

pub type ConnectionPool = Pool<SqliteConnectionManager>;

#[derive(Debug)]
struct TransactionCleanupCustomizer;

impl CustomizeConnection<rusqlite::Connection, rusqlite::Error> for TransactionCleanupCustomizer {
    fn on_acquire(
        &self,
        conn: &mut rusqlite::Connection,
    ) -> std::result::Result<(), rusqlite::Error> {
        let _ = conn.execute("ROLLBACK", []);
        Ok(())
    }
}

#[derive(Debug, Clone)]
pub struct DatabaseService {
    pool: ConnectionPool,
}

/// How many in-memory databases this process has made.
///
/// Only to make each one's name different from the last. A shared-cache
/// in-memory sqlite database IS its name, so a fixed one is a single database
/// wearing every caller's expectations.
static IN_MEMORY_DATABASES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);

/// Put a database file into its journal mode, once, before any pool opens it.
///
/// WAL is a property of the FILE, and switching a file to it takes an
/// exclusive lock that sqlite will not wait for while other connections hold
/// the file — it answers "database is locked" at once, whatever the busy
/// timeout. The pool opens its connections together, so on a new database
/// every one but the first failed to switch it: `ERROR r2d2: database is
/// locked` on every `wsb init`, and connections the pool had to make again.
/// Switched here, on the only connection there is, the pooled ones find it
/// already in WAL and have nothing to contend for.
fn prepare_file(path: &Path, config: &DatabaseConfig) -> rusqlite::Result<()> {
    if !config.db_enable_wal {
        return Ok(());
    }
    let conn = rusqlite::Connection::open(path)?;
    conn.pragma_update(None, "busy_timeout", config.db_busy_timeout_ms)?;
    conn.pragma_update(None, "journal_mode", "WAL")?;
    Ok(())
}

/// What every pooled connection to a database file is set up with.
///
/// `busy_timeout` first, so everything after it waits out a writer rather
/// than failing on one.
fn configure(conn: &rusqlite::Connection, config: &DatabaseConfig) -> rusqlite::Result<()> {
    conn.pragma_update(None, "busy_timeout", config.db_busy_timeout_ms)?;
    if config.db_enable_wal {
        conn.pragma_update(None, "journal_mode", "WAL")?;
    }
    if config.db_enable_foreign_keys {
        conn.pragma_update(None, "foreign_keys", "ON")?;
    }
    conn.pragma_update(None, "synchronous", "NORMAL")?;
    conn.pragma_update(None, "cache_size", 10000)?;
    conn.pragma_update(None, "temp_store", "memory")?;
    conn.pragma_update(None, "mmap_size", 134217728)?;
    conn.pragma_update(None, "wal_autocheckpoint", 1000)?;
    Ok(())
}

impl DatabaseService {
    pub fn new(config: &DatabaseConfig) -> anyhow::Result<Self> {
        if config.db_path == ":memory:" {
            info!("Database: Using in-memory database");
            return Self::in_memory(config.schema, Some(config.db_max_connections));
        }

        if let Some(parent) = Path::new(&config.db_path).parent() {
            std::fs::create_dir_all(parent)
                .map_err(|e| anyhow::anyhow!("Failed to create directory: {}", e))?;
        }

        prepare_file(Path::new(&config.db_path), config)
            .map_err(|e| anyhow::anyhow!("Failed to prepare {}: {}", config.db_path, e))?;

        let config_clone = config.clone();
        let manager = SqliteConnectionManager::file(&config.db_path)
            .with_flags(
                OpenFlags::SQLITE_OPEN_READ_WRITE
                    | OpenFlags::SQLITE_OPEN_CREATE
                    | OpenFlags::SQLITE_OPEN_FULL_MUTEX,
            )
            .with_init(move |conn| configure(conn, &config_clone));

        let pool = Pool::builder()
            .max_size(config.db_max_connections)
            .connection_timeout(std::time::Duration::from_secs(
                config.db_connection_timeout_seconds,
            ))
            .connection_customizer(Box::new(TransactionCleanupCustomizer))
            .build(manager)
            .map_err(|e| anyhow::anyhow!("Failed to create database pool: {}", e))?;

        let service = Self { pool };

        service.run_migrations(config.schema)?;

        Ok(service)
    }

    /// A worktree database in memory.
    pub fn new_in_memory() -> anyhow::Result<Self> {
        Self::in_memory(Schema::Worktree, None)
    }

    /// A worktree database in memory, with room for several connections.
    pub fn new_in_memory_with_config(max_connections: Option<u32>) -> anyhow::Result<Self> {
        Self::in_memory(Schema::Worktree, max_connections)
    }

    /// A sessions database in memory.
    pub fn new_sessions_in_memory() -> anyhow::Result<Self> {
        Self::in_memory(Schema::Sessions, None)
    }

    fn in_memory(schema: Schema, max_connections: Option<u32>) -> anyhow::Result<Self> {
        let pool_size = max_connections.unwrap_or(1);

        // A name of its own, so two of these are two databases.
        //
        // It was `file:memdb?mode=memory&cache=shared` — one fixed name — and a
        // shared-cache in-memory database is identified BY that name. Every
        // caller asking for "an in-memory database" got the same one: tests
        // running in parallel migrated over each other and wrote into each
        // other's tables, and what surfaced was
        //
        //     Migration failed: `error asserting migrations table`,
        //     `database schema is locked: main`
        //
        // from a test that had asked for a database nobody else was using.
        // Intermittent, because it needed two of them to overlap.
        //
        // `cache=shared` stays: the pool may open several connections and they
        // have to see the same memory. It is the NAME that has to differ.
        let name = format!(
            "memdb-{}-{}",
            std::process::id(),
            IN_MEMORY_DATABASES.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
        );
        let manager = SqliteConnectionManager::file(format!("file:{name}?mode=memory&cache=shared"))
            .with_flags(
                OpenFlags::SQLITE_OPEN_READ_WRITE
                    | OpenFlags::SQLITE_OPEN_CREATE
                    | OpenFlags::SQLITE_OPEN_URI
                    | OpenFlags::SQLITE_OPEN_FULL_MUTEX,
            )
            .with_init(|conn| {
                conn.pragma_update(None, "foreign_keys", "ON")?;
                conn.pragma_update(None, "busy_timeout", 30000)?;
                conn.pragma_update(None, "synchronous", "NORMAL")?;
                conn.pragma_update(None, "cache_size", 10000)?;
                conn.pragma_update(None, "temp_store", "memory")?;
                Ok(())
            });

        let pool = Pool::builder()
            .max_size(pool_size)
            .connection_timeout(std::time::Duration::from_secs(5))
            .connection_customizer(Box::new(TransactionCleanupCustomizer))
            .build(manager)
            .map_err(|e| anyhow::anyhow!("Failed to create in-memory database pool: {}", e))?;

        let service = Self { pool };

        service.run_migrations(schema)?;

        Ok(service)
    }

    pub fn get_connection(
        &self,
    ) -> anyhow::Result<r2d2::PooledConnection<SqliteConnectionManager>> {
        self.pool.get().map_err(Into::into)
    }

    pub fn run_migrations(&self, schema: Schema) -> anyhow::Result<()> {
        let conn = self.get_connection()?;
        crate::persist::global::migrations::run_migrations(conn, schema)?;
        tracing::debug!("Database migrations completed successfully");
        Ok(())
    }

    pub fn health_check(&self) -> anyhow::Result<()> {
        let conn = self.get_connection()?;
        let _: i32 = conn.query_row("SELECT 1", [], |row| row.get(0))?;
        Ok(())
    }
}

#[cfg(test)]
mod in_memory_tests {
    //! Two in-memory databases are two databases.

    use super::DatabaseService;

    /// TEST11513: asking twice for an in-memory database gives two of them.
    ///
    /// It gave one. A shared-cache in-memory sqlite database is identified by
    /// its NAME, and every caller used the same fixed one — so every test that
    /// asked for "an in-memory database" got the database every other test was
    /// using. They migrated over each other and wrote into each other's
    /// tables, and what surfaced, intermittently, was a test failing on
    ///
    /// ```text
    /// Migration failed: `error asserting migrations table`,
    /// `database schema is locked: main`
    /// ```
    ///
    /// Intermittent because it took two of them overlapping — which is to say
    /// it got likelier as the suite got faster.
    ///
    /// Written as isolation rather than as "the names differ": what matters is
    /// that a row put in one is not visible in the other, and a name is only
    /// how that is arranged.
    #[test]
    fn test11513_two_in_memory_databases_do_not_share_a_row() {
        let one = DatabaseService::new_in_memory().expect("the first database");
        let two = DatabaseService::new_in_memory().expect("the second database");

        let write = one.get_connection().expect("a connection to the first");
        write
            .execute("CREATE TABLE only_in_one (id INTEGER PRIMARY KEY)", [])
            .expect("the first database takes a table of its own");
        write
            .execute("INSERT INTO only_in_one (id) VALUES (1)", [])
            .expect("and a row in it");

        let read = two.get_connection().expect("a connection to the second");
        let found: Result<i64, _> = read.query_row("SELECT id FROM only_in_one", [], |row| row.get(0));
        assert!(
            found.is_err(),
            "the second database can see the first one's table, so they are one \
             database and every test that asked for its own got everybody's"
        );
    }

    /// TEST11514: an in-memory database is usable by more than one connection.
    ///
    /// The reason the name carries `cache=shared` at all: the pool may open
    /// several connections and they have to see the same memory. Making the
    /// names unique must not have quietly turned that off, which would trade
    /// an intermittent failure for a mysterious empty table.
    #[test]
    fn test11514_one_in_memory_database_is_shared_across_its_connections() {
        let service =
            DatabaseService::new_in_memory_with_config(Some(4)).expect("a database with a pool");

        let write = service.get_connection().expect("one connection");
        write
            .execute("CREATE TABLE shared_rows (id INTEGER PRIMARY KEY)", [])
            .expect("a table");
        write
            .execute("INSERT INTO shared_rows (id) VALUES (7)", [])
            .expect("a row");

        // HELD while the second is taken, so the pool has to open a second
        // one. Returning the first and asking again gets that same connection
        // back — and this then passes with the shared cache switched off,
        // proving nothing about sharing at all.
        let read = service.get_connection().expect("a second, concurrent connection");
        let found: i64 = read
            .query_row("SELECT id FROM shared_rows", [], |row| row.get(0))
            .expect("the same database, seen from another connection");
        assert_eq!(found, 7);
    }
}

#[cfg(test)]
mod configure_tests {
    use super::configure;
    use crate::persist::global::config::DatabaseConfig;

    /// TEST12140: connections configuring a new database together all succeed.
    ///
    /// What the pool does when it opens: several connections at once, on a
    /// new file. Each one switching the file to WAL needs an exclusive lock
    /// sqlite will not wait for while the others hold the file, so all but one
    /// failed — `database is locked`, logged by r2d2 on every `wsb init`. The
    /// file is prepared first, exactly as `DatabaseService::new` does.
    #[test]
    fn test12140_connections_configuring_a_new_database_together_all_succeed() {
        for attempt in 0..5 {
            let dir = tempfile::tempdir().expect("a directory");
            let path = dir.path().join("fresh.db");
            let config = DatabaseConfig::new(&path);
            super::prepare_file(&path, &config).expect("the file is prepared");
            let start = std::sync::Arc::new(std::sync::Barrier::new(8));
            let handles: Vec<_> = (0..8)
                .map(|_| {
                    let (path, config, start) = (path.clone(), config.clone(), start.clone());
                    std::thread::spawn(move || {
                        let conn = rusqlite::Connection::open(&path).expect("the file opens");
                        start.wait();
                        configure(&conn, &config).map_err(|e| e.to_string())
                    })
                })
                .collect();
            for handle in handles {
                let outcome = handle.join().expect("the thread finishes");
                assert!(
                    outcome.is_ok(),
                    "attempt {attempt}: a connection failed to configure a new database: {outcome:?}"
                );
            }
        }
    }
}