use std::{borrow::Cow, collections::BTreeMap, io};
use camino::Utf8PathBuf;
use rusqlite::Connection;
type Migration = Cow<'static, str>;
const VERSION_TABLE: &str = "monarch_db_schema_version";
#[derive(Debug, Clone)]
#[cfg_attr(feature = "serde", derive(serde::Deserialize))]
pub struct ConnectionConfiguration {
#[cfg_attr(feature = "serde", serde(default))]
pub database: Option<Utf8PathBuf>,
}
#[derive(Debug, Clone)]
#[cfg_attr(feature = "serde", derive(serde::Deserialize))]
pub struct MonarchConfiguration {
pub name: String,
pub enable_foreign_keys: bool,
pub migration_directory: Utf8PathBuf,
}
#[derive(Debug, Clone)]
pub struct StaticMonarchConfiguration<const N: usize> {
pub name: &'static str,
pub enable_foreign_keys: bool,
pub migrations: [&'static str; N],
}
impl<const N: usize> From<StaticMonarchConfiguration<N>> for MonarchDB {
fn from(configuration: StaticMonarchConfiguration<N>) -> Self {
MonarchDB {
name: configuration.name.into(),
enable_foreign_keys: configuration.enable_foreign_keys,
migrations: configuration
.migrations
.iter()
.map(|q| Cow::Borrowed(*q))
.collect(),
}
}
}
#[derive(Debug)]
pub struct MonarchDB {
name: Cow<'static, str>,
enable_foreign_keys: bool,
migrations: Vec<Migration>,
}
impl MonarchDB {
pub fn open_in_memory(&self) -> rusqlite::Result<Connection> {
let connection = Connection::open_in_memory()?;
self.migrate(connection)
}
pub fn from_configuration(configuration: MonarchConfiguration) -> io::Result<Self> {
let mut migrations = BTreeMap::new();
for diritem in configuration.migration_directory.read_dir_utf8()? {
let entry = diritem?;
if entry.file_type()?.is_file() {
let query = std::fs::read_to_string(entry.path())?;
migrations.insert(entry.file_name().to_owned(), Cow::from(query));
}
}
Ok(MonarchDB {
name: configuration.name.into(),
enable_foreign_keys: configuration.enable_foreign_keys,
migrations: migrations.into_values().collect(),
})
}
pub fn current_version(&self) -> u32 {
self.migrations.len() as u32
}
fn get_migration(&self, version: u32) -> Option<&str> {
self.migrations
.get(version as usize)
.map(|query| query.as_ref())
}
pub fn create_connection(
&self,
configuration: &ConnectionConfiguration,
) -> rusqlite::Result<Connection> {
let connection = if let Some(path) = configuration.database.as_deref() {
Connection::open(path)?
} else {
Connection::open_in_memory()?
};
self.migrate(connection)
}
pub fn migrate(&self, mut connection: Connection) -> rusqlite::Result<Connection> {
let migrations = Migrations {
connection: &mut connection,
monarch: self,
};
migrations.prepare()?;
Ok(connection)
}
pub fn migrations<'c>(&'c self, connection: &'c mut Connection) -> Migrations<'c> {
Migrations {
connection,
monarch: self,
}
}
}
pub struct Migrations<'c> {
connection: &'c mut Connection,
monarch: &'c MonarchDB,
}
impl<'c> Migrations<'c> {
#[tracing::instrument(level = "trace", skip_all, fields(monarch=%self.monarch.name))]
pub fn prepare(self) -> rusqlite::Result<()> {
if self.monarch.enable_foreign_keys {
tracing::trace!("Set foreign keys");
self.connection.pragma_update(None, "foreign_keys", true)?;
}
self.migrate()?;
Ok(())
}
fn migrate(self) -> rusqlite::Result<()> {
let tx = self.connection.transaction()?;
let mut version = select_schema_version(&tx, &self.monarch.name)?;
while version < self.monarch.current_version() {
let query = self
.monarch
.get_migration(version)
.expect("version <-> migration mismatch");
tracing::trace!("Running migration to version {}", version + 1);
tx.execute_batch(query)?;
version += 1;
}
set_schema_version(&tx, &self.monarch.name, version)?;
tx.commit()?;
tracing::debug!("Migrations complete");
Ok(())
}
}
fn create_schema_version_table(connection: &Connection) -> rusqlite::Result<()> {
let mut stmt = connection.prepare(include_str!("00.versions.sql"))?;
stmt.execute([])?;
Ok(())
}
fn insert_initial_schema_version(connection: &Connection, name: &str) -> rusqlite::Result<()> {
let mut stmt = connection.prepare(&format!(
"INSERT INTO {VERSION_TABLE} (monarch_schema, version) VALUES (:name, 0)"
))?;
stmt.execute(&[(":name", name)])?;
Ok(())
}
fn select_schema_version(connection: &Connection, name: &str) -> rusqlite::Result<u32> {
let mut stmt = connection.prepare("SELECT name FROM sqlite_master WHERE name = :table")?;
let has_version_tbl: Option<Result<String, _>> = stmt
.query_map(&[(":table", VERSION_TABLE)], |row| row.get(0))?
.next();
match has_version_tbl {
Some(Ok(_)) => {}
Some(Err(error)) => {
return Err(error);
}
None => {
tracing::trace!("Create schema version table {VERSION_TABLE}");
create_schema_version_table(connection)?;
insert_initial_schema_version(connection, name)?;
return Ok(0u32);
}
};
let mut stmt = connection.prepare(&format!(
"SELECT version FROM {VERSION_TABLE} WHERE monarch_schema = :name"
))?;
let version: Option<u32> = stmt
.query_map(&[(":name", name)], |row| row.get::<_, u32>(0))?
.next()
.transpose()?;
if let Some(version) = version {
tracing::trace!(%version, "Get schema version");
Ok(version)
} else {
tracing::trace!("Insert new version for {name}");
insert_initial_schema_version(connection, name)?;
Ok(0)
}
}
fn set_schema_version(connection: &Connection, name: &str, version: u32) -> rusqlite::Result<()> {
tracing::trace!(%version, "Set schema version for {name}");
let mut stmt = connection.prepare(&format!(
"UPDATE {VERSION_TABLE} SET version = :version WHERE monarch_schema = :name"
))?;
stmt.execute(rusqlite::named_params! { ":version": version, ":name": name})?;
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_static_monarch_configuration_creation() {
let config = StaticMonarchConfiguration {
name: "test_db",
enable_foreign_keys: true,
migrations: [
"CREATE TABLE users (id INTEGER PRIMARY KEY, name TEXT NOT NULL);",
"ALTER TABLE users ADD COLUMN email TEXT;",
],
};
assert_eq!(config.name, "test_db");
assert!(config.enable_foreign_keys);
assert_eq!(config.migrations.len(), 2);
}
#[test]
fn test_static_configuration_to_monarch_db() {
let config = StaticMonarchConfiguration {
name: "test_db",
enable_foreign_keys: false,
migrations: ["CREATE TABLE posts (id INTEGER PRIMARY KEY, title TEXT NOT NULL);"],
};
let monarch_db: MonarchDB = config.into();
assert_eq!(monarch_db.current_version(), 1);
assert_eq!(monarch_db.name, "test_db");
assert!(!monarch_db.enable_foreign_keys);
}
#[test]
fn test_open_in_memory_with_static_migrations() -> rusqlite::Result<()> {
let config = StaticMonarchConfiguration {
name: "test_memory_db",
enable_foreign_keys: true,
migrations: [
"CREATE TABLE items (id INTEGER PRIMARY KEY, name TEXT NOT NULL);",
"CREATE INDEX idx_items_name ON items(name);",
],
};
let monarch_db: MonarchDB = config.into();
let connection = monarch_db.open_in_memory()?;
let mut stmt = connection
.prepare("SELECT name FROM sqlite_master WHERE type='table' AND name='items'")?;
let table_exists: bool = stmt.query_map([], |_| Ok(true))?.next().is_some();
assert!(table_exists);
let mut stmt = connection.prepare(
"SELECT name FROM sqlite_master WHERE type='index' AND name='idx_items_name'",
)?;
let index_exists: bool = stmt.query_map([], |_| Ok(true))?.next().is_some();
assert!(index_exists);
Ok(())
}
#[test]
fn test_create_connection_with_static_migrations() -> rusqlite::Result<()> {
let config = StaticMonarchConfiguration {
name: "test_file_db",
enable_foreign_keys: false,
migrations: [
"CREATE TABLE products (id INTEGER PRIMARY KEY, name TEXT NOT NULL, price REAL);",
],
};
let monarch_db: MonarchDB = config.into();
let connection_config = ConnectionConfiguration { database: None };
let connection = monarch_db.create_connection(&connection_config)?;
let mut stmt = connection
.prepare("SELECT name FROM sqlite_master WHERE type='table' AND name='products'")?;
let table_exists: bool = stmt.query_map([], |_| Ok(true))?.next().is_some();
assert!(table_exists);
connection.execute(
"INSERT INTO products (name, price) VALUES (?, ?)",
["Test Product", "19.99"],
)?;
let mut stmt = connection.prepare("SELECT COUNT(*) FROM products")?;
let count: i64 = stmt.query_row([], |row| row.get(0))?;
assert_eq!(count, 1);
Ok(())
}
#[test]
fn test_migration_versioning() -> rusqlite::Result<()> {
let config = StaticMonarchConfiguration {
name: "versioning_test",
enable_foreign_keys: false,
migrations: [
"CREATE TABLE v1_table (id INTEGER PRIMARY KEY);",
"CREATE TABLE v2_table (id INTEGER PRIMARY KEY);",
"CREATE TABLE v3_table (id INTEGER PRIMARY KEY);",
],
};
let monarch_db: MonarchDB = config.into();
assert_eq!(monarch_db.current_version(), 3);
let connection = monarch_db.open_in_memory()?;
let table_names = ["v1_table", "v2_table", "v3_table"];
for table_name in table_names {
let mut stmt = connection.prepare(&format!(
"SELECT name FROM sqlite_master WHERE type='table' AND name='{table_name}'"
))?;
let table_exists: bool = stmt.query_map([], |_| Ok(true))?.next().is_some();
assert!(table_exists, "Table {table_name} should exist");
}
Ok(())
}
}