From a1288dc16397487abb80851e0280bacf459d8fd3 Mon Sep 17 00:00:00 2001
From: librelois <c@elo.tf>
Date: Tue, 9 Mar 2021 01:36:03 +0100
Subject: [PATCH] [feat] gva: add query balances

---
 .../gva/gql/src/queries/account_balance.rs    | 78 ++++++++++++++++++-
 1 file changed, 77 insertions(+), 1 deletion(-)

diff --git a/rust-libs/modules/gva/gql/src/queries/account_balance.rs b/rust-libs/modules/gva/gql/src/queries/account_balance.rs
index 6639ac107..2e4e58a9e 100644
--- a/rust-libs/modules/gva/gql/src/queries/account_balance.rs
+++ b/rust-libs/modules/gva/gql/src/queries/account_balance.rs
@@ -19,7 +19,9 @@ use crate::*;
 pub(crate) struct AccountBalanceQuery;
 #[async_graphql::Object]
 impl AccountBalanceQuery {
-    /// Account balance
+    /// Get the balance of an account identified by a public key or a script.
+    ///
+    /// If the balance is null, the account has never been used. If the balance is zero, the account has already transited money in the past.
     async fn balance(
         &self,
         ctx: &async_graphql::Context<'_>,
@@ -43,6 +45,49 @@ impl AccountBalanceQuery {
                 base: balance.0.base() as i32,
             }))
     }
+    /// Get the balance of several accounts in a single request
+    ///
+    /// Each account can be identified by a public key or a script. It is possible to mix the two in the same request.
+    /// The balances are returned in the order of the accounts provided as input. Each account has a balance,
+    /// which is null if the account does not exist.
+    async fn balances(
+        &self,
+        ctx: &async_graphql::Context<'_>,
+        #[graphql(desc = "Accounts scripts or publics keys")] scripts: Vec<String>,
+    ) -> async_graphql::Result<Vec<Option<AmountWithBase>>> {
+        let accounts_scripts = scripts
+            .into_iter()
+            .map(|script| {
+                if let Ok(pubkey) = PublicKey::from_base58(&script) {
+                    Ok(WalletScriptV10::single_sig(pubkey))
+                } else {
+                    dubp::documents_parser::wallet_script_from_str(&script)
+                }
+            })
+            .collect::<Result<Vec<_>, _>>()?;
+
+        let data = ctx.data::<GvaSchemaData>()?;
+        let dbs_reader = data.dbs_reader();
+
+        Ok(data
+            .dbs_pool
+            .execute(move |_| {
+                accounts_scripts
+                    .iter()
+                    .map(|account_script| {
+                        dbs_reader
+                            .get_account_balance(account_script)
+                            .map(|balance_opt| {
+                                balance_opt.map(|balance| AmountWithBase {
+                                    amount: balance.0.amount() as i32,
+                                    base: balance.0.base() as i32,
+                                })
+                            })
+                    })
+                    .collect::<Result<Vec<_>, _>>()
+            })
+            .await??)
+    }
 }
 
 #[cfg(test)]
@@ -81,4 +126,35 @@ mod tests {
         );
         Ok(())
     }
+
+    #[tokio::test]
+    async fn query_balances() -> anyhow::Result<()> {
+        let mut dbs_reader = MockDbsReader::new();
+        dbs_reader
+            .expect_get_account_balance()
+            .withf(|s| {
+                s == &WalletScriptV10::single_sig(
+                    PublicKey::from_base58("DnjL6hYA1k7FavGHbbir79PKQbmzw63d6bsamBBdUULP")
+                        .expect("wrong pubkey"),
+                )
+            })
+            .times(1)
+            .returning(|_| Ok(Some(SourceAmountValV2(SourceAmount::with_base0(38)))));
+        let schema = create_schema(dbs_reader)?;
+        assert_eq!(
+            exec_graphql_request(
+                &schema,
+                r#"{ balances(scripts: ["DnjL6hYA1k7FavGHbbir79PKQbmzw63d6bsamBBdUULP"]) {amount} }"#
+            )
+            .await?,
+            serde_json::json!({
+                "data": {
+                    "balances": [{
+                      "amount": 38
+                    }]
+                }
+            })
+        );
+        Ok(())
+    }
 }
-- 
GitLab