Fix optional store items. (#120)

* Fix optional store items. * Support querying a block hash.
paritytech · Jun 1, 2020 · 26ada75 · 26ada75
1 parent a2eead0
commit 26ada75
Show file tree

Hide file tree

Showing 8 changed files with 94 additions and 38 deletions.
diff --git a/proc-macro/src/store.rs b/proc-macro/src/store.rs
@@ -49,11 +49,16 @@ impl Parse for StoreAttr {
 
 type StoreAttrs = utils::Attrs<StoreAttr>;
 
-fn parse_returns_attr(attr: &syn::Attribute) -> Option<syn::Type> {
+fn parse_returns_attr(attr: &syn::Attribute) -> Option<(syn::Type, syn::Type, bool)> {
     let attrs: StoreAttrs = syn::parse2(attr.tokens.clone()).unwrap();
     attrs.attrs.into_iter().next().map(|attr| {
         let StoreAttr::Returns(attr) = attr;
-        attr.value
+        let ty = attr.value;
+        if let Some(inner) = utils::parse_option(&ty) {
+            (ty, inner, false)
+        } else {
+            (ty.clone(), ty, true)
+        }
     })
 }
 
@@ -72,11 +77,16 @@ pub fn store(s: Structure) -> TokenStream {
     let filtered_fields = utils::filter_fields(&fields, &marker);
     let args = utils::fields_to_args(&filtered_fields);
     let build_struct = utils::build_struct(ident, &fields);
-    let ret = bindings
+    let (ret, store_ret, uses_default) = bindings
         .iter()
         .filter_map(|bi| bi.ast().attrs.iter().filter_map(parse_returns_attr).next())
         .next()
         .expect("#[store(returns = ..)] needs to be specified.");
+    let fetch = if uses_default {
+        quote!(fetch_or_default)
+    } else {
+        quote!(fetch)
+    };
     let store_ty = format_ident!(
         "{}",
         match filtered_fields.len() {
@@ -94,7 +104,7 @@ pub fn store(s: Structure) -> TokenStream {
         impl#generics #subxt::Store<T> for #ident<#(#params),*> {
             const MODULE: &'static str = MODULE;
             const FIELD: &'static str = #store_name;
-            type Returns = #ret;
+            type Returns = #store_ret;
             fn key(
                 &self,
                 metadata: &#subxt::Metadata,
@@ -113,6 +123,7 @@ pub fn store(s: Structure) -> TokenStream {
             fn #store<'a>(
                 &'a self,
                 #args
+                hash: Option<T::Hash>,
             ) -> core::pin::Pin<Box<dyn core::future::Future<Output = Result<#ret, #subxt::Error>> + Send + 'a>>;
         }
 
@@ -125,9 +136,10 @@ pub fn store(s: Structure) -> TokenStream {
             fn #store<'a>(
                 &'a self,
                 #args
+                hash: Option<T::Hash>,
             ) -> core::pin::Pin<Box<dyn core::future::Future<Output = Result<#ret, #subxt::Error>> + Send + 'a>> {
                 let #marker = core::marker::PhantomData::<T>;
-                Box::pin(self.fetch(#build_struct, None))
+                Box::pin(self.#fetch(#build_struct, hash))
             }
         }
     }
@@ -169,6 +181,7 @@ mod tests {
                 fn account<'a>(
                     &'a self,
                     account_id: &'a <T as System>::AccountId,
+                    hash: Option<T::Hash>,
                 ) -> core::pin::Pin<Box<dyn core::future::Future<Output = Result<AccountData<T::Balance>, substrate_subxt::Error>> + Send + 'a>>;
             }
 
@@ -181,10 +194,11 @@ mod tests {
                 fn account<'a>(
                     &'a self,
                     account_id: &'a <T as System>::AccountId,
+                    hash: Option<T::Hash>,
                 ) -> core::pin::Pin<Box<dyn core::future::Future<Output = Result<AccountData<T::Balance>, substrate_subxt::Error>> + Send + 'a>>
                 {
                     let _ = core::marker::PhantomData::<T>;
-                    Box::pin(self.fetch(AccountStore { account_id, }, None))
+                    Box::pin(self.fetch_or_default(AccountStore { account_id, }, hash))
                 }
             }
         };

diff --git a/proc-macro/src/test.rs b/proc-macro/src/test.rs
@@ -330,7 +330,7 @@ impl Step {
                 };
                 let build_struct = quote! {
                     #(
-                        let #state_name = client.fetch(#state, None).await.unwrap();
+                        let #state_name = client.fetch_or_default(#state, None).await.unwrap();
                     )*
                     State { #(#state_name),* }
                 };
@@ -477,11 +477,11 @@ mod tests {
 
                     let pre = {
                         let alice = client
-                            .fetch(AccountStore { account_id: &alice }, None)
+                            .fetch_or_default(AccountStore { account_id: &alice }, None)
                             .await
                             .unwrap();
                         let bob = client
-                            .fetch(AccountStore { account_id: &bob }, None)
+                            .fetch_or_default(AccountStore { account_id: &bob }, None)
                             .await
                             .unwrap();
                         State { alice, bob }
@@ -510,11 +510,11 @@ mod tests {
 
                     let post = {
                         let alice = client
-                            .fetch(AccountStore { account_id: &alice }, None)
+                            .fetch_or_default(AccountStore { account_id: &alice }, None)
                             .await
                             .unwrap();
                         let bob = client
-                            .fetch(AccountStore { account_id: &bob }, None)
+                            .fetch_or_default(AccountStore { account_id: &bob }, None)
                             .await
                             .unwrap();
                         State { alice, bob }

diff --git a/proc-macro/src/utils.rs b/proc-macro/src/utils.rs
@@ -167,6 +167,21 @@ pub fn type_params(generics: &syn::Generics) -> Vec<TokenStream> {
         .collect()
 }
 
+pub fn parse_option(ty: &syn::Type) -> Option<syn::Type> {
+    if let syn::Type::Path(ty_path) = ty {
+        if let Some(seg) = ty_path.path.segments.first() {
+            if seg.ident.to_string() == "Option" {
+                if let syn::PathArguments::AngleBracketed(args) = &seg.arguments {
+                    if let Some(syn::GenericArgument::Type(ty)) = args.args.first() {
+                        return Some(ty.clone())
+                    }
+                }
+            }
+        }
+    }
+    None
+}
+
 #[derive(Debug)]
 pub struct Attrs<A> {
     pub paren: syn::token::Paren,
@@ -209,3 +224,16 @@ pub(crate) fn assert_proc_macro(
     let expected = expected.to_string();
     pretty_assertions::assert_eq!(result, expected);
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_parse_option() {
+        let option_t: syn::Type = syn::parse2(quote!(Option<T>)).unwrap();
+        let t: syn::Type = syn::parse2(quote!(T)).unwrap();
+        assert_eq!(parse_option(&option_t), Some(t.clone()));
+        assert_eq!(parse_option(&t), None);
+    }
+}
diff --git a/proc-macro/tests/balances.rs b/proc-macro/tests/balances.rs
@@ -117,8 +117,8 @@ async fn transfer_balance_example() -> Result<(), Box<dyn std::error::Error>> {
     let alice = AccountKeyring::Alice.to_account_id();
     let bob = AccountKeyring::Bob.to_account_id();
 
-    let alice_account = client.account(&alice).await?;
-    let bob_account = client.account(&bob).await?;
+    let alice_account = client.account(&alice, None).await?;
+    let bob_account = client.account(&bob, None).await?;
     let pre = (alice_account, bob_account);
 
     let _hash = client
@@ -138,8 +138,8 @@ async fn transfer_balance_example() -> Result<(), Box<dyn std::error::Error>> {
         })
     );
 
-    let alice_account = client.account(&alice).await?;
-    let bob_account = client.account(&bob).await?;
+    let alice_account = client.account(&alice, None).await?;
+    let bob_account = client.account(&bob, None).await?;
     let post = (alice_account, bob_account);
 
     assert_eq!(pre.0.free, post.0.free - 10_000);

diff --git a/src/frame/balances.rs b/src/frame/balances.rs
@@ -146,7 +146,7 @@ mod tests {
     async fn test_state_total_issuance() {
         env_logger::try_init().ok();
         let client = test_client().await;
-        let total_issuance = client.total_issuance().await.unwrap();
+        let total_issuance = client.total_issuance(None).await.unwrap();
         assert_ne!(total_issuance, 0);
     }
 
@@ -156,7 +156,7 @@ mod tests {
         env_logger::try_init().ok();
         let client = test_client().await;
         let account = AccountKeyring::Alice.to_account_id();
-        let info = client.account(&account).await.unwrap();
+        let info = client.account(&account, None).await.unwrap();
         assert_ne!(info.data.free, 0);
     }
 }
diff --git a/src/lib.rs b/src/lib.rs
@@ -45,7 +45,10 @@ extern crate substrate_subxt_proc_macro;
 pub use sp_core;
 pub use sp_runtime;
 
-use codec::Encode;
+use codec::{
+    Decode,
+    Encode,
+};
 use futures::future;
 use jsonrpsee::client::Subscription;
 use sc_rpc_api::state::ReadProof;
@@ -200,21 +203,34 @@ impl<T: System, S, E> Client<T, S, E> {
         &self.metadata
     }
 
-    /// Fetch a StorageKey.
-    pub async fn fetch<F: Store<T>>(
+    /// Fetch a StorageKey with default value.
+    pub async fn fetch_or_default<F: Store<T>>(
         &self,
         store: F,
         hash: Option<T::Hash>,
     ) -> Result<F::Returns, Error> {
         let key = store.key(&self.metadata)?;
-        let value = self.rpc.storage::<F::Returns>(key, hash).await?;
-        if let Some(v) = value {
-            Ok(v)
+        if let Some(data) = self.rpc.storage(key, hash).await? {
+            Ok(Decode::decode(&mut &data.0[..])?)
         } else {
             Ok(store.default(&self.metadata)?)
         }
     }
 
+    /// Fetch a StorageKey an optional storage key.
+    pub async fn fetch<F: Store<T>>(
+        &self,
+        store: F,
+        hash: Option<T::Hash>,
+    ) -> Result<Option<F::Returns>, Error> {
+        let key = store.key(&self.metadata)?;
+        if let Some(data) = self.rpc.storage(key, hash).await? {
+            Ok(Some(Decode::decode(&mut &data.0[..])?))
+        } else {
+            Ok(None)
+        }
+    }
+
     /// Query historical storage entries
     pub async fn query_storage(
         &self,
@@ -312,7 +328,7 @@ where
         let account_nonce = if let Some(nonce) = nonce {
             nonce
         } else {
-            self.account(account_id).await?.nonce
+            self.account(account_id, None).await?.nonce
         };
         let spec_version = self.runtime_version.spec_version;
         let tx_version = self.runtime_version.transaction_version;
@@ -453,7 +469,7 @@ mod tests {
 
         let client = test_client().await;
         let nonce = client
-            .account(&AccountKeyring::Alice.to_account_id())
+            .account(&AccountKeyring::Alice.to_account_id(), None)
             .await
             .unwrap()
             .nonce;

diff --git a/src/rpc.rs b/src/rpc.rs
@@ -132,22 +132,15 @@ impl<T: System> Rpc<T> {
     }
 
     /// Fetch a storage key
-    pub async fn storage<V: Decode>(
+    pub async fn storage(
         &self,
         key: StorageKey,
         hash: Option<T::Hash>,
-    ) -> Result<Option<V>, Error> {
+    ) -> Result<Option<StorageData>, Error> {
         let params = Params::Array(vec![to_json_value(key)?, to_json_value(hash)?]);
-        let data: Option<StorageData> =
-            self.client.request("state_getStorage", params).await?;
-        match data {
-            Some(data) => {
-                log::debug!("state_getStorage {:?}", data.0);
-                let value = Decode::decode(&mut &data.0[..])?;
-                Ok(Some(value))
-            }
-            None => Ok(None),
-        }
+        let data = self.client.request("state_getStorage", params).await?;
+        log::debug!("state_getStorage {:?}", data);
+        Ok(data)
     }
 
     /// Query historical storage entries

diff --git a/src/signer.rs b/src/signer.rs
@@ -104,10 +104,15 @@ where
         self.nonce = Some(nonce);
     }
 
-    /// Increment the nonce
+    /// Increment the nonce.
     pub fn increment_nonce(&mut self) {
         self.nonce = self.nonce.map(|nonce| nonce + 1.into());
     }
+
+    /// Returns the signer.
+    pub fn signer(&self) -> &P {
+        &self.signer
+    }
 }
 
 impl<T, S, E, P> Signer<T, S, E> for PairSigner<T, S, E, P>