Add getTokenAccount rpc method and implement logic (metaplex-foundati…

…on#213)

* feat: add getTokenAccount rpc method and implement logic

* add tests for token_accounts indexing

* chore: cleanup

* chore : add checks for getTokenAccounts and cover more scenarios for tests

* show_zero_balance_filter option added (#201)

---------

Co-authored-by: Ahzam Akhtar <[email protected]>

das_api/src/api/api_impl.rs

@@ -9,13 +9,13 @@ use digital_asset_types::{
     dapi::{
         get_asset, get_asset_proofs, get_asset_signatures, get_assets, get_assets_by_authority,
         get_assets_by_creator, get_assets_by_group, get_assets_by_owner, get_proof_for_asset,
-        search_assets,
+        get_token_accounts, search_assets,
     },
     rpc::{
         filter::{AssetSortBy, SearchConditionType},
-        response::GetGroupingResponse,
+        response::{GetGroupingResponse, TokenAccountList},
+        OwnershipModel, RoyaltyModel,
     },
-    rpc::{OwnershipModel, RoyaltyModel},
 };
 use open_rpc_derive::document_rpc;
 use sea_orm::{sea_query::ConditionType, ConnectionTrait, DbBackend, Statement};
@@ -516,4 +516,33 @@ impl ApiContract for DasApi {
             group_size: gs.size,
         })
     }
+
+    async fn get_token_accounts(
+        self: &DasApi,
+        payload: GetTokenAccounts,
+    ) -> Result<TokenAccountList, DasApiError> {
+        let GetTokenAccounts {
+            owner_address,
+            mint_address,
+            limit,
+            page,
+            before,
+            after,
+            options,
+            cursor,
+        } = payload;
+        let owner_address = validate_opt_pubkey(&owner_address)?;
+        let mint_address = validate_opt_pubkey(&mint_address)?;
+        let options = options.unwrap_or_default();
+        let page_options = self.validate_pagination(limit, page, &before, &after, &cursor, None)?;
+        get_token_accounts(
+            &self.db_connection,
+            owner_address,
+            mint_address,
+            &page_options,
+            &options,
+        )
+        .await
+        .map_err(Into::into)
+    }
 }

das_api/src/api/mod.rs

@@ -2,7 +2,7 @@ use crate::error::DasApiError;
 use async_trait::async_trait;
 use digital_asset_types::rpc::filter::{AssetSortDirection, SearchConditionType};
 use digital_asset_types::rpc::options::Options;
-use digital_asset_types::rpc::response::{AssetList, TransactionSignatureList};
+use digital_asset_types::rpc::response::{AssetList, TokenAccountList, TransactionSignatureList};
 use digital_asset_types::rpc::{filter::AssetSorting, response::GetGroupingResponse};
 use digital_asset_types::rpc::{Asset, AssetProof, Interface, OwnershipModel, RoyaltyModel};
 use open_rpc_derive::{document_rpc, rpc};
@@ -163,6 +163,21 @@ pub struct GetAssetSignatures {
     pub sort_direction: Option<AssetSortDirection>,
 }
 
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
+#[serde(deny_unknown_fields, rename_all = "camelCase")]
+pub struct GetTokenAccounts {
+    pub owner_address: Option<String>,
+    pub mint_address: Option<String>,
+    pub limit: Option<u32>,
+    pub page: Option<u32>,
+    pub before: Option<String>,
+    pub after: Option<String>,
+    #[serde(default, alias = "displayOptions")]
+    pub options: Option<Options>,
+    #[serde(default)]
+    pub cursor: Option<String>,
+}
+
 #[document_rpc]
 #[async_trait]
 pub trait ApiContract: Send + Sync + 'static {
@@ -251,4 +266,14 @@ pub trait ApiContract: Send + Sync + 'static {
         summary = "Get a list of assets grouped by a specific authority"
     )]
     async fn get_grouping(&self, payload: GetGrouping) -> Result<GetGroupingResponse, DasApiError>;
+
+    #[rpc(
+        name = "getTokenAccounts",
+        params = "named",
+        summary = "Get a list of token accounts by owner or mint"
+    )]
+    async fn get_token_accounts(
+        &self,
+        payload: GetTokenAccounts,
+    ) -> Result<TokenAccountList, DasApiError>;
 }

das_api/src/builder.rs

@@ -118,6 +118,19 @@ impl RpcApiBuilder {
             Ok(rpc_context.schema())
         })?;
 
+        module.register_async_method(
+            "get_token_accounts",
+            |rpc_params, rpc_context| async move {
+                let payload = rpc_params.parse::<GetTokenAccounts>()?;
+                rpc_context
+                    .get_token_accounts(payload)
+                    .await
+                    .map_err(Into::into)
+            },
+        )?;
+
+        module.register_alias("getTokenAccounts", "get_token_accounts")?;
+
         Ok(module)
     }
 }

