From cd3e33e9a03a1f29eb15032db425e72be72156c5 Mon Sep 17 00:00:00 2001
From: Vladimir Fomene <vladimirfomene@gmail.com>
Date: Thu, 1 Sep 2022 13:27:09 +0300
Subject: [PATCH] Sqlite migrations should either succeed or fail

The current implementation of the `migrate` method for
Sqlite database does not rollback changes when there is
an error while running one of the migration scripts. This
can leave the database in an inconsistent state. This
change ensures that migrations either succeed completely
or fail.
---
 src/database/sqlite.rs | 61 +++++++++++++++++++++++++-----------------
 1 file changed, 36 insertions(+), 25 deletions(-)

diff --git a/src/database/sqlite.rs b/src/database/sqlite.rs
index c875bd947a..189bb1b586 100644
--- a/src/database/sqlite.rs
+++ b/src/database/sqlite.rs
@@ -53,13 +53,11 @@ static MIGRATIONS: &[&str] = &[
     "DELETE FROM utxos;",
     "DROP INDEX idx_txid_vout;",
     "CREATE UNIQUE INDEX idx_utxos_txid_vout ON utxos(txid, vout);",
-    "BEGIN TRANSACTION;\
-    ALTER TABLE utxos RENAME TO utxos_old;\
-    CREATE TABLE utxos (value INTEGER, keychain TEXT, vout INTEGER, txid BLOB, script BLOB, is_spent BOOLEAN DEFAULT 0);\
-    INSERT INTO utxos SELECT value, keychain, vout, txid, script, is_spent FROM utxos_old;\
-    DROP TABLE utxos_old;\
-    CREATE UNIQUE INDEX idx_utxos_txid_vout ON utxos(txid, vout);\
-    COMMIT;"
+    "ALTER TABLE utxos RENAME TO utxos_old;",
+    "CREATE TABLE utxos (value INTEGER, keychain TEXT, vout INTEGER, txid BLOB, script BLOB, is_spent BOOLEAN DEFAULT 0);",
+    "INSERT INTO utxos SELECT value, keychain, vout, txid, script, is_spent FROM utxos_old;",
+    "DROP TABLE utxos_old;",
+    "CREATE UNIQUE INDEX idx_utxos_txid_vout ON utxos(txid, vout);"
 ];
 
 /// Sqlite database stored on filesystem
@@ -921,8 +919,8 @@ impl BatchDatabase for SqliteDatabase {
 }
 
 pub fn get_connection<T: AsRef<Path>>(path: &T) -> Result<Connection, Error> {
-    let connection = Connection::open(path)?;
-    migrate(&connection)?;
+    let mut connection = Connection::open(path)?;
+    migrate(&mut connection)?;
     Ok(connection)
 }
 
@@ -957,28 +955,41 @@ pub fn set_schema_version(conn: &Connection, version: i32) -> rusqlite::Result<u
     )
 }
 
-pub fn migrate(conn: &Connection) -> rusqlite::Result<()> {
+pub fn migrate(conn: &mut Connection) -> Result<(), Error> {
     let version = get_schema_version(conn)?;
     let stmts = &MIGRATIONS[(version as usize)..];
-    let mut i: i32 = version;
 
-    if version == MIGRATIONS.len() as i32 {
+    // begin transaction, all migration statements and new schema version commit or rollback
+    let tx = conn.transaction()?;
+
+    // execute every statement and return `Some` new schema version
+    // if execution fails, return `Error::Rusqlite`
+    // if no statements executed returns `None`
+    let new_version = stmts
+        .iter()
+        .enumerate()
+        .map(|version_stmt| {
+            log::info!(
+                "executing db migration {}: `{}`",
+                version_stmt.0,
+                version_stmt.1
+            );
+            tx.execute(version_stmt.1, [])
+                // map result value to next migration version
+                .map(|_| version_stmt.0 as i32 + 1)
+        })
+        .last()
+        .transpose()?;
+
+    // if `Some` new statement version, set new schema version
+    if let Some(version) = new_version {
+        set_schema_version(&tx, version)?;
+    } else {
         log::info!("db up to date, no migration needed");
-        return Ok(());
     }
 
-    for stmt in stmts {
-        let res = conn.execute(stmt, []);
-        if res.is_err() {
-            println!("migration failed on:\n{}\n{:?}", stmt, res);
-            break;
-        }
-
-        i += 1;
-    }
-
-    set_schema_version(conn, i)?;
-
+    // commit transaction
+    tx.commit()?;
     Ok(())
 }