Improve edit api

rusqlite_migration/src/builder.rs

@@ -1,4 +1,4 @@
-use std::iter::FromIterator;
+use std::{iter::FromIterator, mem::take};
 
 use include_dir::Dir;
 
@@ -7,7 +7,7 @@ use crate::{loader::from_directory, MigrationHook, Result, M};
 /// Allows to build a `Vec<M<'u>>` with additional edits.
 #[derive(Default, Debug)]
 pub struct MigrationsBuilder<'u> {
-    migrations: Vec<M<'u>>,
+    migrations: Vec<Option<M<'u>>>,
 }
 
 impl<'u> MigrationsBuilder<'u> {
@@ -42,28 +42,25 @@ impl<'u> MigrationsBuilder<'u> {
     ///
     /// Panics if no migration with the `id` provided exists.
     #[must_use]
-    pub fn edit(mut self, id: usize, f: impl Fn(&mut M)) -> Self {
+    pub fn edit(mut self, id: usize, f: impl Fn(M) -> M) -> Self {
         if id < 1 {
             panic!("id cannot be equal to 0");
         }
-        f(self
-            .migrations
-            .get_mut(id - 1)
-            .expect("No migration with the given index"));
+        self.migrations[id - 1] = take(&mut self.migrations[id - 1]).map(f);
         self
     }
 
     /// Finalizes the builder and creates either a [`crate::Migrations`] or a
     /// [`crate::AsyncMigrations`] instance.
-    pub fn finalize<T: FromIterator<M<'u>>>(self) -> T {
-        T::from_iter(self.migrations)
+    pub fn finalize<T: FromIterator<M<'u>>>(mut self) -> T {
+        T::from_iter(self.migrations.drain(..).flatten())
     }
 }
 
 impl<'u> FromIterator<M<'u>> for MigrationsBuilder<'u> {
     fn from_iter<T: IntoIterator<Item = M<'u>>>(iter: T) -> Self {
         Self {
-            migrations: Vec::from_iter(iter),
+            migrations: Vec::from_iter(iter.into_iter().map(Some)),
         }
     }
 }
@@ -76,8 +73,9 @@ impl<'u> M<'u> {
     /// Use [`M::up_with_hook`] instead if you're creating a new migration.
     /// This method is meant for editing existing transactions
     /// when using the [`MigrationsBuilder`].
-    pub fn set_up_hook(&mut self, hook: impl MigrationHook + 'static) {
+    pub fn set_up_hook(mut self, hook: impl MigrationHook + 'static) -> Self {
         self.up_hook = Some(hook.clone_box());
+        self
     }
 
     /// Replace the `down_hook` in the given migration with the provided one.
@@ -87,7 +85,8 @@ impl<'u> M<'u> {
     /// Use [`M::down_with_hook`] instead if you're creating a new migration.
     /// This method is meant for editing existing transactions
     /// when using the [`MigrationsBuilder`].
-    pub fn set_down_hook(&mut self, hook: impl MigrationHook + 'static) {
+    pub fn set_down_hook(mut self, hook: impl MigrationHook + 'static) -> Self {
         self.down_hook = Some(hook.clone_box());
+        self
     }
 }

rusqlite_migration/src/lib.rs

@@ -452,9 +452,12 @@ impl<'m> Migrations<'m> {
     /// or don't contain at least a valid `up.sql` file.
     #[cfg(feature = "from-directory")]
     pub fn from_directory(dir: &'static Dir<'static>) -> Result<Self> {
-        Ok(Self {
-            ms: from_directory(dir)?,
-        })
+        let migrations = from_directory(dir)?
+            .into_iter()
+            .collect::<Option<Vec<_>>>()
+            .ok_or(Error::FileLoad("Could not load migrations".to_string()))?;
+
+        Ok(Self { ms: migrations })
     }
 
     /// **Deprecated**: [`Migrations`] now implements [`FromIterator`], so use [`Migrations::from_iter`] instead.

rusqlite_migration/src/loader.rs

@@ -90,7 +90,7 @@ impl<'u> From<&MigrationFile> for M<'u> {
     }
 }
 
-pub(crate) fn from_directory(dir: &'static Dir<'static>) -> Result<Vec<M<'static>>> {
+pub(crate) fn from_directory(dir: &'static Dir<'static>) -> Result<Vec<Option<M<'static>>>> {
     let mut migrations: Vec<Option<M>> = vec![None; dir.dirs().count()];
 
     for dir in dir.dirs() {
@@ -127,8 +127,5 @@ pub(crate) fn from_directory(dir: &'static Dir<'static>) -> Result<Vec<M<'static
     }
 
     // The values are returned in the order of the keys, i.e. of IDs
-    migrations
-        .into_iter()
-        .collect::<Option<Vec<_>>>()
-        .ok_or(Error::FileLoad("Could not load migrations".to_string()))
+    Ok(migrations)
 }

rusqlite_migration/src/tests/builder.rs

@@ -7,13 +7,13 @@ use crate::{MigrationsBuilder, M};
 fn test_non_existing_index() {
     let ms = vec![M::up("CREATE TABLE t(a);")];
 
-    let _ = MigrationsBuilder::from_iter(ms.clone()).edit(100, move |_t| {});
+    let _ = MigrationsBuilder::from_iter(ms.clone()).edit(100, move |t| t);
 }
 
 #[test]
 #[should_panic]
 fn test_0_index() {
     let ms = vec![M::up("CREATE TABLE t(a);")];
 
-    let _ = MigrationsBuilder::from_iter(ms).edit(0, move |_t| {});
+    let _ = MigrationsBuilder::from_iter(ms).edit(0, move |t| t);
 }