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, + /// 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 = + 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> { + let collection: Collection = + 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 = + 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 { @@ -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> { + let collection: Collection = + 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?;