use crate::DbkitError;
use crate::base_handler::{BaseHandler, FetchMode, WriteOp};
use crate::config::Backend;
use sqlx::{AnyPool, Row};
use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};
use tracing::{error, info};
pub struct InitializationHandler {
handler: BaseHandler,
backend: Backend,
}
impl InitializationHandler {
pub fn new(pool: AnyPool, backend: Backend) -> Self {
Self {
handler: BaseHandler::new(pool),
backend,
}
}
fn placeholder(&self, n: usize) -> String {
match self.backend {
Backend::Postgres => format!("${n}"),
Backend::MySql | Backend::Sqlite => "?".to_string(),
}
}
fn tracking_table_ddl(&self) -> &'static str {
match self.backend {
Backend::Postgres => {
"CREATE TABLE IF NOT EXISTS _dbkit_migrations (
name TEXT PRIMARY KEY,
hash TEXT NOT NULL,
applied_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
)"
}
Backend::MySql => {
"CREATE TABLE IF NOT EXISTS _dbkit_migrations (
name VARCHAR(255) PRIMARY KEY,
hash TEXT NOT NULL,
applied_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP
)"
}
Backend::Sqlite => {
"CREATE TABLE IF NOT EXISTS _dbkit_migrations (
name TEXT PRIMARY KEY,
hash TEXT NOT NULL,
applied_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP
)"
}
}
}
fn applied_at_as_text(&self) -> &'static str {
match self.backend {
Backend::MySql => "CAST(applied_at AS CHAR)",
Backend::Postgres | Backend::Sqlite => "CAST(applied_at AS TEXT)",
}
}
async fn ensure_tracking_table(&self) -> Result<(), DbkitError> {
self.handler
.execute_write(WriteOp::BatchDDL {
queries: &[self.tracking_table_ddl()],
})
.await?;
Ok(())
}
fn hash_sql(sql: &str) -> String {
let mut hasher = DefaultHasher::new();
sql.hash(&mut hasher);
format!("{:016x}", hasher.finish())
}
pub async fn run_named_migration(&self, name: &str, sql: &str) -> Result<(), DbkitError> {
self.ensure_tracking_table().await?;
let hash = Self::hash_sql(sql);
let select = format!(
"SELECT hash FROM _dbkit_migrations WHERE name = {}",
self.placeholder(1)
);
let result = self
.handler
.execute_write(WriteOp::Single {
query: select.as_str(),
params: vec![name.into()],
mode: FetchMode::Optional,
})
.await?;
if let Some(row) = result.optional()? {
let existing_hash: String = row.get(0);
if existing_hash == hash {
info!("migration '{}' already applied, skipping", name);
return Ok(());
} else {
return Err(DbkitError::Migration(format!(
"migration '{}' was already applied but content has changed (hash {} → {})",
name, existing_hash, hash
)));
}
}
info!("applying migration '{}'...", name);
let queries: Vec<String> = sql
.split(';')
.map(|s| s.trim().to_string())
.filter(|s| !s.is_empty())
.collect();
let query_refs: Vec<&str> = queries.iter().map(|s| s.as_str()).collect();
match self
.handler
.execute_write(WriteOp::BatchDDL {
queries: &query_refs,
})
.await
{
Ok(_) => {
info!(
"migration '{}': {} DDL statements executed",
name,
query_refs.len()
);
}
Err(e) => {
error!("migration '{}' failed: {:?}", name, e);
return Err(DbkitError::Migration(e.to_string()));
}
}
let insert = format!(
"INSERT INTO _dbkit_migrations (name, hash) VALUES ({}, {})",
self.placeholder(1),
self.placeholder(2)
);
self.handler
.execute_write(WriteOp::Single {
query: insert.as_str(),
params: vec![name.into(), hash.clone().into()],
mode: FetchMode::None,
})
.await?;
info!("migration '{}' recorded", name);
Ok(())
}
pub async fn run_migrations(&self, sql: &str) -> Result<(), DbkitError> {
info!("running database migrations...");
let queries: Vec<String> = sql
.split(';')
.map(|s| s.trim().to_string())
.filter(|s| !s.is_empty())
.collect();
let query_refs: Vec<&str> = queries.iter().map(|s| s.as_str()).collect();
match self
.handler
.execute_write(WriteOp::BatchDDL {
queries: &query_refs,
})
.await
{
Ok(_) => {
info!("{} DDL statements executed", query_refs.len());
}
Err(e) => {
error!("migration failed: {:?}", e);
return Err(DbkitError::Migration(e.to_string()));
}
}
Ok(())
}
pub async fn applied_migrations(&self) -> Result<Vec<(String, String, String)>, DbkitError> {
self.ensure_tracking_table().await?;
let select = format!(
"SELECT name, hash, {} FROM _dbkit_migrations ORDER BY applied_at, name",
self.applied_at_as_text()
);
let result = self
.handler
.execute_write(WriteOp::Single {
query: select.as_str(),
params: vec![],
mode: FetchMode::All,
})
.await?;
let rows = result.all()?;
Ok(rows
.iter()
.map(|row| {
let name: String = row.get(0);
let hash: String = row.get(1);
let applied_at: String = row.get(2);
(name, hash, applied_at)
})
.collect())
}
}