use anyhow::{bail, Context};
use chrono::Utc;
use console::style;
use sqlx::migrate::{AppliedMigration, Migrate, MigrateError, MigrationType, Migrator};
use sqlx::{AnyConnection, Connection};
use std::collections::{HashMap, HashSet};
use std::fs::{self, File};
use std::io::Write;
use std::path::Path;
use std::time::Duration;
fn create_file(
migration_source: &str,
file_prefix: &str,
description: &str,
migration_type: MigrationType,
) -> anyhow::Result<()> {
use std::path::PathBuf;
let mut file_name = file_prefix.to_string();
file_name.push_str("_");
file_name.push_str(&description.replace(' ', "_"));
file_name.push_str(migration_type.suffix());
let mut path = PathBuf::new();
path.push(migration_source);
path.push(&file_name);
println!("Creating {}", style(path.display()).cyan());
let mut file = File::create(&path).context("Failed to create migration file")?;
file.write_all(migration_type.file_content().as_bytes())?;
Ok(())
}
pub async fn add(
migration_source: &str,
description: &str,
reversible: bool,
) -> anyhow::Result<()> {
fs::create_dir_all(migration_source).context("Unable to create migrations directory")?;
let has_existing_migrations = fs::read_dir(migration_source)
.map(|mut dir| dir.next().is_some())
.unwrap_or(false);
let migrator = Migrator::new(Path::new(migration_source)).await?;
for migration in migrator.iter() {
if migration.migration_type.is_reversible() != reversible {
bail!(MigrateError::InvalidMixReversibleAndSimple);
}
}
let dt = Utc::now();
let file_prefix = dt.format("%Y%m%d%H%M%S").to_string();
if reversible {
create_file(
migration_source,
&file_prefix,
description,
MigrationType::ReversibleUp,
)?;
create_file(
migration_source,
&file_prefix,
description,
MigrationType::ReversibleDown,
)?;
} else {
create_file(
migration_source,
&file_prefix,
description,
MigrationType::Simple,
)?;
}
if !has_existing_migrations {
let quoted_source = if migration_source != "migrations" {
format!("{:?}", migration_source)
} else {
"".to_string()
};
print!(
r#"
Congratulations on creating your first migration!
Did you know you can embed your migrations in your application binary?
On startup, after creating your database connection or pool, add:
sqlx::migrate!({}).run(<&your_pool OR &mut your_connection>).await?;
Note that the compiler won't pick up new migrations if no Rust source files have changed.
You can create a Cargo build script to work around this with `sqlx migrate build-script`.
See: https://docs.rs/sqlx/0.5/sqlx/macro.migrate.html
"#,
quoted_source
);
}
Ok(())
}
pub async fn info(migration_source: &str, uri: &str) -> anyhow::Result<()> {
let migrator = Migrator::new(Path::new(migration_source)).await?;
let mut conn = AnyConnection::connect(uri).await?;
conn.ensure_migrations_table().await?;
let applied_migrations: HashMap<_, _> = conn
.list_applied_migrations()
.await?
.into_iter()
.map(|m| (m.version, m))
.collect();
for migration in migrator.iter() {
println!(
"{}: {}/{} {}",
if applied_migrations.contains_key(&migration.version) {
style("Installed").green().bold()
} else {
style("Pending").yellow().bold()
},
migration_source,
style(migration.version).cyan(),
migration.description,
);
}
Ok(())
}
fn validate_applied_migrations(
applied_migrations: &[AppliedMigration],
migrator: &Migrator,
ignore_missing: bool,
) -> Result<(), MigrateError> {
if ignore_missing {
return Ok(());
}
let migrations: HashSet<_> = migrator.iter().map(|m| m.version).collect();
for applied_migration in applied_migrations {
if !migrations.contains(&applied_migration.version) {
return Err(MigrateError::VersionMissing(applied_migration.version));
}
}
Ok(())
}
pub async fn run(
migration_source: &str,
uri: &str,
dry_run: bool,
ignore_missing: bool,
) -> anyhow::Result<()> {
let migrator = Migrator::new(Path::new(migration_source)).await?;
let mut conn = AnyConnection::connect(uri).await?;
conn.ensure_migrations_table().await?;
let version = conn.dirty_version().await?;
if let Some(version) = version {
bail!(MigrateError::Dirty(version));
}
let applied_migrations = conn.list_applied_migrations().await?;
validate_applied_migrations(&applied_migrations, &migrator, ignore_missing)?;
let applied_migrations: HashMap<_, _> = applied_migrations
.into_iter()
.map(|m| (m.version, m))
.collect();
for migration in migrator.iter() {
if migration.migration_type.is_down_migration() {
continue;
}
match applied_migrations.get(&migration.version) {
Some(applied_migration) => {
if migration.checksum != applied_migration.checksum {
bail!(MigrateError::VersionMismatch(migration.version));
}
}
None => {
let elapsed = if dry_run {
Duration::new(0, 0)
} else {
conn.apply(migration).await?
};
let text = if dry_run { "Can apply" } else { "Applied" };
println!(
"{}: {}/{} {} {}",
style(text).green().bold(),
migration_source,
style(migration.version).cyan(),
style(&migration.description),
style(format!("({:?})", elapsed)).dim()
);
}
}
}
Ok(())
}
pub async fn revert(
migration_source: &str,
uri: &str,
dry_run: bool,
ignore_missing: bool,
) -> anyhow::Result<()> {
let migrator = Migrator::new(Path::new(migration_source)).await?;
let mut conn = AnyConnection::connect(uri).await?;
conn.ensure_migrations_table().await?;
let version = conn.dirty_version().await?;
if let Some(version) = version {
bail!(MigrateError::Dirty(version));
}
let applied_migrations = conn.list_applied_migrations().await?;
validate_applied_migrations(&applied_migrations, &migrator, ignore_missing)?;
let applied_migrations: HashMap<_, _> = applied_migrations
.into_iter()
.map(|m| (m.version, m))
.collect();
let mut is_applied = false;
for migration in migrator.iter().rev() {
if !migration.migration_type.is_down_migration() {
continue;
}
if applied_migrations.contains_key(&migration.version) {
let elapsed = if dry_run {
Duration::new(0, 0)
} else {
conn.revert(migration).await?
};
let text = if dry_run { "Can apply" } else { "Applied" };
println!(
"{} {}/{} {} {}",
text,
style(migration.version).cyan(),
style(migration.migration_type.label()).green(),
migration.description,
style(format!("({:?})", elapsed)).dim()
);
is_applied = true;
break;
}
}
if !is_applied {
println!("No migrations available to revert");
}
Ok(())
}
pub fn build_script(migration_source: &str, force: bool) -> anyhow::Result<()> {
anyhow::ensure!(
Path::new("Cargo.toml").exists(),
"must be run in a Cargo project root"
);
anyhow::ensure!(
(force || !Path::new("build.rs").exists()),
"build.rs already exists; use --force to overwrite"
);
let contents = format!(
r#"// generated by `sqlx migrate build-script`
fn main() {{
// trigger recompilation when a new migration is added
println!("cargo:rerun-if-changed={}");
}}"#,
migration_source
);
fs::write("build.rs", contents)?;
println!("Created `build.rs`; be sure to check it into version control!");
Ok(())
}