feat(torii-graphql): subscriptions for tokens and token balances

crates/torii/graphql/src/constants.rs

@@ -36,7 +36,7 @@ pub const MODEL_ORDER_TYPE_NAME: &str = "World__ModelOrder";
 pub const MODEL_ORDER_FIELD_TYPE_NAME: &str = "World__ModelOrderField";
 pub const TOKEN_BALANCE_TYPE_NAME: &str = "Token__Balance";
 pub const TOKEN_TRANSFER_TYPE_NAME: &str = "Token__Transfer";
-pub const TOKEN_TYPE_NAME: &str = "ERC__Token";
+pub const TOKEN_UNION_TYPE_NAME: &str = "ERC__Token";
 // pub const ERC721_METADATA_TYPE_NAME: &str = "ERC721__Metadata";
 
 pub const ERC20_TYPE_NAME: &str = "ERC20__Token";

crates/torii/graphql/src/mapping.rs

@@ -4,7 +4,7 @@ use async_graphql::Name;
 use dojo_types::primitive::Primitive;
 use lazy_static::lazy_static;
 
-use crate::constants::{CONTENT_TYPE_NAME, SOCIAL_TYPE_NAME, TOKEN_TYPE_NAME};
+use crate::constants::{CONTENT_TYPE_NAME, SOCIAL_TYPE_NAME, TOKEN_UNION_TYPE_NAME};
 use crate::types::{GraphqlType, TypeData, TypeMapping};
 
 lazy_static! {
@@ -146,14 +146,14 @@ lazy_static! {
     ]);
 
     pub static ref TOKEN_BALANCE_TYPE_MAPPING: TypeMapping = IndexMap::from([
-        (Name::new("tokenMetadata"), TypeData::Nested((TypeRef::named_nn(TOKEN_TYPE_NAME), IndexMap::new()))),
+        (Name::new("tokenMetadata"), TypeData::Nested((TypeRef::named_nn(TOKEN_UNION_TYPE_NAME), IndexMap::new()))),
     ]);
 
     pub static ref TOKEN_TRANSFER_TYPE_MAPPING: TypeMapping = IndexMap::from([
         (Name::new("from"), TypeData::Simple(TypeRef::named_nn(TypeRef::STRING))),
         (Name::new("to"), TypeData::Simple(TypeRef::named_nn(TypeRef::STRING))),
         (Name::new("executedAt"), TypeData::Simple(TypeRef::named_nn(TypeRef::STRING))),
-        (Name::new("tokenMetadata"), TypeData::Nested((TypeRef::named_nn(TOKEN_TYPE_NAME), IndexMap::new()))),
+        (Name::new("tokenMetadata"), TypeData::Nested((TypeRef::named_nn(TOKEN_UNION_TYPE_NAME), IndexMap::new()))),
         (Name::new("transactionHash"), TypeData::Simple(TypeRef::named_nn(TypeRef::STRING))),
     ]);
 
@@ -190,4 +190,13 @@ lazy_static! {
             TypeData::Simple(TypeRef::named_nn(GraphqlType::DateTime.to_string())),
         ),
     ]);
+
+    pub static ref TOKEN_TYPE_MAPPING: TypeMapping = IndexMap::from([
+        (Name::new("id"), TypeData::Simple(TypeRef::named_nn(TypeRef::STRING))),
+        (Name::new("contractAddress"), TypeData::Simple(TypeRef::named_nn(TypeRef::STRING))),
+        (Name::new("name"), TypeData::Simple(TypeRef::named_nn(TypeRef::STRING))),
+        (Name::new("symbol"), TypeData::Simple(TypeRef::named_nn(TypeRef::STRING))),
+        (Name::new("decimals"), TypeData::Simple(TypeRef::named_nn(TypeRef::INT))),
+        (Name::new("tokenMetadata"), TypeData::Nested((TypeRef::named_nn(TOKEN_UNION_TYPE_NAME), IndexMap::new()))),
+    ]);
 }

crates/torii/graphql/src/object/erc/erc_token.rs

@@ -1,9 +1,15 @@
-use async_graphql::dynamic::FieldValue;
+use async_graphql::dynamic::{
+    Field, FieldValue, SubscriptionField, SubscriptionFieldFuture, TypeRef,
+};
 use async_graphql::{Name, Value};
+use sqlx::{Pool, Row, Sqlite};
+use tokio_stream::StreamExt;
+use torii_sqlite::simple_broker::SimpleBroker;
+use torii_sqlite::types::Token;
 
 use crate::constants::{ERC20_TOKEN_NAME, ERC20_TYPE_NAME, ERC721_TOKEN_NAME, ERC721_TYPE_NAME};
-use crate::mapping::{ERC20_TOKEN_TYPE_MAPPING, ERC721_TOKEN_TYPE_MAPPING};
-use crate::object::BasicObject;
+use crate::mapping::{ERC20_TOKEN_TYPE_MAPPING, ERC721_TOKEN_TYPE_MAPPING, TOKEN_TYPE_MAPPING};
+use crate::object::{BasicObject, ResolvableObject};
 use crate::types::{TypeMapping, ValueMapping};
 
 #[derive(Debug)]
