SPL token balance in transaction metadata (#13673)

* feat: store pre / post token balances

* move helper functions into separate include

* move token balance functionality to transaction-status crate

* fix blockstore processor test

* fix bigtable legacy test

* add caching to decimals

storage-proto/Cargo.toml

@@ -13,6 +13,7 @@ bincode = "1.2.1"
 prost = "0.6.1"
 serde = "1.0.112"
 serde_derive = "1.0.103"
+solana-account-decoder = { path = "../account-decoder", version = "1.5.0" }
 solana-sdk = { path = "../sdk", version = "1.5.0" }
 solana-transaction-status = { path = "../transaction-status", version = "1.5.0" }
 

storage-proto/proto/solana.storage.confirmed_block.rs

@@ -61,6 +61,10 @@ pub struct TransactionStatusMeta {
     pub inner_instructions: ::std::vec::Vec<InnerInstructions>,
     #[prost(string, repeated, tag = "6")]
     pub log_messages: ::std::vec::Vec<std::string::String>,
+    #[prost(message, repeated, tag = "7")]
+    pub pre_token_balances: ::std::vec::Vec<TokenBalance>,
+    #[prost(message, repeated, tag = "8")]
+    pub post_token_balances: ::std::vec::Vec<TokenBalance>,
 }
 #[derive(Clone, PartialEq, ::prost::Message)]
 pub struct TransactionError {
@@ -84,6 +88,24 @@ pub struct CompiledInstruction {
     pub data: std::vec::Vec<u8>,
 }
 #[derive(Clone, PartialEq, ::prost::Message)]
+pub struct TokenBalance {
+    #[prost(uint32, tag = "1")]
+    pub account_index: u32,
+    #[prost(string, tag = "2")]
+    pub mint: std::string::String,
+    #[prost(message, optional, tag = "3")]
+    pub ui_token_amount: ::std::option::Option<UiTokenAmount>,
+}
+#[derive(Clone, PartialEq, ::prost::Message)]
+pub struct UiTokenAmount {
+    #[prost(double, tag = "1")]
+    pub ui_amount: f64,
+    #[prost(uint32, tag = "2")]
+    pub decimals: u32,
+    #[prost(string, tag = "3")]
+    pub amount: std::string::String,
+}
+#[derive(Clone, PartialEq, ::prost::Message)]
 pub struct Reward {
     #[prost(string, tag = "1")]
     pub pubkey: std::string::String,

storage-proto/src/confirmed_block.proto

@@ -41,6 +41,8 @@ message TransactionStatusMeta {
     repeated uint64 post_balances = 4;
     repeated InnerInstructions inner_instructions = 5;
     repeated string log_messages = 6;
+    repeated TokenBalance pre_token_balances = 7;
+    repeated TokenBalance post_token_balances = 8;
 }
 
 message TransactionError {
@@ -58,6 +60,18 @@ message CompiledInstruction {
     bytes data = 3;
 }
 
+message TokenBalance {
+    uint32 account_index = 1;
+    string mint = 2;
+    UiTokenAmount ui_token_amount = 3;
+}
+
+message UiTokenAmount {
+    double ui_amount = 1;
+    uint32 decimals = 2;
+    string amount = 3;
+}
+
 enum RewardType {
     Unspecified = 0;
     Fee = 1;

storage-proto/src/convert.rs

@@ -1,4 +1,5 @@
 use crate::StoredExtendedRewards;
+use solana_account_decoder::parse_token::UiTokenAmount;
 use solana_sdk::{
     hash::Hash,
     instruction::CompiledInstruction,
@@ -9,7 +10,7 @@ use solana_sdk::{
 };
 use solana_transaction_status::{
     ConfirmedBlock, InnerInstructions, Reward, RewardType, TransactionStatusMeta,
-    TransactionWithStatusMeta,
+    TransactionTokenBalance, TransactionWithStatusMeta,
 };
 use std::convert::{TryFrom, TryInto};
 
@@ -260,6 +261,8 @@ impl From<TransactionStatusMeta> for generated::TransactionStatusMeta {
             post_balances,
             inner_instructions,
             log_messages,
+            pre_token_balances,
+            post_token_balances,
         } = value;
         let err = match status {
             Ok(()) => None,
@@ -273,6 +276,17 @@ impl From<TransactionStatusMeta> for generated::TransactionStatusMeta {
             .map(|ii| ii.into())
             .collect();
         let log_messages = log_messages.unwrap_or_default();
+        let pre_token_balances = pre_token_balances
+            .unwrap_or_default()
+            .into_iter()
+            .map(|balance| balance.into())
+            .collect();
+        let post_token_balances = post_token_balances
+            .unwrap_or_default()
+            .into_iter()
+            .map(|balance| balance.into())
+            .collect();
+
         Self {
             err,
             fee,
@@ -280,6 +294,8 @@ impl From<TransactionStatusMeta> for generated::TransactionStatusMeta {
             post_balances,
             inner_instructions,
             log_messages,
+            pre_token_balances,
+            post_token_balances,
         }
     }
 }
@@ -295,6 +311,8 @@ impl TryFrom<generated::TransactionStatusMeta> for TransactionStatusMeta {
             post_balances,
             inner_instructions,
             log_messages,
+            pre_token_balances,
+            post_token_balances,
         } = value;
         let status = match &err {
             None => Ok(()),
@@ -307,6 +325,18 @@ impl TryFrom<generated::TransactionStatusMeta> for TransactionStatusMeta {
                 .collect(),
         );
         let log_messages = Some(log_messages);
+        let pre_token_balances = Some(
+            pre_token_balances
+                .into_iter()
+                .map(|balance| balance.into())
+                .collect(),
+        );
+        let post_token_balances = Some(
+            post_token_balances
+                .into_iter()
+                .map(|balance| balance.into())
+                .collect(),
+        );
         Ok(Self {
             status,
             fee,
@@ -314,6 +344,8 @@ impl TryFrom<generated::TransactionStatusMeta> for TransactionStatusMeta {
             post_balances,
             inner_instructions,
             log_messages,
+            pre_token_balances,
+            post_token_balances,
         })
     }
 }
@@ -336,6 +368,35 @@ impl From<generated::InnerInstructions> for InnerInstructions {
     }
 }
 
+impl From<TransactionTokenBalance> for generated::TokenBalance {
+    fn from(value: TransactionTokenBalance) -> Self {
+        Self {
+            account_index: value.account_index as u32,
+            mint: value.mint,
+            ui_token_amount: Some(generated::UiTokenAmount {
+                ui_amount: value.ui_token_amount.ui_amount,
+                decimals: value.ui_token_amount.decimals as u32,
+                amount: value.ui_token_amount.amount,
+            }),
+        }
+    }
+}
+
+impl From<generated::TokenBalance> for TransactionTokenBalance {
+    fn from(value: generated::TokenBalance) -> Self {
+        let ui_token_amount = value.ui_token_amount.unwrap_or_default();
+        Self {
+            account_index: value.account_index as u8,
+            mint: value.mint,
+            ui_token_amount: UiTokenAmount {
+                ui_amount: ui_token_amount.ui_amount,
+                decimals: ui_token_amount.decimals as u8,
+                amount: ui_token_amount.amount,
+            },
+        }
+    }
+}
+
 impl From<CompiledInstruction> for generated::CompiledInstruction {
     fn from(value: CompiledInstruction) -> Self {
         Self {