digital_asset_types/src/dao/scopes/asset.rs

@@ -4,9 +4,9 @@ use crate::{
         asset_authority, asset_creators, asset_data, asset_grouping, cl_audits_v2,
         extensions::{self, instruction::PascalCase},
         sea_orm_active_enums::Instruction,
-        Cursor, FullAsset, GroupingSize, Pagination,
+        token_accounts, Cursor, FullAsset, GroupingSize, Pagination,
     },
-    rpc::filter::AssetSortDirection,
+    rpc::{filter::AssetSortDirection, options::Options},
 };
 use indexmap::IndexMap;
 use sea_orm::{entity::*, query::*, ConnectionTrait, DbErr, Order};
@@ -553,3 +553,45 @@ fn filter_out_stale_creators(creators: &mut Vec<asset_creators::Model>) {
         }
     }
 }
+
+pub async fn get_token_accounts(
+    conn: &impl ConnectionTrait,
+    owner_address: Option<Vec<u8>>,
+    mint_address: Option<Vec<u8>>,
+    pagination: &Pagination,
+    limit: u64,
+    options: &Options,
+) -> Result<Vec<token_accounts::Model>, DbErr> {
+    let mut condition = Condition::all();
+
+    if options.show_zero_balance {
+        condition = condition.add(token_accounts::Column::Amount.gte(0));
+    } else {
+        condition = condition.add(token_accounts::Column::Amount.gt(0));
+    }
+
+    if owner_address.is_none() && mint_address.is_none() {
+        return Err(DbErr::Custom(
+            "Either 'owner_address' or 'mint_address' must be provided".to_string(),
+        ));
+    }
+
+    if let Some(owner) = owner_address {
+        condition = condition.add(token_accounts::Column::Owner.eq(owner));
+    }
+    if let Some(mint) = mint_address {
+        condition = condition.add(token_accounts::Column::Mint.eq(mint));
+    }
+
+    let token_accounts = paginate(
+        pagination,
+        limit,
+        token_accounts::Entity::find().filter(condition),
+        Order::Asc,
+        token_accounts::Column::Pubkey,
+    )
+    .all(conn)
+    .await?;
+
+    Ok(token_accounts)
+}

digital_asset_types/src/dapi/common/asset.rs

@@ -1,15 +1,18 @@
 use crate::dao::sea_orm_active_enums::SpecificationVersions;
+use crate::dao::token_accounts;
 use crate::dao::FullAsset;
 use crate::dao::PageOptions;
 use crate::dao::Pagination;
 use crate::dao::{asset, asset_authority, asset_creators, asset_data, asset_grouping};
 use crate::rpc::filter::{AssetSortBy, AssetSortDirection, AssetSorting};
 use crate::rpc::options::Options;
+use crate::rpc::response::TokenAccountList;
 use crate::rpc::response::TransactionSignatureList;
-use crate::rpc::response::{AssetError, AssetList};
+use crate::rpc::response::{AssetList, DasError};
 use crate::rpc::{
     Asset as RpcAsset, Authority, Compression, Content, Creator, File, Group, Interface,
-    MetadataMap, MplCoreInfo, Ownership, Royalty, Scope, Supply, Uses,
+    MetadataMap, MplCoreInfo, Ownership, Royalty, Scope, Supply, TokenAccount as RpcTokenAccount,
+    Uses,
 };
 use jsonpath_lib::JsonPathError;
 use log::warn;
