Pagination queries for balances endpoint

CHANGELOG.md

@@ -28,6 +28,7 @@ and this project adheres to [Semantic Versioning](http://semver.org/).
 - [2327](https://github.com/FuelLabs/fuel-core/pull/2327): Add more services tests and more checks of the pool. Also add an high level documentation for users of the pool and contributors.
 - [2416](https://github.com/FuelLabs/fuel-core/issues/2416): Define the `GasPriceServiceV1` task.
 - [2033](https://github.com/FuelLabs/fuel-core/pull/2033): Remove `Option<BlockHeight>` in favor of `BlockHeightQuery` where applicable.
+- [2490](https://github.com/FuelLabs/fuel-core/pull/2490): Added pagination support for the `balances` GraphQL query, available only when 'balances indexation' is enabled.
 - [2439](https://github.com/FuelLabs/fuel-core/pull/2439): Add gas costs for the two new zk opcodes `ecop` and `eadd` and the benches that allow to calibrate them.
 - [2472](https://github.com/FuelLabs/fuel-core/pull/2472): Added the `amountU128` field to the `Balance` GraphQL schema, providing the total balance as a `U128`. The existing `amount` field clamps any balance exceeding `U64` to `u64::MAX`.
 - [2526](https://github.com/FuelLabs/fuel-core/pull/2526): Add possibility to not have any cache set for RocksDB. Add an option to either load the RocksDB columns families on creation of the database or when the column is used.
@@ -61,6 +62,7 @@ and this project adheres to [Semantic Versioning](http://semver.org/).
 - [2154](https://github.com/FuelLabs/fuel-core/pull/2154): Transaction graphql endpoints use `TransactionType` instead of `fuel_tx::Transaction`.
 - [2446](https://github.com/FuelLabs/fuel-core/pull/2446): Use graphiql instead of graphql-playground due to known vulnerability and stale development.
 - [2379](https://github.com/FuelLabs/fuel-core/issues/2379): Change `kv_store::Value` to be `Arc<[u8]>` instead of `Arc<Vec<u8>>`.
+- [2490](https://github.com/FuelLabs/fuel-core/pull/2490): Updated GraphQL complexity calculation for `balances` query to account for pagination (`first`/`last`) and nested field complexity (`child_complexity`). Queries with large pagination values or deeply nested fields may have higher complexity costs.
 - [2526](https://github.com/FuelLabs/fuel-core/pull/2526): By default the cache of RocksDB is now disabled instead of being `1024 * 1024 * 1024`.
 
 ## [Version 0.40.2]

crates/chain-config/src/config/chain.rs

@@ -123,6 +123,15 @@ impl ChainConfig {
             ..Default::default()
         }
     }
+
+    #[cfg(feature = "test-helpers")]
+    pub fn local_testnet_with_consensus_parameters(cp: &ConsensusParameters) -> Self {
+        Self {
+            chain_name: LOCAL_TESTNET.to_string(),
+            consensus_parameters: cp.clone(),
+            ..Default::default()
+        }
+    }
 }
 
 impl GenesisCommitment for ChainConfig {

crates/fuel-core/src/graphql_api/ports.rs

@@ -80,12 +80,13 @@ pub trait OffChainDatabase: Send + Sync {
         base_asset_id: &AssetId,
     ) -> StorageResult<TotalBalanceAmount>;
 
-    fn balances(
-        &self,
+    fn balances<'a>(
+        &'a self,
         owner: &Address,
-        base_asset_id: &AssetId,
+        start: Option<AssetId>,
+        base_asset_id: &'a AssetId,
         direction: IterDirection,
-    ) -> BoxedIter<'_, StorageResult<(AssetId, TotalBalanceAmount)>>;
+    ) -> BoxedIter<'a, StorageResult<(AssetId, TotalBalanceAmount)>>;
 
     fn owned_coins_ids(
         &self,

crates/fuel-core/src/query/balance.rs

@@ -69,12 +69,14 @@ impl ReadView {
     pub fn balances<'a>(
         &'a self,
         owner: &'a Address,
+        start: Option<AssetId>,
         direction: IterDirection,
         base_asset_id: &'a AssetId,
     ) -> impl Stream<Item = StorageResult<AddressBalance>> + 'a {
         if self.balances_enabled {
             futures::future::Either::Left(self.balances_with_cache(
                 owner,
+                start,
                 base_asset_id,
                 direction,
             ))
@@ -140,17 +142,21 @@ impl ReadView {
     fn balances_with_cache<'a>(
         &'a self,
         owner: &'a Address,
-        base_asset_id: &AssetId,
+        start: Option<AssetId>,
+        base_asset_id: &'a AssetId,
         direction: IterDirection,
     ) -> impl Stream<Item = StorageResult<AddressBalance>> + 'a {
-        stream::iter(self.off_chain.balances(owner, base_asset_id, direction))
-            .map(move |result| {
-                result.map(|(asset_id, amount)| AddressBalance {
-                    owner: *owner,
-                    asset_id,
-                    amount,
-                })
+        stream::iter(
+            self.off_chain
+                .balances(owner, start, base_asset_id, direction),
+        )
+        .map(move |result| {
+            result.map(|(asset_id, amount)| AddressBalance {
+                owner: *owner,
+                asset_id,
+                amount,
             })
-            .yield_each(self.batch_size)
+        })
+        .yield_each(self.batch_size)
     }
 }

crates/fuel-core/src/schema/balance.rs

@@ -79,9 +79,11 @@ impl BalanceQuery {
         Ok(balance)
     }
 
-    // TODO: This API should be migrated to the indexer for better support and
-    //  discontinued within fuel-core.
-    #[graphql(complexity = "query_costs().balance_query")]
+    #[graphql(
+        complexity = "query_costs().balance_query + query_costs().storage_iterator \
+        + (query_costs().storage_read + first.unwrap_or_default() as usize) * child_complexity \
+        + (query_costs().storage_read + last.unwrap_or_default() as usize) * child_complexity"
+    )]
     async fn balances(
         &self,
         ctx: &Context<'_>,
@@ -92,18 +94,21 @@ impl BalanceQuery {
         before: Option<String>,
     ) -> async_graphql::Result<Connection<AssetId, Balance, EmptyFields, EmptyFields>>
     {
-        if before.is_some() || after.is_some() {
-            return Err(anyhow!("pagination is not yet supported").into())
-        }
         let query = ctx.read_view()?;
+        if !query.balances_enabled && (before.is_some() || after.is_some()) {
+            return Err(anyhow!(
+                "Can not use pagination when balances indexation is not available"
+            )
+            .into())
+        }
         let base_asset_id = *ctx
             .data_unchecked::<ConsensusProvider>()
             .latest_consensus_params()
             .base_asset_id();
         let owner = filter.owner.into();
-        crate::schema::query_pagination(after, before, first, last, |_, direction| {
+        crate::schema::query_pagination(after, before, first, last, |start, direction| {
             Ok(query
-                .balances(&owner, direction, &base_asset_id)
+                .balances(&owner, (*start).map(Into::into), direction, &base_asset_id)
                 .map(|result| {
                     result.map(|balance| (balance.asset_id.into(), balance.into()))
                 }))

crates/fuel-core/src/service/adapters/graphql_api/off_chain.rs

@@ -233,15 +233,49 @@ impl OffChainDatabase for OffChainIterableKeyValueView {
         }
     }
 
-    fn balances(
-        &self,
+    fn balances<'a>(
+        &'a self,
         owner: &Address,
-        base_asset_id: &AssetId,
+        start: Option<AssetId>,
+        base_asset_id: &'a AssetId,
         direction: IterDirection,
-    ) -> BoxedIter<'_, StorageResult<(AssetId, TotalBalanceAmount)>> {
+    ) -> BoxedIter<'a, StorageResult<(AssetId, TotalBalanceAmount)>> {
+        match (direction, start) {
+            (IterDirection::Forward, None) => {
+                self.base_asset_first(owner, base_asset_id, direction)
+            }
+            (IterDirection::Forward, Some(asset_id)) => {
+                if asset_id == *base_asset_id {
+                    self.base_asset_first(owner, base_asset_id, direction)
+                } else {
+                    let start = (asset_id != *base_asset_id)
+                        .then_some(CoinBalancesKey::new(owner, &asset_id));
+                    self.non_base_asset_balances(owner, start, direction, base_asset_id)
+                }
+            }
+            (IterDirection::Reverse, None) => {
+                self.non_base_assets_first(start, owner, base_asset_id, direction)
+            }
+            (IterDirection::Reverse, Some(asset_id)) => {
+                if asset_id == *base_asset_id {
+                    self.base_asset_balance(base_asset_id, owner)
+                } else {
+                    self.non_base_assets_first(start, owner, base_asset_id, direction)
+                }
+            }
+        }
+    }
+}
+
+impl OffChainIterableKeyValueView {
+    fn base_asset_balance(
+        &self,
+        base_asset_id: &AssetId,
+        owner: &Address,
+    ) -> BoxedIter<'_, Result<(AssetId, u128), StorageError>> {
         let base_asset_id = *base_asset_id;
         let base_balance = self.balance(owner, &base_asset_id, &base_asset_id);
-        let base_asset_balance = match base_balance {
+        match base_balance {
             Ok(base_asset_balance) => {
                 if base_asset_balance != 0 {
                     iter::once(Ok((base_asset_id, base_asset_balance))).into_boxed()
@@ -250,37 +284,68 @@ impl OffChainDatabase for OffChainIterableKeyValueView {
                 }
             }
             Err(err) => iter::once(Err(err)).into_boxed(),
-        };
+        }
+    }
 
-        let non_base_asset_balance = self
-            .iter_all_filtered_keys::<CoinBalances, _>(Some(owner), None, Some(direction))
-            .filter_map(move |result| match result {
-                Ok(key) if *key.asset_id() != base_asset_id => Some(Ok(key)),
-                Ok(_) => None,
-                Err(err) => Some(Err(err)),
-            })
-            .map(move |result| {
-                result.and_then(|key| {
-                    let asset_id = key.asset_id();
-                    let coin_balance =
-                        self.storage_as_ref::<CoinBalances>()
-                            .get(&key)?
-                            .unwrap_or_default()
-                            .into_owned() as TotalBalanceAmount;
-                    Ok((*asset_id, coin_balance))
-                })
+    fn non_base_asset_balances<'a>(
+        &'a self,
+        owner: &Address,
+        start: Option<CoinBalancesKey>,
+        direction: IterDirection,
+        base_asset_id: &'a AssetId,
+    ) -> BoxedIter<'_, Result<(AssetId, u128), StorageError>> {
+        self.iter_all_filtered_keys::<CoinBalances, _>(
+            Some(owner),
+            start.as_ref(),
+            Some(direction),
+        )
+        .filter_map(move |result| match result {
+            Ok(key) if *key.asset_id() != *base_asset_id => Some(Ok(key)),
+            Ok(_) => None,
+            Err(err) => Some(Err(err)),
+        })
+        .map(move |result| {
+            result.and_then(|key| {
+                let asset_id = key.asset_id();
+                let coin_balance =
+                    self.storage_as_ref::<CoinBalances>()
+                        .get(&key)?
+                        .unwrap_or_default()
+                        .into_owned() as TotalBalanceAmount;
+                Ok((*asset_id, coin_balance))
             })
-            .into_boxed();
+        })
+        .into_boxed()
+    }
 
-        if direction == IterDirection::Forward {
-            base_asset_balance
-                .chain(non_base_asset_balance)
-                .into_boxed()
-        } else {
-            non_base_asset_balance
-                .chain(base_asset_balance)
-                .into_boxed()
-        }
+    fn non_base_assets_first<'a>(
+        &'a self,
+        start: Option<AssetId>,
+        owner: &Address,
+        base_asset_id: &'a AssetId,
+        direction: IterDirection,
+    ) -> BoxedIter<'_, Result<(AssetId, u128), StorageError>> {
+        let start = start.map(|asset_id| CoinBalancesKey::new(owner, &asset_id));
+        let base_asset_balance = self.base_asset_balance(base_asset_id, owner);
+        let non_base_asset_balance =
+            self.non_base_asset_balances(owner, start, direction, base_asset_id);
+        non_base_asset_balance
+            .chain(base_asset_balance)
+            .into_boxed()
+    }
+
+    fn base_asset_first<'a>(
+        &'a self,
+        owner: &Address,
+        base_asset_id: &'a AssetId,
+        direction: IterDirection,
+    ) -> BoxedIter<'_, Result<(AssetId, u128), StorageError>> {
+        let base_asset_balance = self.base_asset_balance(base_asset_id, owner);
+        let non_base_asset_balances =
+            self.non_base_asset_balances(owner, None, direction, base_asset_id);
+        base_asset_balance
+            .chain(non_base_asset_balances)
+            .into_boxed()
     }
 }
 

crates/services/gas_price_service/src/common/utils.rs

@@ -17,7 +17,7 @@ pub enum Error {
     },
     #[error("Failed to initialize updater: {0:?}")]
     CouldNotInitUpdater(anyhow::Error),
-    #[error("Failed to convert metadata to concrete type. THere is no migration path for this metadata version")]
+    #[error("Failed to convert metadata to concrete type. There is no migration path for this metadata version")]
     CouldNotConvertMetadata, // todo(https://github.com/FuelLabs/fuel-core/issues/2286)
     #[error("Failed to commit to storage: {0:?}")]
     CouldNotCommit(anyhow::Error),

crates/services/src/service.rs

@@ -427,7 +427,7 @@ where
     match shutdown.catch_unwind().await {
         Ok(Ok(_)) => {}
         Ok(Err(e)) => {
-            tracing::error!("Go an error during shutdown of the task: {e}");
+            tracing::error!("Got an error during shutdown of the task: {e}");
         }
         Err(e) => {
             if got_panic.is_some() {