cdk/crates/cdk-sqlite/src/common.rs

use std::sync::Arc;
use std::time::Duration;

use rusqlite::{params, Connection};

use crate::pool::{Pool, ResourceManager};

/// The config need to create a new SQLite connection
#[derive(Debug)]
pub struct Config {
    path: Option<String>,
    password: Option<String>,
}

/// Sqlite connection manager
#[derive(Debug)]
pub struct SqliteConnectionManager;

impl ResourceManager for SqliteConnectionManager {
    type Config = Config;

    type Resource = Connection;

    type Error = rusqlite::Error;

    fn new_resource(
        config: &Self::Config,
    ) -> Result<Self::Resource, crate::pool::Error<Self::Error>> {
        let conn = if let Some(path) = config.path.as_ref() {
            Connection::open(path)?
        } else {
            Connection::open_in_memory()?
        };

        if let Some(password) = config.password.as_ref() {
            conn.execute_batch(&format!("pragma key = '{password}';"))?;
        }

        conn.execute_batch(
            r#"
            pragma busy_timeout = 10000;
            pragma journal_mode = WAL;
            pragma synchronous = normal;
            pragma temp_store = memory;
            pragma mmap_size = 30000000000;
            pragma cache = shared;
            "#,
        )?;

        Ok(conn)
    }
}

/// Create a configured rusqlite connection to a SQLite database.
/// For SQLCipher support, enable the "sqlcipher" feature and pass a password.
pub fn create_sqlite_pool(
    path: &str,
    #[cfg(feature = "sqlcipher")] password: String,
) -> Arc<Pool<SqliteConnectionManager>> {
    #[cfg(feature = "sqlcipher")]
    let password = Some(password);

    #[cfg(not(feature = "sqlcipher"))]
    let password = None;

    let (config, max_size) = if path.contains(":memory:") {
        (
            Config {
                path: None,
                password,
            },
            1,
        )
    } else {
        (
            Config {
                path: Some(path.to_owned()),
                password,
            },
            20,
        )
    };

    Pool::new(config, max_size, Duration::from_secs(5))
}

/// Migrates the migration generated by `build.rs`
pub fn migrate(conn: &mut Connection, migrations: &[(&str, &str)]) -> Result<(), rusqlite::Error> {
    let tx = conn.transaction()?;
    tx.execute(
        r#"
           CREATE TABLE IF NOT EXISTS migrations (
               name TEXT PRIMARY KEY,
               applied_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
           )
           "#,
        [],
    )?;

    if tx.query_row(
        r#"select count(*) from sqlite_master where name = '_sqlx_migrations'"#,
        [],
        |row| row.get::<_, i32>(0),
    )? == 1
    {
        tx.execute_batch(
            r#"
        INSERT INTO migrations
        SELECT
            version || '_' ||  REPLACE(description, ' ', '_') || '.sql',
            execution_time
        FROM _sqlx_migrations;
        DROP TABLE _sqlx_migrations;
        "#,
        )?;
    }

    // Apply each migration if it hasn’t been applied yet
    for (name, sql) in migrations {
        let already_applied: bool = tx.query_row(
            "SELECT EXISTS(SELECT 1 FROM migrations WHERE name = ?1)",
            params![name],
            |row| row.get(0),
        )?;

        if !already_applied {
            tx.execute_batch(sql)?;
            tx.execute("INSERT INTO migrations (name) VALUES (?1)", params![name])?;
        }
    }

    tx.commit()?;

    Ok(())
}