delete zero balances

iotaledger · Mar 21, 2024 · 8115acb · 8115acb
1 parent 925d142
commit 8115acb
Showing 1 changed file with 18 additions and 6 deletions.
diff --git a/src/db/mongodb/collections/analytics/address_balance.rs b/src/db/mongodb/collections/analytics/address_balance.rs
@@ -91,12 +91,24 @@ pub struct DistributionStat {
 impl AddressBalanceCollection {
     /// Insert a balance for an address.
     pub async fn insert_balance(&self, address: &Address, balance: u64) -> Result<(), DbError> {
-        self.update_one(
-            doc! { "_id": AddressDto::from(address) },
-            doc! { "$set": { "balance": balance.to_string() } },
-            UpdateOptions::builder().upsert(true).build(),
-        )
-        .await?;
+        if balance == 0 {
+            self.delete_balance(address).await?;
+        } else {
+            self.update_one(
+                doc! { "_id": AddressDto::from(address) },
+                doc! { "$set": { "balance": balance.to_string() } },
+                UpdateOptions::builder().upsert(true).build(),
+            )
+            .await?;
+        }
+        Ok(())
+    }
+
+    /// Delete a balance for an address.
+    pub async fn delete_balance(&self, address: &Address) -> Result<(), DbError> {
+        self.collection
+            .delete_one(doc! { "_id": AddressDto::from(address) }, None)
+            .await?;
         Ok(())
     }