Pagination queries for balances endpoint

CHANGELOG.md

@@ -31,6 +31,7 @@ and this project adheres to [Semantic Versioning](http://semver.org/).
 - [2447](https://github.com/FuelLabs/fuel-core/pull/2447): Use new `expiration` policy in the transaction pool. Add a mechanism to prune the transactions when they expired.
 - [1922](https://github.com/FuelLabs/fuel-core/pull/1922): Added support for posting blocks to the shared sequencer.
 - [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.
@@ -67,6 +68,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.
 - [2463](https://github.com/FuelLabs/fuel-core/pull/2463): 'CoinsQueryError::MaxCoinsReached` variant has been removed. The `InsufficientCoins` variant has been renamed to `InsufficientCoinsForTheMax` and it now contains the additional `max` field
 - [2463](https://github.com/FuelLabs/fuel-core/pull/2463): The number of excluded ids in the `coinsToSpend` GraphQL query is now limited to the maximum number of inputs allowed in transaction.
 - [2463](https://github.com/FuelLabs/fuel-core/pull/2463): The `coinsToSpend` GraphQL query may now return different coins, depending whether the indexation is enabled or not. However, regardless of the differences, the returned coins will accurately reflect the current state of the database within the context of the query.

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

@@ -88,12 +88,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_indexation_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_indexation_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

@@ -248,55 +248,43 @@ 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)>> {
-        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 {
-            Ok(base_asset_balance) => {
-                if base_asset_balance != 0 {
-                    iter::once(Ok((base_asset_id, base_asset_balance))).into_boxed()
+    ) -> 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 {
-                    iter::empty().into_boxed()
+                    let start = CoinBalancesKey::new(owner, &asset_id);
+                    self.non_base_asset_balances(
+                        owner,
+                        Some(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)
                 }
             }
-            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))
-                })
-            })
-            .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()
         }
     }
+
     // TODO: Return error if indexation is not available: https://github.com/FuelLabs/fuel-core/issues/2499
     fn coins_to_spend_index(
         &self,
@@ -332,6 +320,88 @@ impl OffChainDatabase for OffChainIterableKeyValueView {
     }
 }
 
+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);
+        match base_balance {
+            Ok(base_asset_balance) => {
+                if base_asset_balance != 0 {
+                    iter::once(Ok((base_asset_id, base_asset_balance))).into_boxed()
+                } else {
+                    iter::empty().into_boxed()
+                }
+            }
+            Err(err) => iter::once(Err(err)).into_boxed(),
+        }
+    }
+
+    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()
+    }
+
+    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()
+    }
+}
+
 impl worker::OffChainDatabase for Database<OffChain> {
     type Transaction<'a>
         = StorageTransaction<&'a mut Self>

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() {