@@ -107,3 +113,141 @@
         }
     }
 }
+
+#[derive(Debug)]
+pub struct TokenObject;
+
+impl BasicObject for TokenObject {
+    fn name(&self) -> (&str, &str) {
+        ("tokens", "token")
+    }
+
+    fn type_name(&self) -> &str {
+        "Token"
+    }
+
+    fn type_mapping(&self) -> &TypeMapping {
+        &TOKEN_TYPE_MAPPING
+    }
+}
+
+impl ResolvableObject for TokenObject {
+    fn resolvers(&self) -> Vec<Field> {
+        vec![]
+    }
+
+    fn subscriptions(&self) -> Option<Vec<SubscriptionField>> {
+        Some(vec![SubscriptionField::new(
+            "tokenUpdated",
+            TypeRef::named_nn(self.type_name()),
+            |ctx| {
+                SubscriptionFieldFuture::new(async move {
+                    let pool = ctx.data::<Pool<Sqlite>>()?;
+                    Ok(SimpleBroker::<Token>::subscribe()
+                        .then(move |token| {
+                            let pool = pool.clone();
+                            async move {
+                                // Fetch complete token data including contract type
+                                let query = "SELECT t.*, c.contract_type 
+                                               FROM tokens t 
+                                               JOIN contracts c ON t.contract_address = \
+                                             c.contract_address 
+                                               WHERE t.id = ?";
+
+                                let row =
+                                    match sqlx::query(query).bind(&token.id).fetch_one(&pool).await
+                                    {
+                                        Ok(row) => row,
+                                        Err(_) => return None,
+                                    };
+
+                                let contract_type: String = row.get("contract_type");
+                                let token_metadata = match contract_type.to_lowercase().as_str() {
+                                    "erc20" => {
+                                        let token = Erc20Token {
+                                            contract_address: row.get("contract_address"),
+                                            name: row.get("name"),
+                                            symbol: row.get("symbol"),
+                                            decimals: row.get("decimals"),
+                                            amount: "0".to_string(), // New token has no balance
+                                        };
+                                        ErcTokenType::Erc20(token)
+                                    }
+                                    "erc721" => {
+                                        let metadata_str: String = row.get("metadata");
+                                        let (
+                                            metadata_str,
+                                            metadata_name,
+                                            metadata_description,
+                                            metadata_attributes,
+                                            image_path,
+                                        ) = if metadata_str.is_empty() {
+                                            (String::new(), None, None, None, String::new())
+                                        } else {
+                                            let metadata: serde_json::Value =
+                                                serde_json::from_str(&metadata_str)
+                                                    .expect("metadata is always json");
+                                            let metadata_name = metadata.get("name").map(|v| {
+                                                v.to_string().trim_matches('"').to_string()
+                                            });
+                                            let metadata_description =
+                                                metadata.get("description").map(|v| {
+                                                    v.to_string().trim_matches('"').to_string()
+                                                });
+                                            let metadata_attributes =
+                                                metadata.get("attributes").map(|v| {
+                                                    v.to_string().trim_matches('"').to_string()
+                                                });
+
+                                            let contract_address: String =
+                                                row.get("contract_address");
+                                            let image_path = format!("{}/image", contract_address);
+
+                                            (
+                                                metadata_str,
+                                                metadata_name,
+                                                metadata_description,
+                                                metadata_attributes,
+                                                image_path,
+                                            )
+                                        };
+
+                                        let token = Erc721Token {
+                                            name: row.get("name"),
+                                            metadata: metadata_str,
+                                            contract_address: row.get("contract_address"),
+                                            symbol: row.get("symbol"),
+                                            token_id: "0".to_string(), /* New token has no
+                                                                        * specific token_id */
+                                            metadata_name,
+                                            metadata_description,
+                                            metadata_attributes,
+                                            image_path,
+                                        };
+                                        ErcTokenType::Erc721(token)
+                                    }
+                                    _ => return None,
+                                };
+
+                                Some(Ok(FieldValue::value(Value::Object(ValueMapping::from([
+                                    (Name::new("id"), Value::String(token.id)),
+                                    (
+                                        Name::new("contractAddress"),
+                                        Value::String(token.contract_address),
+                                    ),
+                                    (Name::new("name"), Value::String(token.name)),
+                                    (Name::new("symbol"), Value::String(token.symbol)),
+                                    (Name::new("decimals"), Value::Number(token.decimals.into())),
+                                    (
+                                        Name::new("tokenMetadata"),
+                                        token_metadata.to_field_value().as_value().unwrap().clone(),
+                                    ),
+                                ])))))
+                            }
+                        })
+                        .filter_map(|result| result))
+                })
+            },
+        )])
+    }
+}

crates/torii/graphql/src/object/erc/token_balance.rs

@@ -1,11 +1,16 @@
 use async_graphql::connection::PageInfo;
