feat(api): implement balance/ endpoint (#388)

* Calculate balance of an address

* Impl balance calc for ledger updates

* Cleanup

* Remove internal tagging for Address enum

* Improve address balance query

* Slight refactor

* Match signature locked outputs

* Fix referencing amount in document

* Fix balance query

* Format

* Nits

* Fix syntax

* Add 'address' and 'is_trivial_unlock' to output collection

* Move balances calculation to outputs collection

* Create partial index

* Fix refactor

* Include 'ledger_index' to balances api response

* Address comments

* Remove pub(crate) visibility if not used from collections

* Refactor newly added 'OutputDocument' fields into 'OutputDetails' struct

* Update balance query

* Address comments

bin/inx-chronicle/src/api/stardust/history/responses.rs

@@ -69,3 +69,13 @@ impl From<LedgerUpdateByMilestoneRecord> for LedgerUpdateByMilestoneResponse {
  }
  }
 }
+
+#[derive(Clone, Debug, Serialize, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct BalanceResponse {
+ pub total_balance: u64,
+ pub sig_locked_balance: u64,
+ pub ledger_index: MilestoneIndex,
+}
+
+impl_success_response!(BalanceResponse);

bin/inx-chronicle/src/api/stardust/history/routes.rs

@@ -16,19 +16,22 @@ use super::{
  LedgerUpdatesByMilestonePagination,
  },
  responses::{
- LederUpdatesByAddressResponse, LedgerUpdateByAddressResponse, LedgerUpdateByMilestoneResponse,
+ BalanceResponse, LederUpdatesByAddressResponse, LedgerUpdateByAddressResponse, LedgerUpdateByMilestoneResponse,
  LedgerUpdatesByMilestoneResponse,
  },
 };
 use crate::api::{responses::SyncDataDto, ApiError, ApiResult};
 
 pub fn routes() -> Router {
- Router::new().route("/gaps", get(sync)).nest(
- "/ledger/updates",
- Router::new()
- .route("/by-address/:address", get(ledger_updates_by_address))
- .route("/by-milestone/:milestone_id", get(ledger_updates_by_milestone)),
- )
+ Router::new()
+ .route("/gaps", get(sync))
+ .route("/balance/:address", get(balance))
+ .nest(
+ "/ledger/updates",
+ Router::new()
+ .route("/by-address/:address", get(ledger_updates_by_address))
+ .route("/by-milestone/:milestone_id", get(ledger_updates_by_milestone)),
+ )
 }
 
 async fn sync(database: Extension<MongoDb>) -> ApiResult<SyncDataDto> {
@@ -113,3 +116,17 @@ async fn ledger_updates_by_milestone(
  cursor,
  })
 }
+
+async fn balance(database: Extension<MongoDb>, Path(address): Path<String>) -> ApiResult<BalanceResponse> {
+ let address = Address::from_str(&address).map_err(ApiError::bad_parse)?;
+ let res = database
+ .sum_balances_owned_by_address(address)
+ .await?
+ .ok_or(ApiError::NoResults)?;
+
+ Ok(BalanceResponse {
+ total_balance: res.total_balance,
+ sig_locked_balance: res.sig_locked_balance,
+ ledger_index: res.ledger_index,
+ })
+}

src/db/collections/ledger_update.rs

@@ -105,7 +105,7 @@ impl MongoDb {
  Ok(())
  }
 
