pezkuwi-sdk/bizinikiwi/client/rpc-spec-v2/src/common/storage.rs

// This file is part of Bizinikiwi.

// Copyright (C) Parity Technologies (UK) Ltd.
// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0

// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.

// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with this program. If not, see <https://www.gnu.org/licenses/>.

//! Storage queries for the RPC-V2 spec.

use std::{marker::PhantomData, sync::Arc};

use pezsc_client_api::{Backend, ChildInfo, StorageKey, StorageProvider};
use pezsp_runtime::traits::Block as BlockT;
use tokio::sync::mpsc;

use super::events::{StorageQuery, StorageQueryType, StorageResult, StorageResultType};
use crate::hex_string;

/// Call into the storage of blocks.
pub struct Storage<Client, Block, BE> {
	/// Bizinikiwi client.
	client: Arc<Client>,
	_phandom: PhantomData<(BE, Block)>,
}

impl<Client, Block, BE> Clone for Storage<Client, Block, BE> {
	fn clone(&self) -> Self {
		Self { client: self.client.clone(), _phandom: PhantomData }
	}
}

impl<Client, Block, BE> Storage<Client, Block, BE> {
	/// Constructs a new [`Storage`].
	pub fn new(client: Arc<Client>) -> Self {
		Self { client, _phandom: PhantomData }
	}
}

/// Query to iterate over storage.
#[derive(Debug)]
pub struct QueryIter {
	/// The key from which the iteration was started.
	pub query_key: StorageKey,
	/// The key after which pagination should resume.
	pub pagination_start_key: Option<StorageKey>,
	/// The type of the query (either value or hash).
	pub ty: IterQueryType,
}

/// The query type of an iteration.
#[derive(Debug)]
pub enum IterQueryType {
	/// Iterating over (key, value) pairs.
	Value,
	/// Iterating over (key, hash) pairs.
	Hash,
}

/// The result of making a query call.
pub type QueryResult = Result<Option<StorageResult>, String>;

impl<Client, Block, BE> Storage<Client, Block, BE>
where
	Block: BlockT + 'static,
	BE: Backend<Block> + 'static,
	Client: StorageProvider<Block, BE> + 'static,
{
	/// Fetch the value from storage.
	pub fn query_value(
		&self,
		hash: Block::Hash,
		key: &StorageKey,
		child_key: Option<&ChildInfo>,
	) -> QueryResult {
		let result = if let Some(child_key) = child_key {
			self.client.child_storage(hash, child_key, key)
		} else {
			self.client.storage(hash, key)
		};

		result
			.map(|opt| {
				QueryResult::Ok(opt.map(|storage_data| StorageResult {
					key: hex_string(&key.0),
					result: StorageResultType::Value(hex_string(&storage_data.0)),
					child_trie_key: child_key.map(|c| hex_string(&c.storage_key())),
				}))
			})
			.unwrap_or_else(|error| QueryResult::Err(error.to_string()))
	}

	/// Fetch the hash of a value from storage.
	pub fn query_hash(
		&self,
		hash: Block::Hash,
		key: &StorageKey,
		child_key: Option<&ChildInfo>,
	) -> QueryResult {
		let result = if let Some(child_key) = child_key {
			self.client.child_storage_hash(hash, child_key, key)
		} else {
			self.client.storage_hash(hash, key)
		};

		result
			.map(|opt| {
				QueryResult::Ok(opt.map(|storage_data| StorageResult {
					key: hex_string(&key.0),
					result: StorageResultType::Hash(hex_string(&storage_data.as_ref())),
					child_trie_key: child_key.map(|c| hex_string(&c.storage_key())),
				}))
			})
			.unwrap_or_else(|error| QueryResult::Err(error.to_string()))
	}

	/// Fetch the closest merkle value.
	pub fn query_merkle_value(
		&self,
		hash: Block::Hash,
		key: &StorageKey,
		child_key: Option<&ChildInfo>,
	) -> QueryResult {
		let result = if let Some(ref child_key) = child_key {
			self.client.child_closest_merkle_value(hash, child_key, key)
		} else {
			self.client.closest_merkle_value(hash, key)
		};

		result
			.map(|opt| {
				QueryResult::Ok(opt.map(|storage_data| {
					let result = match &storage_data {
						pezsc_client_api::MerkleValue::Node(data) => hex_string(&data.as_slice()),
						pezsc_client_api::MerkleValue::Hash(hash) => hex_string(&hash.as_ref()),
					};

					StorageResult {
						key: hex_string(&key.0),
						result: StorageResultType::ClosestDescendantMerkleValue(result),
						child_trie_key: child_key.map(|c| hex_string(&c.storage_key())),
					}
				}))
			})
			.unwrap_or_else(|error| QueryResult::Err(error.to_string()))
	}

	/// Iterate over the storage keys and send the results to the provided sender.
	///
	/// Because this relies on a bounded channel, it will pause the storage iteration
	// if the channel is becomes full which in turn provides backpressure.
	pub fn query_iter_pagination_with_producer(
		&self,
		query: QueryIter,
		hash: Block::Hash,
		child_key: Option<&ChildInfo>,
		tx: &mpsc::Sender<QueryResult>,
	) {
		let QueryIter { ty, query_key, pagination_start_key } = query;

		let maybe_storage = if let Some(child_key) = child_key {
			self.client.child_storage_keys(
				hash,
				child_key.to_owned(),
				Some(&query_key),
				pagination_start_key.as_ref(),
			)
		} else {
			self.client.storage_keys(hash, Some(&query_key), pagination_start_key.as_ref())
		};

		let keys_iter = match maybe_storage {
			Ok(keys_iter) => keys_iter,
			Err(error) => {
				_ = tx.blocking_send(Err(error.to_string()));
				return;
			},
		};

		for key in keys_iter {
			let result = match ty {
				IterQueryType::Value => self.query_value(hash, &key, child_key),
				IterQueryType::Hash => self.query_hash(hash, &key, child_key),
			};

			if tx.blocking_send(result).is_err() {
				break;
			}
		}
	}

	/// Raw iterator over the keys.
	pub fn raw_keys_iter(
		&self,
		hash: Block::Hash,
		child_key: Option<ChildInfo>,
	) -> Result<impl Iterator<Item = StorageKey>, String> {
		let keys_iter = if let Some(child_key) = child_key {
			self.client.child_storage_keys(hash, child_key, None, None)
		} else {
			self.client.storage_keys(hash, None, None)
		};

		keys_iter.map_err(|err| err.to_string())
	}
}

