From a8dc065dce8b3083b440b7f68abd280b3aa8a611 Mon Sep 17 00:00:00 2001
From: Alexis Sellier <alexis@radicle.xyz>
Date: Sun, 6 Nov 2022 12:40:19 +0100
Subject: [PATCH] Use database transactions where appropriate

Signed-off-by: Alexis Sellier <alexis@radicle.xyz>
---
 radicle-node/src/address/store.rs | 79 ++++++++++++++++---------------
 radicle-node/src/lib.rs           |  1 +
 radicle-node/src/sql.rs           | 21 ++++++++
 3 files changed, 64 insertions(+), 37 deletions(-)
 create mode 100644 radicle-node/src/sql.rs

diff --git a/radicle-node/src/address/store.rs b/radicle-node/src/address/store.rs
index 97307269..48ba31e5 100644
--- a/radicle-node/src/address/store.rs
+++ b/radicle-node/src/address/store.rs
@@ -11,6 +11,7 @@ use crate::address::{KnownAddress, Source};
 use crate::clock::Timestamp;
 use crate::prelude::Address;
 use crate::service::NodeId;
+use crate::sql::transaction;
 use crate::wire::message::AddressType;
 
 #[derive(Error, Debug)]
@@ -120,52 +121,56 @@ impl Store for Book {
         timestamp: Timestamp,
         addrs: impl IntoIterator<Item = KnownAddress>,
     ) -> Result<bool, Error> {
-        let mut stmt = self.db.prepare(
-            "INSERT INTO nodes (id, features, alias, timestamp)
-             VALUES (?1, ?2, ?3, ?4)
-             ON CONFLICT DO UPDATE
-             SET features = ?2, alias = ?3, timestamp = ?4
-             WHERE timestamp < ?4",
-        )?;
-
-        stmt.bind(1, node)?;
-        stmt.bind(2, features)?;
-        stmt.bind(3, alias)?;
-        stmt.bind(4, timestamp as i64)?;
-        stmt.next()?;
-
-        for addr in addrs {
-            let mut stmt = self.db.prepare(
-                "INSERT INTO addresses (node, type, value, source, timestamp)
-                 VALUES (?1, ?2, ?3, ?4, ?5)
+        transaction(&self.db, move |db| {
+            let mut stmt = db.prepare(
+                "INSERT INTO nodes (id, features, alias, timestamp)
+                 VALUES (?1, ?2, ?3, ?4)
                  ON CONFLICT DO UPDATE
-                 SET timestamp = ?5
-                 WHERE timestamp < ?5",
+                 SET features = ?2, alias = ?3, timestamp = ?4
+                 WHERE timestamp < ?4",
             )?;
+
             stmt.bind(1, node)?;
-            stmt.bind(2, AddressType::from(&addr.addr))?;
-            stmt.bind(3, addr.addr)?;
-            stmt.bind(4, addr.source)?;
-            stmt.bind(5, timestamp as i64)?;
+            stmt.bind(2, features)?;
+            stmt.bind(3, alias)?;
+            stmt.bind(4, timestamp as i64)?;
             stmt.next()?;
-        }
-        Ok(self.db.change_count() > 0)
+
+            for addr in addrs {
+                let mut stmt = db.prepare(
+                    "INSERT INTO addresses (node, type, value, source, timestamp)
+                     VALUES (?1, ?2, ?3, ?4, ?5)
+                     ON CONFLICT DO UPDATE
+                     SET timestamp = ?5
+                     WHERE timestamp < ?5",
+                )?;
+                stmt.bind(1, node)?;
+                stmt.bind(2, AddressType::from(&addr.addr))?;
+                stmt.bind(3, addr.addr)?;
+                stmt.bind(4, addr.source)?;
+                stmt.bind(5, timestamp as i64)?;
+                stmt.next()?;
+            }
+            Ok(db.change_count() > 0)
+        })
+        .map_err(Error::from)
     }
 
     fn remove(&mut self, node: &NodeId) -> Result<bool, Error> {
-        self.db
-            .prepare("DELETE FROM nodes WHERE id = ?")?
-            .into_cursor()
-            .bind(&[(*node).into()])?
-            .next();
+        transaction(&self.db, move |db| {
+            db.prepare("DELETE FROM nodes WHERE id = ?")?
+                .into_cursor()
+                .bind(&[(*node).into()])?
+                .next();
 
-        self.db
-            .prepare("DELETE FROM addresses WHERE node = ?")?
-            .into_cursor()
-            .bind(&[(*node).into()])?
-            .next();
+            db.prepare("DELETE FROM addresses WHERE node = ?")?
+                .into_cursor()
+                .bind(&[(*node).into()])?
+                .next();
 
-        Ok(self.db.change_count() > 0)
+            Ok(db.change_count() > 0)
+        })
+        .map_err(Error::from)
     }
 
     fn entries(&self) -> Result<Box<dyn Iterator<Item = (NodeId, KnownAddress)>>, Error> {
diff --git a/radicle-node/src/lib.rs b/radicle-node/src/lib.rs
index 1523734a..ac0e4982 100644
--- a/radicle-node/src/lib.rs
+++ b/radicle-node/src/lib.rs
@@ -5,6 +5,7 @@ pub mod control;
 pub mod decoder;
 pub mod logger;
 pub mod service;
+pub mod sql;
 #[cfg(test)]
 pub mod test;
 pub mod transport;
diff --git a/radicle-node/src/sql.rs b/radicle-node/src/sql.rs
new file mode 100644
index 00000000..81d80850
--- /dev/null
+++ b/radicle-node/src/sql.rs
@@ -0,0 +1,21 @@
+use sqlite as sql;
+
+/// Run an SQL query inside a transaction.
+/// Commits the transaction on success, and rolls back on error.
+pub fn transaction<T>(
+    db: &sql::Connection,
+    query: impl FnOnce(&sql::Connection) -> Result<T, sql::Error>,
+) -> Result<T, sql::Error> {
+    db.execute("BEGIN")?;
+
+    match query(db) {
+        Ok(result) => {
+            db.execute("COMMIT")?;
+            Ok(result)
+        }
+        Err(err) => {
+            db.execute("ROLLBACK")?;
+            Err(err)
+        }
+    }
+}