@@ -455,18 +458,102 @@ pub fn asset_to_rpc(asset: FullAsset, options: &Options) -> Result<RpcAsset, DbE
 pub fn asset_list_to_rpc(
     asset_list: Vec<FullAsset>,
     options: &Options,
-) -> (Vec<RpcAsset>, Vec<AssetError>) {
+) -> (Vec<RpcAsset>, Vec<DasError>) {
     asset_list
         .into_iter()
         .fold((vec![], vec![]), |(mut assets, mut errors), asset| {
             let id = bs58::encode(asset.asset.id.clone()).into_string();
             match asset_to_rpc(asset, options) {
                 Ok(rpc_asset) => assets.push(rpc_asset),
-                Err(e) => errors.push(AssetError {
+                Err(e) => errors.push(DasError {
                     id,
                     error: e.to_string(),
                 }),
             }
             (assets, errors)
         })
 }
+
+pub fn token_account_to_rpc(
+    token_account: token_accounts::Model,
+    _options: &Options,
+) -> Result<RpcTokenAccount, DbErr> {
+    let address = bs58::encode(token_account.pubkey.clone()).into_string();
+    let mint = bs58::encode(token_account.mint.clone()).into_string();
+    let owner = bs58::encode(token_account.owner.clone()).into_string();
+    let delegate = token_account
+        .delegate
+        .map(|d| bs58::encode(d).into_string());
+    let close_authority = token_account
+        .close_authority
+        .map(|d| bs58::encode(d).into_string());
+
+    Ok(RpcTokenAccount {
+        address,
+        mint,
+        amount: token_account.amount as u64,
+        owner,
+        frozen: token_account.frozen,
+        delegate,
+        delegated_amount: token_account.delegated_amount as u64,
+        close_authority,
+        extensions: None,
+    })
+}
+
+pub fn token_account_list_to_rpc(
+    token_accounts: Vec<token_accounts::Model>,
+    options: &Options,
+) -> (Vec<RpcTokenAccount>, Vec<DasError>) {
+    token_accounts.into_iter().fold(
+        (vec![], vec![]),
+        |(mut accounts, mut errors), token_account| {
+            let id = bs58::encode(token_account.pubkey.clone()).into_string();
+            match token_account_to_rpc(token_account, options) {
+                Ok(rpc_token_account) => accounts.push(rpc_token_account),
+                Err(e) => errors.push(DasError {
+                    id,
+                    error: e.to_string(),
+                }),
+            }
+            (accounts, errors)
+        },
+    )
+}
+
+pub fn build_token_list_response(
+    token_accounts: Vec<token_accounts::Model>,
+    limit: u64,
+    pagination: &Pagination,
+    options: &Options,
+) -> TokenAccountList {
+    let total = token_accounts.len() as u32;
+    let (page, before, after, cursor) = match pagination {
+        Pagination::Keyset { before, after } => {
+            let bef = before.clone().and_then(|x| String::from_utf8(x).ok());
+            let aft = after.clone().and_then(|x| String::from_utf8(x).ok());
+            (None, bef, aft, None)
+        }
+        Pagination::Page { page } => (Some(*page as u32), None, None, None),
+        Pagination::Cursor(_) => {
+            if let Some(last_token_account) = token_accounts.last() {
+                let cursor_str = bs58::encode(&last_token_account.pubkey.clone()).into_string();
+                (None, None, None, Some(cursor_str))
+            } else {
+                (None, None, None, None)
+            }
+        }
+    };
+
+    let (items, errors) = token_account_list_to_rpc(token_accounts, options);
+    TokenAccountList {
+        total,
+        limit: limit as u32,
+        page,
+        before,
+        after,
+        token_accounts: items,
+        cursor,
+        errors,
+    }
+}

digital_asset_types/src/dapi/get_token_accounts.rs

@@ -0,0 +1,35 @@
+use sea_orm::{DatabaseConnection, DbErr};
+
+use crate::{
+    dao::PageOptions,
+    rpc::{options::Options, response::TokenAccountList},
+};
+
+use super::common::{build_token_list_response, create_pagination};
+
+pub async fn get_token_accounts(
+    db: &DatabaseConnection,
+    owner_address: Option<Vec<u8>>,
+    mint_address: Option<Vec<u8>>,
+    page_options: &PageOptions,
+    options: &Options,
+) -> Result<TokenAccountList, DbErr> {
+    let pagination = create_pagination(page_options)?;
+
+    let token_accounts = crate::dao::scopes::asset::get_token_accounts(
+        db,
+        owner_address,
+        mint_address,
+        &pagination,
+        page_options.limit,
+        options,
+    )
+    .await?;
+
+    Ok(build_token_list_response(
+        token_accounts,
+        page_options.limit,
+        &pagination,
+        options,
+    ))
+}

digital_asset_types/src/dapi/mod.rs

@@ -5,6 +5,7 @@ mod assets_by_owner;
 mod change_logs;
 mod get_asset;
 mod get_asset_signatures;
+mod get_token_accounts;
 mod search_assets;
 
 pub mod common;
@@ -16,4 +17,5 @@ pub use assets_by_owner::*;
 pub use change_logs::*;
 pub use get_asset::*;
 pub use get_asset_signatures::*;
+pub use get_token_accounts::*;
 pub use search_assets::*;

digital_asset_types/src/rpc/asset.rs

@@ -399,3 +399,17 @@ pub struct Asset {
     #[serde(skip_serializing_if = "Option::is_none")]
     pub unknown_external_plugins: Option<Value>,
 }
+
+#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Default, JsonSchema)]
+#[serde(default)]
+pub struct TokenAccount {
+    pub address: String,
+    pub mint: String,
+    pub amount: u64,
+    pub owner: String,
+    pub frozen: bool,
+    pub delegate: Option<String>,
+    pub delegated_amount: u64,
+    pub close_authority: Option<String>,
+    pub extensions: Option<Value>,
+}

digital_asset_types/src/rpc/options.rs

@@ -6,4 +6,6 @@ use serde::{Deserialize, Serialize};
 pub struct Options {
     #[serde(default)]
     pub show_unverified_collections: bool,
+    #[serde(default)]
+    pub show_zero_balance: bool,
 }