sui-graphql-rpc: add balance function (#13)

Author: stefan-mysten

crates/sui-graphql-client/src/lib.rs

@@ -7,12 +7,13 @@ pub mod query_types;
 
 use base64ct::Encoding;
 use query_types::{
-    ChainIdentifierQuery, CheckpointArgs, CheckpointId, CheckpointQuery, CoinMetadata,
-    CoinMetadataArgs, CoinMetadataQuery, EpochSummaryArgs, EpochSummaryQuery, EventFilter,
-    EventsQuery, EventsQueryArgs, ObjectFilter, ObjectQuery, ObjectQueryArgs, ObjectsQuery,
-    ObjectsQueryArgs, PageInfo, ProtocolConfigQuery, ProtocolConfigs, ProtocolVersionArgs,
-    ServiceConfig, ServiceConfigQuery, TransactionBlockArgs, TransactionBlockQuery,
-    TransactionBlocksQuery, TransactionBlocksQueryArgs, TransactionsFilter, Uint53,
+    BalanceArgs, BalanceQuery, ChainIdentifierQuery, CheckpointArgs, CheckpointId, CheckpointQuery,
+    CoinMetadata, CoinMetadataArgs, CoinMetadataQuery, EpochSummaryArgs, EpochSummaryQuery,
+    EventFilter, EventsQuery, EventsQueryArgs, ObjectFilter, ObjectQuery, ObjectQueryArgs,
+    ObjectsQuery, ObjectsQueryArgs, PageInfo, ProtocolConfigQuery, ProtocolConfigs,
+    ProtocolVersionArgs, ServiceConfig, ServiceConfigQuery, TransactionBlockArgs,
+    TransactionBlockQuery, TransactionBlocksQuery, TransactionBlocksQueryArgs, TransactionsFilter,
+    Uint53,
 };
 use reqwest::Url;
 use sui_types::types::{
@@ -194,6 +195,38 @@ impl Client {
             .ok_or_else(|| Error::msg("No data in response"))
     }
 
+    // ===========================================================================
+    // Balance API
+    // ===========================================================================
+
+    /// Get the balance of all the coins owned by address for the provided coin type.
+    /// Coin type will default to `0x2::coin::Coin<0x2::sui::SUI>` if not provided.
+    pub async fn balance(
+        &self,
+        address: Address,
+        coin_type: Option<&str>,
+    ) -> Result<Option<u128>, Error> {
+        let operation = BalanceQuery::build(BalanceArgs {
+            address: address.into(),
+            coin_type: coin_type.map(|x| x.to_string()),
+        });
+        let response = self.run_query(&operation).await?;
+
+        if let Some(errors) = response.errors {
+            return Err(Error::msg(format!("{:?}", errors)));
+        }
+
+        let total_balance = response
+            .data
+            .map(|b| b.owner.and_then(|o| o.balance.map(|b| b.total_balance)))
+            .ok_or_else(|| Error::msg("No data in response"))?
+            .flatten()
+            .map(|x| x.0.parse::<u128>())
+            .transpose()
+            .map_err(|e| Error::msg(format!("Cannot parse balance into u128: {e}")))?;
+        Ok(total_balance)
+    }
+
     // ===========================================================================
     // Coin API
     // ===========================================================================
@@ -576,6 +609,15 @@ mod tests {
         assert!(client.set_rpc_server("9125/graphql").is_err());
     }
 
+    #[tokio::test]
+    async fn test_balance_query() {
+        for (n, _) in NETWORKS.iter() {
+            let client = Client::new(n).unwrap();
+            let balance = client.balance("0x1".parse().unwrap(), None).await;
+            assert!(balance.is_ok(), "Balance query failed for network: {n}");
+        }
+    }
+
     #[tokio::test]
     async fn test_chain_id() {
         for (n, id) in NETWORKS.iter() {

crates/sui-graphql-client/src/query_types/balance.rs

@@ -0,0 +1,31 @@
+// Copyright (c) Mysten Labs, Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+use super::SuiAddress;
+use crate::query_types::{schema, BigInt};
+
+#[derive(cynic::QueryFragment, Debug)]
+#[cynic(schema = "rpc", graphql_type = "Query", variables = "BalanceArgs")]
+pub struct BalanceQuery {
+    #[arguments(address: $address)]
+    pub owner: Option<Owner>,
+}
+
+#[derive(cynic::QueryFragment, Debug)]
+#[cynic(schema = "rpc", graphql_type = "Owner", variables = "BalanceArgs")]
+pub struct Owner {
+    #[arguments(type: $coin_type)]
+    pub balance: Option<Balance>,
+}
+
+#[derive(cynic::QueryFragment, Debug)]
+#[cynic(schema = "rpc", graphql_type = "Balance")]
+pub struct Balance {
+    pub total_balance: Option<BigInt>,
+}
+
+#[derive(cynic::QueryVariables, Debug)]
+pub struct BalanceArgs {
+    pub address: SuiAddress,
+    pub coin_type: Option<String>,
+}

crates/sui-graphql-client/src/query_types/mod.rs

@@ -2,6 +2,7 @@
 // SPDX-License-Identifier: Apache-2.0
 use std::str::FromStr;
 
+mod balance;
 mod chain;
 mod checkpoint;
 mod coin;
@@ -13,6 +14,7 @@ mod service_config;
 mod transaction;
 
 use anyhow::{anyhow, Error};
+pub use balance::{Balance, BalanceArgs, BalanceQuery, Owner};
 pub use chain::ChainIdentifierQuery;
 pub use checkpoint::{CheckpointArgs, CheckpointId, CheckpointQuery};
 pub use coin::{CoinMetadata, CoinMetadataArgs, CoinMetadataQuery};
@@ -48,7 +50,7 @@ pub struct BigInt(pub String);
 #[cynic(graphql_type = "DateTime")]
 pub struct DateTime(pub String);
 
-#[derive(cynic::Scalar, Debug, Clone)]
+#[derive(cynic::Scalar, Debug)]
 pub struct SuiAddress(pub String);
 
 #[derive(cynic::Scalar, Debug, Clone)]