bittensor_rs/queries/
account.rs1use crate::api::api;
6use crate::error::BittensorError;
7use crate::types::Balance;
8use crate::AccountId;
9use std::str::FromStr;
10use subxt::OnlineClient;
11use subxt::PolkadotConfig;
12
13pub async fn get_balance(
37 client: &OnlineClient<PolkadotConfig>,
38 address: &str,
39) -> Result<Balance, BittensorError> {
40 let account_id = AccountId::from_str(address).map_err(|_| BittensorError::InvalidHotkey {
41 hotkey: address.to_string(),
42 })?;
43
44 let storage = api::storage().system().account(account_id);
45
46 let account_info = client
47 .storage()
48 .at_latest()
49 .await
50 .map_err(|e| BittensorError::RpcError {
51 message: format!("Failed to get storage: {}", e),
52 })?
53 .fetch(&storage)
54 .await
55 .map_err(|e| BittensorError::StorageQueryError {
56 key: "account".to_string(),
57 message: format!("Failed to fetch account: {}", e),
58 })?;
59
60 match account_info {
61 Some(info) => Ok(Balance::from_rao(info.data.free)),
62 None => Ok(Balance::zero()),
63 }
64}
65
66pub async fn get_stake(
75 client: &OnlineClient<PolkadotConfig>,
76 hotkey: &str,
77 coldkey: &str,
78 netuid: u16,
79) -> Result<Balance, BittensorError> {
80 let hotkey_id = AccountId::from_str(hotkey).map_err(|_| BittensorError::InvalidHotkey {
81 hotkey: hotkey.to_string(),
82 })?;
83
84 let coldkey_id = AccountId::from_str(coldkey).map_err(|_| BittensorError::InvalidHotkey {
85 hotkey: coldkey.to_string(),
86 })?;
87
88 let storage = api::storage()
89 .subtensor_module()
90 .alpha(hotkey_id, coldkey_id, netuid);
91
92 let stake = client
93 .storage()
94 .at_latest()
95 .await
96 .map_err(|e| BittensorError::RpcError {
97 message: format!("Failed to get storage: {}", e),
98 })?
99 .fetch(&storage)
100 .await
101 .map_err(|e| BittensorError::StorageQueryError {
102 key: "stake".to_string(),
103 message: format!("Failed to fetch stake: {}", e),
104 })?;
105
106 match stake {
107 Some(amount) => {
110 let raw_bits: u128 = amount.bits;
112 let rao = (raw_bits >> 64) as u64;
114 Ok(Balance::from_rao(rao))
115 }
116 None => Ok(Balance::zero()),
117 }
118}
119
120pub async fn get_total_network_stake(
128 client: &OnlineClient<PolkadotConfig>,
129) -> Result<Balance, BittensorError> {
130 let storage = api::storage().subtensor_module().total_stake();
131
132 let stake = client
133 .storage()
134 .at_latest()
135 .await
136 .map_err(|e| BittensorError::RpcError {
137 message: format!("Failed to get storage: {}", e),
138 })?
139 .fetch(&storage)
140 .await
141 .map_err(|e| BittensorError::StorageQueryError {
142 key: "total_stake".to_string(),
143 message: format!("Failed to fetch total stake: {}", e),
144 })?;
145
146 match stake {
147 Some(amount) => Ok(Balance::from_rao(amount)),
148 None => Ok(Balance::zero()),
149 }
150}
151
152#[derive(Debug, Clone)]
154pub struct StakeInfo {
155 pub hotkey: AccountId,
157 pub coldkey: AccountId,
159 pub netuid: u16,
161 pub stake: Balance,
163 pub locked: Balance,
165 pub emission: Balance,
167 pub tao_emission: Balance,
169 pub drain: Balance,
171 pub is_registered: bool,
173}
174
175pub async fn get_stake_info_for_coldkey(
182 client: &OnlineClient<PolkadotConfig>,
183 coldkey: &str,
184) -> Result<Vec<StakeInfo>, BittensorError> {
185 let coldkey_id = AccountId::from_str(coldkey).map_err(|_| BittensorError::InvalidHotkey {
186 hotkey: coldkey.to_string(),
187 })?;
188
189 let runtime_api =
190 client
191 .runtime_api()
192 .at_latest()
193 .await
194 .map_err(|e| BittensorError::RpcError {
195 message: format!("Failed to get runtime API: {}", e),
196 })?;
197
198 let stake_infos = runtime_api
199 .call(
200 api::runtime_apis::stake_info_runtime_api::StakeInfoRuntimeApi
201 .get_stake_info_for_coldkey(coldkey_id.clone()),
202 )
203 .await
204 .map_err(|e| BittensorError::RpcMethodError {
205 method: "get_stake_info_for_coldkey".to_string(),
206 message: e.to_string(),
207 })?;
208
209 let result = stake_infos
210 .into_iter()
211 .map(|info| StakeInfo {
212 hotkey: info.hotkey,
213 coldkey: info.coldkey,
214 netuid: info.netuid,
215 stake: Balance::from_rao(info.stake),
216 locked: Balance::from_rao(info.locked),
217 emission: Balance::from_rao(info.emission),
218 tao_emission: Balance::from_rao(info.tao_emission),
219 drain: Balance::from_rao(info.drain),
220 is_registered: info.is_registered,
221 })
222 .collect();
223
224 Ok(result)
225}
226
227#[cfg(test)]
228mod tests {
229 #[test]
230 fn test_balance_type() {
231 use crate::types::Balance;
232 let balance = Balance::from_rao(1_000_000_000);
233 assert_eq!(balance.as_tao(), 1.0);
234 }
235
236 #[test]
237 fn test_stake_info_struct() {
238 use super::*;
239 use subxt::utils::AccountId32;
240
241 let _info = StakeInfo {
242 hotkey: AccountId32::from([0u8; 32]),
243 coldkey: AccountId32::from([1u8; 32]),
244 netuid: 1,
245 stake: Balance::from_rao(1000),
246 locked: Balance::from_rao(0),
247 emission: Balance::from_rao(100),
248 tao_emission: Balance::from_rao(50),
249 drain: Balance::from_rao(0),
250 is_registered: true,
251 };
252 }
253}