[graphql-alt] Load balanceChanges from simulate and execute transaction

Author: tpham-mysten

crates/sui-indexer-alt-e2e-tests/tests/graphql_simulate_transaction_tests.rs

@@ -995,3 +995,84 @@ async fn test_package_resolver_finds_newly_published_package() {
 
     graphql_cluster.stopped().await;
 }
+
+#[tokio::test]
+async fn test_simulate_transaction_balance_changes() {
+    let validator_cluster = TestClusterBuilder::new().build().await;
+    let graphql_cluster = GraphQlTestCluster::new(&validator_cluster).await;
+
+    // Create a transfer transaction that will cause balance changes
+    let recipient = SuiAddress::random_for_testing_only();
+    let transfer_amount = 1_000_000u64;
+
+    let signed_tx = make_transfer_sui_transaction(
+        &validator_cluster.wallet,
+        Some(recipient),
+        Some(transfer_amount),
+    )
+    .await;
+    let (tx_bytes, _signatures) = signed_tx.to_tx_bytes_and_signatures();
+
+    let result = graphql_cluster
+        .execute_graphql(
+            r#"
+            query($txData: JSON!) {
+                simulateTransaction(transaction: $txData) {
+                    effects {
+                        status
+                        balanceChanges {
+                            nodes {
+                                coinType {
+                                    repr
+                                }
+                                amount
+                            }
+                        }
+                    }
+                    error
+                }
+            }
+        "#,
+            json!({
+                "txData": {
+                    "bcs": {
+                        "value": tx_bytes.encoded()
+                    }
+                }
+            }),
+        )
+        .await
+        .expect("GraphQL request failed");
+
+    // Verify balance changes are populated from execution context
+    let balance_changes = result
+        .pointer("/data/simulateTransaction/effects/balanceChanges/nodes")
+        .expect("balanceChanges should be present")
+        .as_array()
+        .unwrap();
+
+    // Should have balance changes for both sender and recipient
+    assert_eq!(balance_changes.len(), 2, "Should have 2 balance changes");
+
+    // Verify structure matches expected format
+    insta::assert_json_snapshot!(result.pointer("/data/simulateTransaction/effects/balanceChanges"), @r#"
+    {
+      "nodes": [
+        {
+          "coinType": {
+            "repr": "0x0000000000000000000000000000000000000000000000000000000000000002::sui::SUI"
+          },
+          "amount": "-3976000"
+        },
+        {
+          "coinType": {
+            "repr": "0x0000000000000000000000000000000000000000000000000000000000000002::sui::SUI"
+          },
+          "amount": "1000000"
+        }
+      ]
+    }
+    "#);
+
+    graphql_cluster.stopped().await;
+}

crates/sui-indexer-alt-graphql/src/api/types/balance_change.rs

@@ -4,48 +4,65 @@
 use std::str::FromStr;
 
 use anyhow::Context as _;
-use async_graphql::Object;
+use async_graphql::SimpleObject;
 use sui_indexer_alt_schema::transactions::BalanceChange as StoredBalanceChange;
-use sui_types::{TypeTag, object::Owner as NativeOwner};
+use sui_rpc::proto::sui::rpc::v2::BalanceChange as GrpcBalanceChange;
+use sui_types::{TypeTag, object::Owner};
 
 use crate::{api::scalars::big_int::BigInt, error::RpcError, scope::Scope};
 
 use super::{address::Address, move_type::MoveType};
 
-#[derive(Clone)]
+/// Effects to the balance (sum of coin values per coin type) of addresses and objects.
+#[derive(Clone, SimpleObject)]
 pub(crate) struct BalanceChange {
-    pub(crate) scope: Scope,
-    pub(crate) stored: StoredBalanceChange,
+    /// The address or object whose balance has changed.
+    pub(crate) owner: Option<Address>,
+    /// The inner type of the coin whose balance has changed (e.g. `0x2::sui::SUI`).
+    pub(crate) coin_type: Option<MoveType>,
+    /// The signed balance change.
+    pub(crate) amount: Option<BigInt>,
 }
 
-/// Effects to the balance (sum of coin values per coin type) of addresses and objects.
-#[Object]
 impl BalanceChange {
-    /// The address or object whose balance has changed.
-    async fn owner(&self) -> Option<Address> {
-        use NativeOwner as O;
-        let StoredBalanceChange::V1 { owner, .. } = &self.stored;
-
-        match owner {
-            O::AddressOwner(addr)
-            | O::ObjectOwner(addr)
-            | O::ConsensusAddressOwner { owner: addr, .. } => {
-                Some(Address::with_address(self.scope.clone(), *addr))
-            }
-            O::Shared { .. } | O::Immutable => None,
-        }
-    }
+    /// Create a BalanceChange from a gRPC BalanceChange.
+    pub(crate) fn from_grpc(scope: Scope, grpc: &GrpcBalanceChange) -> Result<Self, RpcError> {
+        let address = grpc.address().parse().context("Failed to parse address")?;
+        let coin_type: TypeTag = grpc
+            .coin_type()
+            .parse()
+            .context("Failed to parse coin type")?;
+        let amount: i128 = grpc.amount().parse().context("Failed to parse amount")?;
 
-    /// The inner type of the coin whose balance has changed (e.g. `0x2::sui::SUI`).
-    async fn coin_type(&self) -> Result<Option<MoveType>, RpcError> {
-        let StoredBalanceChange::V1 { coin_type, .. } = &self.stored;
-        let type_ = TypeTag::from_str(coin_type).context("Failed to parse coin type")?;
-        Ok(Some(MoveType::from_native(type_, self.scope.clone())))
+        Ok(Self {
+            owner: Some(Address::with_address(scope.clone(), address)),
+            coin_type: Some(MoveType::from_native(coin_type, scope)),
+            amount: Some(BigInt::from(amount)),
+        })
     }
 
-    /// The signed balance change.
-    async fn amount(&self) -> Option<BigInt> {
-        let StoredBalanceChange::V1 { amount, .. } = &self.stored;
-        Some(BigInt::from(*amount))
+    /// Create a BalanceChange from a stored BalanceChange (database).
+    pub(crate) fn from_stored(scope: Scope, stored: StoredBalanceChange) -> Result<Self, RpcError> {
+        let StoredBalanceChange::V1 {
+            owner,
+            coin_type,
+            amount,
+        } = stored;
+
+        // Extract address from owner
+        let address = match owner {
+            Owner::AddressOwner(addr)
+            | Owner::ObjectOwner(addr)
+            | Owner::ConsensusAddressOwner { owner: addr, .. } => Some(addr),
+            Owner::Shared { .. } | Owner::Immutable => None,
+        };
+
+        let coin_type = TypeTag::from_str(&coin_type).context("Failed to parse coin type")?;
+
+        Ok(Self {
+            owner: address.map(|addr| Address::with_address(scope.clone(), addr)),
+            coin_type: Some(MoveType::from_native(coin_type, scope)),
+            amount: Some(BigInt::from(amount)),
+        })
     }
 }

crates/sui-indexer-alt-graphql/src/api/types/transaction_effects.rs

@@ -11,7 +11,7 @@ use sui_indexer_alt_reader::{
     pg_reader::PgReader,
     tx_balance_changes::TxBalanceChangeKey,
 };
-use sui_indexer_alt_schema::transactions::BalanceChange as NativeBalanceChange;
+use sui_indexer_alt_schema::transactions::BalanceChange as StoredBalanceChange;
 use sui_rpc::proto::sui::rpc::v2::ExecutedTransaction;
 use sui_types::{
     digests::TransactionDigest,
@@ -217,9 +217,21 @@ impl EffectsContents {
             return Ok(Some(Connection::new(false, false)));
         };
 
-        let transaction_digest = content.digest()?;
+        let pagination: &PaginationConfig = ctx.data()?;
+        let limits = pagination.limits("TransactionEffects", "balanceChanges");
+        let page = Page::from_params(limits, first, after, last, before)?;
 
-        // Load balance changes from database using DataLoader
+        // First try to get balance changes from execution context
+        if let Some(grpc_balance_changes) = content.balance_changes() {
+            return page
+                .paginate_indices(grpc_balance_changes.len(), |i| {
+                    BalanceChange::from_grpc(self.scope.clone(), &grpc_balance_changes[i])
+                })
+                .map(Some);
+        }
+
+        // Fall back to loading from database
+        let transaction_digest = content.digest()?;
         let pg_loader: &Arc<DataLoader<PgReader>> = ctx.data()?;
         let key = TxBalanceChangeKey(transaction_digest);
 
@@ -232,19 +244,12 @@ impl EffectsContents {
         };
 
         // Deserialize balance changes from BCS bytes
-        let balance_changes: Vec<NativeBalanceChange> =
+        let balance_changes: Vec<StoredBalanceChange> =
             bcs::from_bytes(&stored_balance_changes.balance_changes)
                 .context("Failed to deserialize balance changes")?;
 
-        let pagination: &PaginationConfig = ctx.data()?;
-        let limits = pagination.limits("TransactionEffects", "balanceChanges");
-        let page = Page::from_params(limits, first, after, last, before)?;
-
         page.paginate_indices(balance_changes.len(), |i| {
-            Ok(BalanceChange {
-                scope: self.scope.clone(),
-                stored: balance_changes[i].clone(),
-            })
+            BalanceChange::from_stored(self.scope.clone(), balance_changes[i].clone())
         })
         .map(Some)
     }

crates/sui-indexer-alt-reader/src/kv_loader.rs

@@ -9,7 +9,7 @@ use sui_indexer_alt_schema::transactions::StoredTransaction;
 use sui_kvstore::{
     TransactionData as KVTransactionData, TransactionEventsData as KVTransactionEventsData,
 };
-use sui_rpc::proto::sui::rpc::v2::ExecutedTransaction;
+use sui_rpc::proto::sui::rpc::v2 as grpc;
 use sui_types::{
     base_types::ObjectID,
     crypto::AuthorityQuorumSignInfo,
@@ -54,6 +54,7 @@ pub enum TransactionContents {
         events: Option<Vec<Event>>,
         transaction_data: Box<TransactionData>,
         signatures: Vec<GenericSignature>,
+        balance_changes: Vec<grpc::BalanceChange>,
     },
 }
 
@@ -218,7 +219,7 @@ impl KvLoader {
 impl TransactionContents {
     /// Create a TransactionContents from an ExecutedTransaction.
     pub fn from_executed_transaction(
-        executed_transaction: &ExecutedTransaction,
+        executed_transaction: &grpc::ExecutedTransaction,
         transaction_data: TransactionData,
         signatures: Vec<GenericSignature>,
     ) -> anyhow::Result<Self> {
@@ -240,11 +241,15 @@ impl TransactionContents {
             .transpose()?
             .map(|events: TransactionEvents| events.data);
 
+        // Extract balance changes from the gRPC response
+        let balance_changes = executed_transaction.balance_changes.clone();
+
         Ok(Self::ExecutedTransaction {
             effects: Box::new(effects),
             events,
             transaction_data: Box::new(transaction_data),
             signatures,
+            balance_changes,
         })
     }
 
@@ -311,6 +316,15 @@ impl TransactionContents {
         }
     }
 
+    pub fn balance_changes(&self) -> Option<Vec<grpc::BalanceChange>> {
+        match self {
+            Self::ExecutedTransaction {
+                balance_changes, ..
+            } => Some(balance_changes.clone()),
+            _ => None,
+        }
+    }
+
     pub fn raw_transaction(&self) -> anyhow::Result<Vec<u8>> {
         match self {
             Self::Pg(stored) => Ok(stored.raw_transaction.clone()),