Improve edit api

Author: czocher

rusqlite_migration/Cargo.toml

@@ -25,6 +25,7 @@ include_dir = { version = "0.7.3", optional = true }
 tokio = { version = "1.25", features = ["macros"], optional = true }
 tokio-rusqlite = { version = "0.4.0", optional = true }
 log = "0.4"
+take_mut = "0.2.2"
 
 [dependencies.rusqlite]
 version = ">=0.29.0"

rusqlite_migration/src/builder.rs

@@ -1,6 +1,7 @@
 use std::iter::FromIterator;
 
 use include_dir::Dir;
+use take_mut::take;
 
 use crate::{loader::from_directory, MigrationHook, Result, M};
 
@@ -42,14 +43,11 @@ 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"));
+        take(&mut self.migrations[id - 1], f);
         self
     }
 
@@ -76,8 +74,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 +86,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/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);
 }