-use async_graphql::dynamic::{Field, FieldFuture, FieldValue, InputValue, TypeRef};
+use async_graphql::dynamic::{
+    Field, FieldFuture, FieldValue, InputValue, SubscriptionField, SubscriptionFieldFuture, TypeRef,
+};
 use convert_case::{Case, Casing};
 use serde::Deserialize;
 use sqlx::sqlite::SqliteRow;
 use sqlx::{FromRow, Pool, Row, Sqlite, SqliteConnection};
 use starknet_crypto::Felt;
+use tokio_stream::StreamExt;
 use torii_sqlite::constants::TOKEN_BALANCE_TABLE;
+use torii_sqlite::simple_broker::SimpleBroker;
+use torii_sqlite::types::TokenBalance;
 use torii_sqlite::utils::felt_to_sql_string;
 use tracing::warn;
 
@@ -85,6 +90,76 @@
         field = connection_arguments(field);
         vec![field]
     }
+
+    fn subscriptions(&self) -> Option<Vec<SubscriptionField>> {
+        Some(vec![
+            SubscriptionField::new(
+                "tokenBalanceUpdated",
+                TypeRef::named_nn(format!("{}Connection", self.type_name())),
+                |ctx| {
+                    SubscriptionFieldFuture::new(async move {
+                        let address = match ctx.args.get("accountAddress") {
+                            Some(addr) => Some(addr.string()?.to_string()),
+                            None => None,
+                        };
+
+                        let pool = ctx.data::<Pool<Sqlite>>()?;
+                        Ok(SimpleBroker::<TokenBalance>::subscribe()
+                            .then(move |token_balance| {
+                                let address = address.clone();
+                                let pool = pool.clone();
+                                async move {
+                                    // Filter by account address if provided
+                                    if let Some(addr) = &address {
+                                        if token_balance.account_address != *addr {
+                                            return None;
+                                        }
+                                    }
+                                    // Fetch associated token data
+                                    let query = format!(
+                                        "SELECT b.id, t.contract_address, t.name, t.symbol, \
+                                         t.decimals, b.balance, b.token_id, t.metadata, \
+                                         c.contract_type
+                                        FROM {} b
+                                        JOIN tokens t ON b.token_id = t.id
+                                        JOIN contracts c ON t.contract_address = \
+                                         c.contract_address
+                                        WHERE b.id = ?",
+                                        TOKEN_BALANCE_TABLE
+                                    );
+
+                                    let row = match sqlx::query(&query)
+                                        .bind(&token_balance.id)
+                                        .fetch_one(&pool)
+                                        .await
+                                    {
+                                        Ok(row) => row,
+                                        Err(_) => return None,
+                                    };
+
+                                    // Create connection with single edge
+                                    match token_balances_connection_output(
+                                        &[row],
+                                        1, // total_count
+                                        PageInfo {
+                                            has_previous_page: false,
+                                            has_next_page: false,
+                                            start_cursor: None,
+                                            end_cursor: None,
+                                        },
+                                    ) {
+                                        Ok(value) => Some(Ok(value)),
+                                        Err(_) => None,
+                                    }
+                                }
+                            })
+                            .filter_map(|result| result))
+                    })
+                },
+            )
+            .argument(InputValue::new("accountAddress", TypeRef::named(TypeRef::STRING))),
+        ])
+    }
 }
 
 async fn fetch_token_balances(

crates/torii/graphql/src/schema.rs

@@ -12,11 +12,11 @@ use super::types::ScalarType;
 use super::utils;
 use crate::constants::{
     EMPTY_TYPE_NAME, ERC20_TYPE_NAME, ERC721_TYPE_NAME, QUERY_TYPE_NAME, SUBSCRIPTION_TYPE_NAME,
-    TOKEN_TYPE_NAME,
+    TOKEN_UNION_TYPE_NAME,
 };
 use crate::object::controller::ControllerObject;
 use crate::object::empty::EmptyObject;
-use crate::object::erc::erc_token::{Erc20TokenObject, Erc721TokenObject};
+use crate::object::erc::erc_token::{Erc20TokenObject, Erc721TokenObject, TokenObject};
 use crate::object::erc::token_balance::ErcBalanceObject;
 use crate::object::erc::token_transfer::ErcTransferObject;
 use crate::object::event_message::EventMessageObject;
@@ -131,15 +131,17 @@ async fn build_objects(pool: &SqlitePool) -> Result<(Vec<ObjectVariant>, Vec<Uni
         ObjectVariant::Basic(Box::new(Erc721TokenObject)),
         ObjectVariant::Basic(Box::new(Erc20TokenObject)),
         ObjectVariant::Basic(Box::new(EmptyObject)),
+        ObjectVariant::Resolvable(Box::new(TokenObject)),
     ];
 
     // model union object
     let mut unions: Vec<Union> = Vec::new();
     let mut model_union = Union::new("ModelUnion");
 
     // erc_token union object
-    let erc_token_union =
-        Union::new(TOKEN_TYPE_NAME).possible_type(ERC20_TYPE_NAME).possible_type(ERC721_TYPE_NAME);
+    let erc_token_union = Union::new(TOKEN_UNION_TYPE_NAME)
+        .possible_type(ERC20_TYPE_NAME)
+        .possible_type(ERC721_TYPE_NAME);
 
     unions.push(erc_token_union);