From dff3361e3d8b3a5aa8de1a668cbf71efc087a33a Mon Sep 17 00:00:00 2001 From: stefan-mysten <135084671+stefan-mysten@users.noreply.github.com> Date: Sun, 29 Sep 2024 09:03:07 -0700 Subject: [PATCH] sui-graphql-rpc: add balance function (#13) --- crates/sui-graphql-client/src/lib.rs | 54 ++++++++++++++++--- .../src/query_types/balance.rs | 31 +++++++++++ .../sui-graphql-client/src/query_types/mod.rs | 4 +- 3 files changed, 82 insertions(+), 7 deletions(-) create mode 100644 crates/sui-graphql-client/src/query_types/balance.rs diff --git a/crates/sui-graphql-client/src/lib.rs b/crates/sui-graphql-client/src/lib.rs index fd7f417ff..99d770902 100644 --- a/crates/sui-graphql-client/src/lib.rs +++ b/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, 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::()) + .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() { diff --git a/crates/sui-graphql-client/src/query_types/balance.rs b/crates/sui-graphql-client/src/query_types/balance.rs new file mode 100644 index 000000000..02dfee4f1 --- /dev/null +++ b/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, +} + +#[derive(cynic::QueryFragment, Debug)] +#[cynic(schema = "rpc", graphql_type = "Owner", variables = "BalanceArgs")] +pub struct Owner { + #[arguments(type: $coin_type)] + pub balance: Option, +} + +#[derive(cynic::QueryFragment, Debug)] +#[cynic(schema = "rpc", graphql_type = "Balance")] +pub struct Balance { + pub total_balance: Option, +} + +#[derive(cynic::QueryVariables, Debug)] +pub struct BalanceArgs { + pub address: SuiAddress, + pub coin_type: Option, +} diff --git a/crates/sui-graphql-client/src/query_types/mod.rs b/crates/sui-graphql-client/src/query_types/mod.rs index e45edfa28..34accfd64 100644 --- a/crates/sui-graphql-client/src/query_types/mod.rs +++ b/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)]