/// Generates storage events for `chainHead_storage` and `archive_storage` subscriptions.
pub struct StorageSubscriptionClient<Client, Block, BE> {
	/// Storage client.
	client: Storage<Client, Block, BE>,
	_phandom: PhantomData<(BE, Block)>,
}

impl<Client, Block, BE> Clone for StorageSubscriptionClient<Client, Block, BE> {
	fn clone(&self) -> Self {
		Self { client: self.client.clone(), _phandom: PhantomData }
	}
}

impl<Client, Block, BE> StorageSubscriptionClient<Client, Block, BE> {
	/// Constructs a new [`StorageSubscriptionClient`].
	pub fn new(client: Arc<Client>) -> Self {
		Self { client: Storage::new(client), _phandom: PhantomData }
	}
}

impl<Client, Block, BE> StorageSubscriptionClient<Client, Block, BE>
where
	Block: BlockT + 'static,
	BE: Backend<Block> + 'static,
	Client: StorageProvider<Block, BE> + Send + Sync + 'static,
{
	/// Generate storage events to the provided sender.
	pub async fn generate_events(
		&mut self,
		hash: Block::Hash,
		items: Vec<StorageQuery<StorageKey>>,
		child_key: Option<ChildInfo>,
		tx: mpsc::Sender<QueryResult>,
	) -> Result<(), tokio::task::JoinError> {
		let this = self.clone();

		tokio::task::spawn_blocking(move || {
			for item in items {
				match item.query_type {
					StorageQueryType::Value => {
						let rp = this.client.query_value(hash, &item.key, child_key.as_ref());
						if tx.blocking_send(rp).is_err() {
							break;
						}
					},
					StorageQueryType::Hash => {
						let rp = this.client.query_hash(hash, &item.key, child_key.as_ref());
						if tx.blocking_send(rp).is_err() {
							break;
						}
					},
					StorageQueryType::ClosestDescendantMerkleValue => {
						let rp =
							this.client.query_merkle_value(hash, &item.key, child_key.as_ref());
						if tx.blocking_send(rp).is_err() {
							break;
						}
					},
					StorageQueryType::DescendantsValues => {
						let query = QueryIter {
							query_key: item.key,
							ty: IterQueryType::Value,
							pagination_start_key: item.pagination_start_key,
						};
						this.client.query_iter_pagination_with_producer(
							query,
							hash,
							child_key.as_ref(),
							&tx,
						)
					},
					StorageQueryType::DescendantsHashes => {
						let query = QueryIter {
							query_key: item.key,
							ty: IterQueryType::Hash,
							pagination_start_key: item.pagination_start_key,
						};
						this.client.query_iter_pagination_with_producer(
							query,
							hash,
							child_key.as_ref(),
							&tx,
						)
					},
				}
			}
		})
		.await?;

		Ok(())
	}
}