- /// Upserts a [`Output`](crate::types::stardust::block::Output) together with its associated
+ /// Upserts an [`Output`](crate::types::stardust::block::Output) together with its associated
  /// [`OutputMetadata`](crate::types::ledger::OutputMetadata).
  pub async fn insert_ledger_updates(
  &self,

src/db/collections/milestone.rs

@@ -26,7 +26,7 @@ use crate::{
 
 /// A milestone's metadata.
 #[derive(Clone, Debug, Serialize, Deserialize)]
-pub(crate) struct MilestoneDocument {
+struct MilestoneDocument {
  /// The milestone index.
  milestone_index: MilestoneIndex,
  /// The [`MilestoneId`](MilestoneId) of the milestone.
@@ -41,7 +41,7 @@ pub(crate) struct MilestoneDocument {
 
 impl MilestoneDocument {
  /// The stardust milestone collection name.
- pub(crate) const COLLECTION: &'static str = "stardust_milestones";
+ const COLLECTION: &'static str = "stardust_milestones";
 }
 
 #[derive(Clone, Debug, Serialize, Deserialize)]

src/db/collections/outputs/mod.rs

@@ -19,7 +19,7 @@ use crate::{
  db::MongoDb,
  types::{
  ledger::{MilestoneIndexTimestamp, OutputMetadata, OutputWithMetadata, SpentMetadata},
- stardust::block::{BlockId, Output, OutputId},
+ stardust::block::{Address, BlockId, Output, OutputId},
  tangle::MilestoneIndex,
  },
 };
@@ -30,19 +30,35 @@ struct OutputDocument {
  output_id: OutputId,
  output: Output,
  metadata: OutputMetadata,
+ details: OutputDetails,
 }
 
 impl OutputDocument {
  /// The stardust outputs collection name.
  const COLLECTION: &'static str = "stardust_outputs";
 }
 
+/// Precalculated info and other output details.
+#[derive(Clone, Debug, Serialize, Deserialize)]
+struct OutputDetails {
+ #[serde(skip_serializing_if = "Option::is_none")]
+ address: Option<Address>,
+ is_trivial_unlock: bool,
+}
+
 impl From<OutputWithMetadata> for OutputDocument {
  fn from(rec: OutputWithMetadata) -> Self {
+ let address = rec.output.owning_address().copied();
+ let is_trivial_unlock = rec.output.is_trivial_unlock();
+
  Self {
  output_id: rec.metadata.output_id,
  output: rec.output,
  metadata: rec.metadata,
+ details: OutputDetails {
+ address,
+ is_trivial_unlock,
+ },
  }
  }
 }
@@ -64,6 +80,14 @@ pub struct OutputWithMetadataResult {
  pub metadata: OutputMetadataResult,
 }
 
+#[derive(Clone, Debug)]
+#[allow(missing_docs)]
+pub struct BalancesResult {
+ pub total_balance: u64,
+ pub sig_locked_balance: u64,
+ pub ledger_index: MilestoneIndex,
+}
+
 /// Implements the queries for the core API.
 impl MongoDb {
  /// Creates output indexes.
@@ -85,6 +109,24 @@ impl MongoDb {
  )
  .await?;
 
+ collection
+ .create_index(
+ IndexModel::builder()
+ .keys(doc! { "details.address": 1 })
+ .options(
+ IndexOptions::builder()
+ .unique(false)
+ .name("address_index".to_string())
+ .partial_filter_expression(doc! {
+ "details.address": { "$exists": true } ,
+ })
+ .build(),
+ )
+ .build(),
+ None,
+ )
+ .await?;
+
  self.create_indexer_output_indexes().await?;
 
  Ok(())
@@ -219,4 +261,70 @@ impl MongoDb {
 
  Ok(metadata)
  }
+
+ /// Sums the amounts of all outputs owned by the given [`Address`](crate::types::stardust::block::Address).
+ pub async fn sum_balances_owned_by_address(&self, address: Address) -> Result<Option<BalancesResult>, Error> {
+ #[derive(Deserialize, Default)]
+ struct Amount {
+ amount: f64,
+ }
+
+ #[derive(Deserialize, Default)]
+ struct Balances {
+ total_balance: Amount,
+ sig_locked_balance: Amount,
+ }
+
+ let ledger_index = self.get_ledger_index().await?;
+ if let Some(ledger_index) = ledger_index {
+ let balances = self
+ .0
+ .collection::<Balances>(OutputDocument::COLLECTION)
+ .aggregate(
+ vec![
+ // Look at all (at ledger index o'clock) unspent output documents for the given address.
+ doc! { "$match": {
+ "details.address": &address,
+ "metadata.booked.milestone_index": { "$lte": ledger_index },
+ "$or": [
+ { "metadata.spent_metadata.spent": null },
+ { "metadata.spent_metadata.spent.milestone_index": { "$gt": ledger_index } },
+ ]
+ } },
+ doc! { "$facet": {
+ // Sum all output amounts (total balance).
+ "total_balance": [
+ { "$group" : {
+ "_id": "null",
+ "amount": { "$sum": { "$toDouble": "$output.amount" } },
+ }},
+ ],
+ // Sum only trivially unlockable output amounts (signature locked balance).
+ "sig_locked_balance": [
+ { "$match": { "details.is_trivial_unlock": true } },
+ { "$group" : {
+ "_id": "null",
+ "amount": { "$sum": { "$toDouble": "$output.amount" } },
+ } },
+ ],
+ } },
+ ],
+ None,
+ )
+ .await?
+ .try_next()
+ .await?
+ .map(bson::from_document::<Balances>)
+ .transpose()?
+ .unwrap_or_default();
+
+ Ok(Some(BalancesResult {
+ total_balance: balances.total_balance.amount as u64,
+ sig_locked_balance: balances.sig_locked_balance.amount as u64,
+ ledger_index,
+ }))
+ } else {
+ Ok(None)
+ }
+ }
 }