properly rollback failed migrations

m1guelpf · Sep 2, 2023 · f5aa0e0 · f5aa0e0
1 parent 7d26156
commit f5aa0e0
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 19 deletions.
diff --git a/ensemble/src/migrations/migrator.rs b/ensemble/src/migrations/migrator.rs
@@ -95,14 +95,23 @@ impl Migrator {
                 .map_err(|_| Error::Lock)?
                 .replace(self.connection);
 
-            migration.up().await?;
+            let migration_result = migration.up().await;
 
             self.connection = MIGRATE_CONN
                 .try_lock()
                 .map_err(|_| Error::Lock)?
                 .take()
                 .ok_or(Error::Lock)?;
 
+            if let Err(e) = migration_result {
+                self.connection
+                    .exec("rollback", vec![])
+                    .await
+                    .map_err(|e| Error::Database(e.to_string()))?;
+
+                return Err(e);
+            }
+
             self.connection
                 .exec(
                     "insert into migrations (migration, batch) values (?, ?)",
@@ -217,17 +226,17 @@ fn migrations_table_query() -> &'static str {
     match connection::which_db() {
         Database::MySQL => {
             "create table if not exists migrations (
-            id int unsigned not null auto_increment primary key,
-            migration varchar(255) not null,
-            batch int not null
-        )"
+                id int unsigned not null auto_increment primary key,
+                migration varchar(255) not null unique,
+                batch int not null
+            )"
         }
         Database::PostgreSQL => {
             "create table if not exists migrations (
-            id serial primary key,
-            migration varchar(255) not null,
-            batch int not null
-        )"
+                id serial primary key,
+                migration varchar(255) not null unique,
+                batch int not null
+            )"
         }
     }
 }
diff --git a/ensemble/src/migrations/schema/mod.rs b/ensemble/src/migrations/schema/mod.rs
@@ -55,15 +55,15 @@ impl Schema {
         );
 
         tracing::debug!(sql = sql.as_str(), "Running CREATE TABLE SQL query");
-
-        conn.exec(&sql, vec![])
-            .await
-            .map_err(|e| Error::Database(e.to_string()))?;
+        let query_result = conn.exec(&sql, vec![]).await;
 
         conn_lock.replace(conn);
         drop(conn_lock);
 
-        Ok(())
+        match query_result {
+            Ok(_) => Ok(()),
+            Err(e) => Err(Error::Database(e.to_string())),
+        }
     }
 
     /// Drops a table.
@@ -72,16 +72,21 @@ impl Schema {
     ///
     /// Returns an error if the table cannot be dropped, or if a connection to the database cannot be established.
     pub async fn drop(table_name: &str) -> Result<(), Error> {
-        let mut conn = connection::get().await?;
+        let mut conn_lock = MIGRATE_CONN.try_lock().map_err(|_| Error::Lock)?;
+        let mut conn = conn_lock.take().ok_or(Error::Lock)?;
+
         let (sql, bindings) = (format!("DROP TABLE ?"), vec![to_value!(table_name)]);
 
         tracing::debug!(sql = sql.as_str(), bindings = ?bindings, "Running DROP TABLE SQL query");
+        let query_result = conn.exec(&sql, bindings).await;
 
-        conn.exec(&sql, bindings)
-            .await
-            .map_err(|e| Error::Database(e.to_string()))?;
+        conn_lock.replace(conn);
+        drop(conn_lock);
 
-        Ok(())
+        match query_result {
+            Ok(_) => Ok(()),
+            Err(e) => Err(Error::Database(e.to_string())),
+        }
     }
 
     fn get_schema<F>(table_name: String, callback: F) -> Result<(Vec<Column>, Vec<Command>), Error>