@@ -43,10 +43,12 @@ type KeyPair = (StorageKey, StorageData);
4343
4444const  LOG_TARGET :  & str  = "remote-ext" ; 
4545const  DEFAULT_TARGET :  & str  = "wss://rpc.polkadot.io" ; 
46- const  BATCH_SIZE :  usize  = 512 ; 
46+ const  BATCH_SIZE :  usize  = 1000 ; 
4747
4848jsonrpsee_proc_macros:: rpc_client_api! { 
4949	RpcApi <B :  BlockT > { 
50+ 		#[ rpc( method = "state_getStorage" ,  positional_params) ] 
51+ 		fn  get_storage( prefix:  StorageKey ,  hash:  Option <B :: Hash >)  -> StorageData ; 
5052		#[ rpc( method = "state_getKeysPaged" ,  positional_params) ] 
5153		fn  get_keys_paged( 
5254			prefix:  Option <StorageKey >, 
@@ -107,7 +109,7 @@ impl From<String> for Transport {
107109/// A state snapshot config may be present and will be written to in that case. 
108110#[ derive( Clone ) ]  
109111pub  struct  OnlineConfig < B :  BlockT >  { 
110- 	/// The block number  at which to connect . Will be latest finalized head if not provided. 
112+ 	/// The block hash  at which to get the runtime state . Will be latest finalized head if not provided. 
111113 	pub  at :  Option < B :: Hash > , 
112114	/// An optional state snapshot file to WRITE to, not for reading. Not written if set to `None`. 
113115 	pub  state_snapshot :  Option < SnapshotConfig > , 
@@ -159,8 +161,11 @@ impl Default for SnapshotConfig {
159161pub  struct  Builder < B :  BlockT >  { 
160162	/// Custom key-pairs to be injected into the externalities. 
161163 	inject :  Vec < KeyPair > , 
162- 	/// Storage entry key prefixes to be injected into the externalities. The *hashed* prefix must be given. 
164+ 	/// Storage entry key prefixes to be injected into the externalities. The *hashed* prefix must 
165+  	/// be given. 
163166 	hashed_prefixes :  Vec < Vec < u8 > > , 
167+ 	/// Storage entry keys to be injected into the externalities. The *hashed* key must be given. 
168+  	hashed_keys :  Vec < Vec < u8 > > , 
164169	/// connectivity mode, online or offline. 
165170 	mode :  Mode < B > , 
166171} 
@@ -169,7 +174,12 @@ pub struct Builder<B: BlockT> {
169174// that. 
170175impl < B :  BlockT >  Default  for  Builder < B >  { 
171176	fn  default ( )  -> Self  { 
172- 		Self  {  inject :  Default :: default ( ) ,  mode :  Default :: default ( ) ,  hashed_prefixes :  Default :: default ( )  } 
177+ 		Self  { 
178+ 			inject :  Default :: default ( ) , 
179+ 			mode :  Default :: default ( ) , 
180+ 			hashed_prefixes :  Default :: default ( ) , 
181+ 			hashed_keys :  Default :: default ( ) , 
182+ 		} 
173183	} 
174184} 
175185
@@ -192,6 +202,17 @@ impl<B: BlockT> Builder<B> {
192202
193203// RPC methods 
194204impl < B :  BlockT >  Builder < B >  { 
205+ 	async  fn  rpc_get_storage ( 
206+ 		& self , 
207+ 		key :  StorageKey , 
208+ 		maybe_at :  Option < B :: Hash > , 
209+ 	)  -> Result < StorageData ,  & ' static  str >  { 
210+ 		trace ! ( target:  LOG_TARGET ,  "rpc: get_storage" ) ; 
211+ 		RpcApi :: < B > :: get_storage ( self . as_online ( ) . rpc_client ( ) ,  key,  maybe_at) . await . map_err ( |e| { 
212+ 			error ! ( "Error = {:?}" ,  e) ; 
213+ 			"rpc get_storage failed." 
214+ 		} ) 
215+ 	} 
195216	/// Get the latest finalized head. 
196217 	async  fn  rpc_get_head ( & self )  -> Result < B :: Hash ,  & ' static  str >  { 
197218		trace ! ( target:  LOG_TARGET ,  "rpc: finalized_head" ) ; 
@@ -281,7 +302,7 @@ impl<B: BlockT> Builder<B> {
281302			let  values = client. batch_request :: < Option < StorageData > > ( batch) 
282303				. await 
283304				. map_err ( |e| { 
284- 					log:: error!( target:  LOG_TARGET ,  "failed to execute batch {:?} due to  {:?}" ,  chunk_keys,  e) ; 
305+ 					log:: error!( target:  LOG_TARGET ,  "failed to execute batch:  {:?}. Error:  {:?}" ,  chunk_keys,  e) ; 
285306					"batch failed." 
286307				} ) ?; 
287308			assert_eq ! ( chunk_keys. len( ) ,  values. len( ) ) ; 
@@ -356,11 +377,23 @@ impl<B: BlockT> Builder<B> {
356377		} ; 
357378
358379		for  prefix in  & self . hashed_prefixes  { 
359- 			info ! ( target:  LOG_TARGET ,  "adding data for hashed prefix: {:?}" ,  HexDisplay :: from( prefix) ) ; 
360- 			let  additional_key_values = self . rpc_get_pairs_paged ( StorageKey ( prefix. to_vec ( ) ) ,  at) . await ?; 
380+ 			debug ! ( 
381+ 				target:  LOG_TARGET , 
382+ 				"adding data for hashed prefix: {:?}" , 
383+ 				HexDisplay :: from( prefix) 
384+ 			) ; 
385+ 			let  additional_key_values =
386+ 				self . rpc_get_pairs_paged ( StorageKey ( prefix. to_vec ( ) ) ,  at) . await ?; 
361387			keys_and_values. extend ( additional_key_values) ; 
362388		} 
363389
390+ 		for  key in  & self . hashed_keys  { 
391+ 			let  key = StorageKey ( key. to_vec ( ) ) ; 
392+ 			debug ! ( target:  LOG_TARGET ,  "adding data for hashed key: {:?}" ,  HexDisplay :: from( & key) ) ; 
393+ 			let  value = self . rpc_get_storage ( key. clone ( ) ,  Some ( at) ) . await ?; 
394+ 			keys_and_values. push ( ( key,  value) ) ; 
395+ 		} 
396+ 
364397		Ok ( keys_and_values) 
365398	} 
366399
@@ -400,7 +433,7 @@ impl<B: BlockT> Builder<B> {
400433
401434		info ! ( 
402435			target:  LOG_TARGET , 
403- 			"extending externalities with {} manually injected keys " , 
436+ 			"extending externalities with {} manually injected key-values " , 
404437			self . inject. len( ) 
405438		) ; 
406439		base_kv. extend ( self . inject . clone ( ) ) ; 
@@ -416,19 +449,29 @@ impl<B: BlockT> Builder<B> {
416449	} 
417450
418451	/// Inject a manual list of key and values to the storage. 
419-  	pub  fn  inject ( mut  self ,  injections :  & [ KeyPair ] )  -> Self  { 
452+  	pub  fn  inject_key_value ( mut  self ,  injections :  & [ KeyPair ] )  -> Self  { 
420453		for  i in  injections { 
421454			self . inject . push ( i. clone ( ) ) ; 
422455		} 
423456		self 
424457	} 
425458
426- 	/// Inject a hashed prefix. This is treated as-is, and should be pre-hashed.  
459+ 	/// Inject a hashed prefix. This is treated as-is, and should be pre-hashed. 
460+  	/// 
461+  	/// This should be used to inject a "PREFIX", like a storage (double) map. 
427462 	pub  fn  inject_hashed_prefix ( mut  self ,  hashed :  & [ u8 ] )  -> Self  { 
428463		self . hashed_prefixes . push ( hashed. to_vec ( ) ) ; 
429464		self 
430465	} 
431466
467+ 	/// Inject a hashed key to scrape. This is treated as-is, and should be pre-hashed. 
468+  	/// 
469+  	/// This should be used to inject a "KEY", like a storage value. 
470+  	pub  fn  inject_hashed_key ( mut  self ,  hashed :  & [ u8 ] )  -> Self  { 
471+ 		self . hashed_keys . push ( hashed. to_vec ( ) ) ; 
472+ 		self 
473+ 	} 
474+ 
432475	/// Configure a state snapshot to be used. 
433476 	pub  fn  mode ( mut  self ,  mode :  Mode < B > )  -> Self  { 
434477		self . mode  = mode; 
0 commit comments