From 663934ca374983308fa3cceccbbb58f59b6785b3 Mon Sep 17 00:00:00 2001
From: Andrew Jones <ascjones@gmail.com>
Date: Mon, 3 Aug 2020 12:27:31 +0100
Subject: [PATCH] Fetch paged storage keys

---
 proc-macro/src/store.rs | 24 ++++++++++++++++++++++--
 src/frame/mod.rs        |  2 ++
 src/lib.rs              | 21 +++++++++++++++++++++
 src/metadata.rs         | 10 +++++-----
 src/rpc.rs              | 21 +++++++++++++++++++++
 5 files changed, 71 insertions(+), 7 deletions(-)

diff --git a/proc-macro/src/store.rs b/proc-macro/src/store.rs
index 47294589a1..5dcabe0d53 100644
--- a/proc-macro/src/store.rs
+++ b/proc-macro/src/store.rs
@@ -116,6 +116,16 @@ pub fn store(s: Structure) -> TokenStream {
             const MODULE: &'static str = MODULE;
             const FIELD: &'static str = #store_name;
             type Returns = #store_ret;
+
+            fn prefix(
+                metadata: &#subxt::Metadata,
+            ) -> Result<#subxt::sp_core::storage::StorageKey, #subxt::MetadataError> {
+                Ok(metadata
+                    .module(Self::MODULE)?
+                    .storage(Self::FIELD)?
+                    .prefix())
+            }
+
             fn key(
                 &self,
                 metadata: &#subxt::Metadata,
@@ -130,7 +140,7 @@ pub fn store(s: Structure) -> TokenStream {
 
         /// Store extension trait.
         pub trait #store_trait<T: #module> {
-            /// Retrive the store element.
+            /// Retrieve the store element.
             fn #store<'a>(
                 &'a self,
                 #args
@@ -169,6 +179,16 @@ mod tests {
                 const MODULE: &'static str = MODULE;
                 const FIELD: &'static str = "Account";
                 type Returns = AccountData<T::Balance>;
+
+                fn prefix(
+                    metadata: &substrate_subxt::Metadata,
+                ) -> Result<substrate_subxt::sp_core::storage::StorageKey, substrate_subxt::MetadataError> {
+                    Ok(metadata
+                        .module(Self::MODULE)?
+                        .storage(Self::FIELD)?
+                        .prefix())
+                }
+
                 fn key(
                     &self,
                     metadata: &substrate_subxt::Metadata,
@@ -183,7 +203,7 @@ mod tests {
 
             /// Store extension trait.
             pub trait AccountStoreExt<T: Balances> {
-                /// Retrive the store element.
+                /// Retrieve the store element.
                 fn account<'a>(
                     &'a self,
                     account_id: &'a <T as System>::AccountId,
diff --git a/src/frame/mod.rs b/src/frame/mod.rs
index ab0fbf6f24..1d7e262ed7 100644
--- a/src/frame/mod.rs
+++ b/src/frame/mod.rs
@@ -42,6 +42,8 @@ pub trait Store<T>: Encode {
     const FIELD: &'static str;
     /// Return type.
     type Returns: Decode;
+    /// Returns the key prefix for storage maps
+    fn prefix(metadata: &Metadata) -> Result<StorageKey, MetadataError>;
     /// Returns the `StorageKey`.
     fn key(&self, metadata: &Metadata) -> Result<StorageKey, MetadataError>;
     /// Returns the default value.
diff --git a/src/lib.rs b/src/lib.rs
index 5d8148e353..b3c823b235 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -228,6 +228,20 @@ impl<T: Runtime> Client<T> {
         }
     }
 
+	/// Fetch up to `count` keys for a storage map in lexicographic order.
+	///
+	/// Supports pagination by passing a value to `start_key`.
+    pub async fn fetch_keys<F: Store<T>>(
+        &self,
+        count: u32,
+        start_key: Option<StorageKey>,
+        hash: Option<T::Hash>,
+    ) -> Result<Vec<StorageKey>, Error> {
+        let prefix = <F as Store<T>>::prefix(&self.metadata)?;
+        let keys = self.rpc.storage_keys_paged(Some(prefix), count, start_key, hash).await?;
+        Ok(keys)
+    }
+
     /// Query historical storage entries
     pub async fn query_storage(
         &self,
@@ -612,4 +626,11 @@ mod tests {
         let mut blocks = client.subscribe_finalized_blocks().await.unwrap();
         blocks.next().await;
     }
+
+    #[async_std::test]
+    async fn test_fetch_keys() {
+        let (client, _) = test_client().await;
+        let keys = client.fetch_keys::<system::AccountStore<_>>(4, None, None).await.unwrap();
+        assert_eq!(keys.len(), 4)
+    }
 }
diff --git a/src/metadata.rs b/src/metadata.rs
index 899aed4c23..bbaf478d8a 100644
--- a/src/metadata.rs
+++ b/src/metadata.rs
@@ -249,10 +249,10 @@ pub struct StorageMetadata {
 }
 
 impl StorageMetadata {
-    pub fn prefix(&self) -> Vec<u8> {
+    pub fn prefix(&self) -> StorageKey {
         let mut bytes = sp_core::twox_128(self.module_prefix.as_bytes()).to_vec();
         bytes.extend(&sp_core::twox_128(self.storage_prefix.as_bytes())[..]);
-        bytes
+        StorageKey(bytes)
     }
 
     pub fn default<V: Decode>(&self) -> Result<V, MetadataError> {
@@ -286,7 +286,7 @@ impl StorageMetadata {
         match &self.ty {
             StorageEntryType::Plain(_) => {
                 Ok(StoragePlain {
-                    prefix: self.prefix(),
+                    prefix: self.prefix().0,
                 })
             }
             _ => Err(MetadataError::StorageTypeError),
@@ -298,7 +298,7 @@ impl StorageMetadata {
             StorageEntryType::Map { hasher, .. } => {
                 Ok(StorageMap {
                     _marker: PhantomData,
-                    prefix: self.prefix(),
+                    prefix: self.prefix().0,
                     hasher: hasher.clone(),
                 })
             }
@@ -317,7 +317,7 @@ impl StorageMetadata {
             } => {
                 Ok(StorageDoubleMap {
                     _marker: PhantomData,
-                    prefix: self.prefix(),
+                    prefix: self.prefix().0,
                     hasher1: hasher.clone(),
                     hasher2: key2_hasher.clone(),
                 })
diff --git a/src/rpc.rs b/src/rpc.rs
index 08662d0ad6..ecb2a055bf 100644
--- a/src/rpc.rs
+++ b/src/rpc.rs
@@ -131,6 +131,27 @@ impl<T: Runtime> Rpc<T> {
         Ok(data)
     }
 
+    /// Returns the keys with prefix with pagination support.
+	/// Up to `count` keys will be returned.
+	/// If `start_key` is passed, return next keys in storage in lexicographic order.
+    pub async fn storage_keys_paged(
+        &self,
+        prefix: Option<StorageKey>,
+        count: u32,
+        start_key: Option<StorageKey>,
+        hash: Option<T::Hash>,
+    ) -> Result<Vec<StorageKey>, Error> {
+        let params = Params::Array(vec![
+            to_json_value(prefix)?,
+            to_json_value(count)?,
+            to_json_value(start_key)?,
+            to_json_value(hash)?,
+        ]);
+        let data = self.client.request("state_getKeysPaged", params).await?;
+        log::debug!("state_getKeysPaged {:?}", data);
+        Ok(data)
+    }
+
     /// Query historical storage entries
     pub async fn query_storage(
         &self,