From a88dbd99afde2b0d287575b6fbd27b24132c48e1 Mon Sep 17 00:00:00 2001
From: Eric Woolsey <ewoolsey@ualberta.ca>
Date: Tue, 13 Feb 2024 11:34:18 -0800
Subject: [PATCH] Delete recoveries PRO-534 (#691)

* delete-recoveries
---
 src/database/mod.rs                           | 28 +++++++++--------
 src/task_monitor/tasks/finalize_identities.rs | 30 ++++++++-----------
 2 files changed, 28 insertions(+), 30 deletions(-)

diff --git a/src/database/mod.rs b/src/database/mod.rs
index c198a69e..8126e8bc 100644
--- a/src/database/mod.rs
+++ b/src/database/mod.rs
@@ -671,7 +671,7 @@ pub trait DatabaseExt<'a>: Executor<'a, Database = Postgres> {
         Ok(())
     }
 
-    async fn get_recoveries(self) -> Result<Vec<RecoveryEntry>, Error> {
+    async fn get_all_recoveries(self) -> Result<Vec<RecoveryEntry>, Error> {
         Ok(
             sqlx::query_as::<_, RecoveryEntry>("SELECT * FROM recoveries")
                 .fetch_all(self)
@@ -679,21 +679,22 @@ pub trait DatabaseExt<'a>: Executor<'a, Database = Postgres> {
         )
     }
 
-    async fn find_recoveries_by_prev_commit(
+    async fn delete_recoveries<I: IntoIterator<Item = T>, T: Into<U256>>(
         self,
-        prev_commits: &[U256],
+        prev_commits: I,
     ) -> Result<Vec<RecoveryEntry>, Error> {
         // TODO: upstream PgHasArrayType impl to ruint
         let prev_commits = prev_commits
-            .iter()
-            .map(|c| c.to_be_bytes())
+            .into_iter()
+            .map(|c| c.into().to_be_bytes())
             .collect::<Vec<[u8; 32]>>();
 
         let res = sqlx::query_as::<_, RecoveryEntry>(
             r#"
-            SELECT *
+            DELETE
             FROM recoveries
             WHERE existing_commitment = ANY($1)
+            RETURNING *
             "#,
         )
         .bind(&prev_commits)
@@ -1099,7 +1100,7 @@ mod test {
         db.insert_new_recovery(&existing_commitment, &new_commitment)
             .await?;
 
-        let recoveries = db.get_recoveries().await?;
+        let recoveries = db.get_all_recoveries().await?;
 
         assert_eq!(recoveries.len(), 1);
         assert_eq!(recoveries[0].existing_commitment, existing_commitment);
@@ -1287,14 +1288,14 @@ mod test {
             db.insert_new_recovery(&old, &new).await?;
         }
 
-        let recoveries = db.get_recoveries().await?;
+        let recoveries = db.get_all_recoveries().await?;
         assert_eq!(recoveries.len(), 3);
 
         Ok(())
     }
 
     #[tokio::test]
-    async fn test_find_recoveries_from_prev() -> anyhow::Result<()> {
+    async fn test_delete_recoveries() -> anyhow::Result<()> {
         let (db, _db_container) = setup_db().await?;
 
         let old_identities = mock_identities(3);
@@ -1304,10 +1305,13 @@ mod test {
             db.insert_new_recovery(&old, &new).await?;
         }
 
-        let recoveries = db
-            .find_recoveries_by_prev_commit(&old_identities[0..2])
+        let deleted_recoveries = db
+            .delete_recoveries(old_identities[0..2].iter().cloned())
             .await?;
-        assert_eq!(recoveries.len(), 2);
+        assert_eq!(deleted_recoveries.len(), 2);
+
+        let remaining = db.get_all_recoveries().await?;
+        assert_eq!(remaining.len(), 1);
 
         Ok(())
     }
diff --git a/src/task_monitor/tasks/finalize_identities.rs b/src/task_monitor/tasks/finalize_identities.rs
index 89828a21..50b34f1b 100644
--- a/src/task_monitor/tasks/finalize_identities.rs
+++ b/src/task_monitor/tasks/finalize_identities.rs
@@ -245,8 +245,6 @@ fn extract_roots_from_secondary_logs(logs: &[Log]) -> Vec<U256> {
     roots
 }
 
-use crate::identity_tree::Hash;
-
 async fn update_eligible_recoveries(
     database: &Database,
     identity_manager: &IdentityManager,
@@ -266,16 +264,6 @@ async fn update_eligible_recoveries(
         .map(std::convert::Into::into)
         .collect();
 
-    // Check if any deleted commitments correspond with entries in the
-    // recoveries table and insert the new commitment into the unprocessed
-    // identities table with the proper eligibility timestamp
-    let recoveries = database
-        .get_recoveries()
-        .await?
-        .iter()
-        .map(|f| (f.existing_commitment, f.new_commitment))
-        .collect::<HashMap<Hash, Hash>>();
-
     // Fetch the root history expiry time on chain
     let root_history_expiry = identity_manager.root_history_expiry().await?;
 
@@ -289,15 +277,21 @@ async fn update_eligible_recoveries(
 
     let eligibility_timestamp = Utc::now() + delay;
 
+    let mut tx = database.begin().await?;
+
+    // Check if any deleted commitments correspond with entries in the
+    // recoveries table and insert the new commitment into the unprocessed
+    // identities table with the proper eligibility timestamp
+    let deleted_recoveries = tx.delete_recoveries(commitments).await?;
+
     // For each deletion, if there is a corresponding recovery, insert a new
     // identity with the specified eligibility timestamp
-    for prev_commitment in commitments {
-        if let Some(new_commitment) = recoveries.get(&prev_commitment.into()) {
-            database
-                .insert_new_identity(*new_commitment, eligibility_timestamp)
-                .await?;
-        }
+    for recovery in deleted_recoveries {
+        tx.insert_new_identity(recovery.new_commitment, eligibility_timestamp)
+            .await?;
     }
 
+    tx.commit().await?;
+
     Ok(())
 }