Skip to content

Commit

Permalink
Iris DB methods
Browse files Browse the repository at this point in the history
  • Loading branch information
Dzejkop committed Mar 7, 2024
1 parent 2f8274b commit b7e04a4
Show file tree
Hide file tree
Showing 2 changed files with 77 additions and 6 deletions.
2 changes: 1 addition & 1 deletion bin/migrate-codes/migrate.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 =
Expand Down
81 changes: 76 additions & 5 deletions src/iris_db.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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")]
Expand Down Expand Up @@ -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> {
Expand All @@ -71,14 +121,32 @@ 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?;

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,
Expand All @@ -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?;
Expand Down

0 comments on commit b7e04a4

Please sign in to comment.