From b7e04a41d6115aefb990a17c4e068c219822f557 Mon Sep 17 00:00:00 2001
From: Dzejkop <jakubtrad@gmail.com>
Date: Thu, 7 Mar 2024 11:07:17 +0100
Subject: [PATCH] Iris DB methods

---
 bin/migrate-codes/migrate.rs |  2 +-
 src/iris_db.rs               | 81 +++++++++++++++++++++++++++++++++---
 2 files changed, 77 insertions(+), 6 deletions(-)

diff --git a/bin/migrate-codes/migrate.rs b/bin/migrate-codes/migrate.rs
index 9f403fd..b391188 100644
--- a/bin/migrate-codes/migrate.rs
+++ b/bin/migrate-codes/migrate.rs
@@ -76,7 +76,7 @@ async fn main() -> eyre::Result<()> {
 
     let iris_db = IrisDb::new(args.iris_code_db).await?;
 
-    let num_iris_codes = iris_db.count_iris_codes(latest_serial_id).await?;
+    let num_iris_codes = iris_db.count_whitelisted_iris_codes(latest_serial_id).await?;
     tracing::info!("Processing {} iris codes", num_iris_codes);
 
     let pb =
diff --git a/src/iris_db.rs b/src/iris_db.rs
index 3c19917..f513182 100644
--- a/src/iris_db.rs
+++ b/src/iris_db.rs
@@ -14,6 +14,7 @@ pub const FINAL_RESULT_STATUS: &str = "COMPLETED";
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct IrisCodeEntry {
     pub signup_id: String,
+    /// Internal serial id of the iris code db
     pub serial_id: u64,
     pub iris_code_left: Bits,
     pub mask_code_left: Bits,
@@ -24,11 +25,13 @@ pub struct IrisCodeEntry {
 
 #[derive(Serialize, Deserialize, Debug)]
 pub struct FinalResult {
+    /// Should always by "COMPLETED"
     pub status: String,
 
-    #[serde(skip_serializing_if = "Option::is_none")]
+    /// The MPC serial id associated with this signup
     pub serial_id: Option<u64>,
 
+    /// A unique signup id string
     pub signup_id: String,
 
     #[serde(skip_serializing_if = "Option::is_none")]
@@ -62,7 +65,54 @@ impl IrisDb {
     }
 
     #[tracing::instrument(skip(self))]
-    pub async fn count_iris_codes(
+    pub async fn save_final_results(
+        &self,
+        final_results: &[FinalResult],
+    ) -> eyre::Result<()> {
+        let collection: Collection<FinalResult> =
+            self.db.collection(FINAL_RESULT_COLLECTION_NAME);
+
+        collection.insert_many(final_results, None).await?;
+
+        Ok(())
+    }
+
+    #[tracing::instrument(skip(self))]
+    pub async fn get_final_result_by_serial_id(
+        &self,
+        serial_id: u64,
+    ) -> eyre::Result<Option<FinalResult>> {
+        let collection: Collection<FinalResult> =
+            self.db.collection(FINAL_RESULT_COLLECTION_NAME);
+
+        let final_result = collection
+            .find_one(doc! { "serial_id": serial_id as i64 }, None)
+            .await?;
+
+        Ok(final_result)
+    }
+
+    /// Removes all final result entries with serial id larger than the given one
+    #[tracing::instrument(skip(self))]
+    pub async fn prune_final_results(
+        &self,
+        serial_id: u64,
+    ) -> eyre::Result<()> {
+        let collection: Collection<FinalResult> =
+            self.db.collection(FINAL_RESULT_COLLECTION_NAME);
+
+        collection
+            .delete_many(
+                doc! { "serial_id": { "$gt": serial_id as i64 } },
+                None,
+            )
+            .await?;
+
+        Ok(())
+    }
+
+    #[tracing::instrument(skip(self))]
+    pub async fn count_whitelisted_iris_codes(
         &self,
         last_serial_id: u64,
     ) -> eyre::Result<u64> {
@@ -71,7 +121,10 @@ impl IrisDb {
 
         let count = collection
             .count_documents(
-                doc! {"mpc_serial_id": {"$gt": last_serial_id as i64}},
+                doc! {
+                    "serial_id": {"$gt": last_serial_id as i64},
+                    "whitelisted": true,
+                },
                 None,
             )
             .await?;
@@ -79,6 +132,21 @@ impl IrisDb {
         Ok(count)
     }
 
+    #[tracing::instrument(skip(self))]
+    pub async fn get_entry_by_signup_id(
+        &self,
+        signup_id: &str,
+    ) -> eyre::Result<Option<IrisCodeEntry>> {
+        let collection: Collection<IrisCodeEntry> =
+            self.db.collection(COLLECTION_NAME);
+
+        let iris_code_entry = collection
+            .find_one(doc! {"signup_id": signup_id}, None)
+            .await?;
+
+        Ok(iris_code_entry)
+    }
+
     #[tracing::instrument(skip(self))]
     pub async fn stream_whitelisted_iris_codes(
         &self,
@@ -88,14 +156,17 @@ impl IrisDb {
     > {
         let find_options = mongodb::options::FindOptions::builder()
             .batch_size(IRIS_CODE_BATCH_SIZE)
-            .sort(doc! {"mpc_serial_id": 1})
+            .sort(doc! { "serial_id": 1 })
             .build();
 
         let collection = self.db.collection(COLLECTION_NAME);
 
         let cursor = collection
             .find(
-                doc! {"mpc_serial_id": {"$gt": last_serial_id as i64}},
+                doc! {
+                    "serial_id": {"$gt": last_serial_id as i64},
+                    "whitelisted": true
+                },
                 find_options,
             )
             .await?;