Allow passing a custom database when creating the Service (#3957)

* Put the DB configuration in an enum

* Allow passing a custom database to client-db

* Clean-ups in client-db

* Fix client tests

* Fix service tests

* Hopefully fix tests for good this time 😩

* Address review

substrate/Cargo.lock

@@ -4934,6 +4934,7 @@ dependencies = [
  "sr-primitives 2.0.0",
  "sr-std 2.0.0",
  "sr-version 2.0.0",
+ "substrate-client-db 2.0.0",
  "substrate-consensus-common 2.0.0",
  "substrate-executor 2.0.0",
  "substrate-header-metadata 2.0.0",

substrate/core/cli/src/lib.rs

@@ -28,7 +28,7 @@ pub mod informant;
 
 use client::ExecutionStrategies;
 use service::{
-	config::Configuration,
+	config::{Configuration, DatabaseConfig},
 	ServiceBuilderExport, ServiceBuilderImport, ServiceBuilderRevert,
 	RuntimeGenesis, ChainSpecExtension, PruningMode, ChainSpec,
 };
@@ -351,7 +351,9 @@ impl<'a> ParseAndPrepareExport<'a> {
 	{
 		let config = create_config_with_db_path(spec_factory, &self.params.shared_params, self.version)?;
 
-		info!("DB path: {}", config.database_path.display());
+		if let DatabaseConfig::Path { ref path, .. } = &config.database {
+			info!("DB path: {}", path.display());
+		}
 		let from = self.params.from.unwrap_or(1);
 		let to = self.params.to;
 		let json = self.params.json;
@@ -430,7 +432,13 @@ impl<'a> ParseAndPreparePurge<'a> {
 		let config = create_config_with_db_path::<(), _, _, _>(
 			spec_factory, &self.params.shared_params, self.version
 		)?;
-		let db_path = config.database_path;
+		let db_path = match config.database {
+			DatabaseConfig::Path { path, .. } => path,
+			_ => {
+				eprintln!("Cannot purge custom database implementation");
+				return Ok(());
+			}
+		};
 
 		if !self.params.yes {
 			print!("Are you sure to remove {:?}? [y/N]: ", &db_path);
@@ -455,7 +463,7 @@ impl<'a> ParseAndPreparePurge<'a> {
 				Ok(())
 			},
 			Result::Err(ref err) if err.kind() == ErrorKind::NotFound => {
-				println!("{:?} did not exist.", &db_path);
+				eprintln!("{:?} did not exist.", &db_path);
 				Ok(())
 			},
 			Result::Err(err) => Result::Err(err.into())
@@ -664,8 +672,10 @@ where
 		|| keystore_path(&base_path, config.chain_spec.id())
 	);
 
-	config.database_path = db_path(&base_path, config.chain_spec.id());
-	config.database_cache_size = cli.database_cache_size;
+	config.database = DatabaseConfig::Path {
+		path: db_path(&base_path, config.chain_spec.id()),
+		cache_size: cli.database_cache_size,
+	};
 	config.state_cache_size = cli.state_cache_size;
 
 	let is_dev = cli.shared_params.dev;
@@ -829,7 +839,10 @@ where
 	let base_path = base_path(cli, version);
 
 	let mut config = service::Configuration::default_with_spec(spec.clone());
-	config.database_path = db_path(&base_path, spec.id());
+	config.database = DatabaseConfig::Path {
+		path: db_path(&base_path, spec.id()),
+		cache_size: None,
+	};
 
 	Ok(config)
 }

substrate/core/client/Cargo.toml

@@ -32,6 +32,7 @@ header-metadata = { package = "substrate-header-metadata", path = "header-metada
 [dev-dependencies]
 env_logger = "0.7.0"
 tempfile = "3.1.0"
+client-db = { package = "substrate-client-db", path = "./db", features = ["kvdb-rocksdb"] }
 test-client = { package = "substrate-test-runtime-client", path = "../test-runtime/client" }
 kvdb-memorydb = { git = "https://github.com/paritytech/parity-common", rev="b0317f649ab2c665b7987b8475878fc4d2e1f81d" }
 panic-handler = { package = "substrate-panic-handler", path = "../panic-handler" }

substrate/core/client/db/src/lib.rs

@@ -83,6 +83,9 @@ const DEFAULT_CHILD_RATIO: (usize, usize) = (1, 10);
 /// DB-backed patricia trie state, transaction type is an overlay of changes to commit.
 pub type DbState = state_machine::TrieBackend<Arc<dyn state_machine::Storage<Blake2Hasher>>, Blake2Hasher>;
 
+/// Re-export the KVDB trait so that one can pass an implementation of it.
+pub use kvdb;
+
 /// A reference tracking state.
 ///
 /// It makes sure that the hash we are using stays pinned in storage
@@ -191,16 +194,28 @@ impl<B: BlockT> StateBackend<Blake2Hasher> for RefTrackingState<B> {
 
 /// Database settings.
 pub struct DatabaseSettings {
-	/// Cache size in bytes. If `None` default is used.
-	pub cache_size: Option<usize>,
 	/// State cache size.
 	pub state_cache_size: usize,
 	/// Ratio of cache size dedicated to child tries.
 	pub state_cache_child_ratio: Option<(usize, usize)>,
-	/// Path to the database.
-	pub path: PathBuf,
 	/// Pruning mode.
 	pub pruning: PruningMode,
+	/// Where to find the database.
+	pub source: DatabaseSettingsSrc,
+}
+
+/// Where to find the database..
+pub enum DatabaseSettingsSrc {
+	/// Load a database from a given path. Recommended for most uses.
+	Path {
+		/// Path to the database.
+		path: PathBuf,
+		/// Cache size in bytes. If `None` default is used.
+		cache_size: Option<usize>,
+	},
+
+	/// Use a custom already-open database.
+	Custom(Arc<dyn KeyValueDB>),
 }
 
 /// Create an instance of db-backed client.
@@ -775,17 +790,7 @@ impl<Block: BlockT<Hash=H256>> Backend<Block> {
 	///
 	/// The pruning window is how old a block must be before the state is pruned.
 	pub fn new(config: DatabaseSettings, canonicalization_delay: u64) -> ClientResult<Self> {
-		Self::new_inner(config, canonicalization_delay)
-	}
-
-	fn new_inner(config: DatabaseSettings, canonicalization_delay: u64) -> ClientResult<Self> {
-		#[cfg(feature = "kvdb-rocksdb")]
 		let db = crate::utils::open_database(&config, columns::META, "full")?;
-		#[cfg(not(feature = "kvdb-rocksdb"))]
-		let db = {
-			log::warn!("Running without the RocksDB feature. The database will NOT be saved.");
-			Arc::new(kvdb_memorydb::create(crate::utils::NUM_COLUMNS))
-		};
 		Self::from_kvdb(db as Arc<_>, canonicalization_delay, &config)
 	}
 
@@ -793,25 +798,14 @@ impl<Block: BlockT<Hash=H256>> Backend<Block> {
 	#[cfg(any(test, feature = "test-helpers"))]
 	pub fn new_test(keep_blocks: u32, canonicalization_delay: u64) -> Self {
 		let db = Arc::new(kvdb_memorydb::create(crate::utils::NUM_COLUMNS));
-		Self::new_test_db(keep_blocks, canonicalization_delay, db as Arc<_>)
-	}
-
-	/// Creates a client backend with test settings.
-	#[cfg(any(test, feature = "test-helpers"))]
-	pub fn new_test_db(keep_blocks: u32, canonicalization_delay: u64, db: Arc<dyn KeyValueDB>) -> Self {
-
 		let db_setting = DatabaseSettings {
-			cache_size: None,
 			state_cache_size: 16777216,
 			state_cache_child_ratio: Some((50, 100)),
-			path: Default::default(),
 			pruning: PruningMode::keep_blocks(keep_blocks),
+			source: DatabaseSettingsSrc::Custom(db),
 		};
-		Self::from_kvdb(
-			db,
-			canonicalization_delay,
-			&db_setting,
-		).expect("failed to create test-db")
+
+		Self::new(db_setting, canonicalization_delay).expect("failed to create test-db")
 	}
 
 	fn from_kvdb(
@@ -1636,7 +1630,12 @@ mod tests {
 			db.storage.db.clone()
 		};
 
-		let backend = Backend::<Block>::new_test_db(1, 0, backing);
+		let backend = Backend::<Block>::new(DatabaseSettings {
+			state_cache_size: 16777216,
+			state_cache_child_ratio: Some((50, 100)),
+			pruning: PruningMode::keep_blocks(1),
+			source: DatabaseSettingsSrc::Custom(backing),
+		}, 0).unwrap();
 		assert_eq!(backend.blockchain().info().best_number, 9);
 		for i in 0..10 {
 			assert!(backend.blockchain().hash(i).unwrap().is_some())

substrate/core/client/db/src/light.rs

@@ -70,22 +70,10 @@ impl<Block> LightStorage<Block>
 {
 	/// Create new storage with given settings.
 	pub fn new(config: DatabaseSettings) -> ClientResult<Self> {
-		Self::new_inner(config)
-	}
-
-	#[cfg(feature = "kvdb-rocksdb")]
-	fn new_inner(config: DatabaseSettings) -> ClientResult<Self> {
 		let db = crate::utils::open_database(&config, columns::META, "light")?;
 		Self::from_kvdb(db as Arc<_>)
 	}
 
-	#[cfg(not(feature = "kvdb-rocksdb"))]
-	fn new_inner(_config: DatabaseSettings) -> ClientResult<Self> {
-		log::warn!("Running without the RocksDB feature. The database will NOT be saved.");
-		let db = Arc::new(kvdb_memorydb::create(crate::utils::NUM_COLUMNS));
-		Self::from_kvdb(db as Arc<_>)
-	}
-
 	/// Create new memory-backed `LightStorage` for tests.
 	#[cfg(any(test, feature = "test-helpers"))]
 	pub fn new_test() -> Self {

substrate/core/client/db/src/utils.rs

@@ -17,7 +17,6 @@
 //! Db-based backend utility structures and functions, used by both
 //! full and light storages.
 
-#[cfg(feature = "kvdb-rocksdb")]
 use std::sync::Arc;
 use std::{io, convert::TryInto};
 
@@ -34,8 +33,7 @@ use sr_primitives::traits::{
 	Block as BlockT, Header as HeaderT, Zero,
 	UniqueSaturatedFrom, UniqueSaturatedInto,
 };
-#[cfg(feature = "kvdb-rocksdb")]
-use crate::DatabaseSettings;
+use crate::{DatabaseSettings, DatabaseSettingsSrc};
 
 /// Number of columns in the db. Must be the same for both full && light dbs.
 /// Otherwise RocksDb will fail to open database && check its type.
@@ -206,16 +204,26 @@ pub fn db_err(err: io::Error) -> client::error::Error {
 }
 
 /// Open RocksDB database.
-#[cfg(feature = "kvdb-rocksdb")]
 pub fn open_database(
 	config: &DatabaseSettings,
 	col_meta: Option<u32>,
 	db_type: &str
 ) -> client::error::Result<Arc<dyn KeyValueDB>> {
-	let mut db_config = DatabaseConfig::with_columns(Some(NUM_COLUMNS));
-	db_config.memory_budget = config.cache_size;
-	let path = config.path.to_str().ok_or_else(|| client::error::Error::Backend("Invalid database path".into()))?;
-	let db = Database::open(&db_config, &path).map_err(db_err)?;
+	let db: Arc<dyn KeyValueDB> = match &config.source {
+		#[cfg(feature = "kvdb-rocksdb")]
+		DatabaseSettingsSrc::Path { path, cache_size } => {
+			let mut db_config = DatabaseConfig::with_columns(Some(NUM_COLUMNS));
+			db_config.memory_budget = *cache_size;
+			let path = path.to_str().ok_or_else(|| client::error::Error::Backend("Invalid database path".into()))?;
+			Arc::new(Database::open(&db_config, &path).map_err(db_err)?)
+		},
+		#[cfg(not(feature = "kvdb-rocksdb"))]
+		DatabaseSettingsSrc::Path { .. } => {
+			let msg = "Try to open RocksDB database with RocksDB disabled".into();
+			return Err(client::error::Error::Backend(msg));
+		},
+		DatabaseSettingsSrc::Custom(db) => db.clone(),
+	};
 
 	// check database type
 	match db.get(col_meta, meta_keys::TYPE).map_err(db_err)? {
@@ -232,7 +240,7 @@ pub fn open_database(
 		},
 	}
 
-	Ok(Arc::new(db))
+	Ok(db)
 }
 
 /// Read database column entry for the given block.

substrate/core/client/src/client.rs

@@ -1871,7 +1871,7 @@ pub(crate) mod tests {
 	use consensus::{BlockOrigin, SelectChain};
 	use test_client::{
 		prelude::*,
-		client_db::{Backend, DatabaseSettings, PruningMode},
+		client_db::{Backend, DatabaseSettings, DatabaseSettingsSrc, PruningMode},
 		runtime::{self, Block, Transfer, RuntimeApi, TestAPI},
 	};
 
@@ -2755,11 +2755,13 @@ pub(crate) mod tests {
 		// states
 		let backend = Arc::new(Backend::new(
 			DatabaseSettings {
-				cache_size: None,
 				state_cache_size: 1 << 20,
 				state_cache_child_ratio: None,
-				path: tmp.path().into(),
 				pruning: PruningMode::ArchiveAll,
+				source: DatabaseSettingsSrc::Path {
+					path: tmp.path().into(),
+					cache_size: None,
+				}
 			},
 			u64::max_value(),
 		).unwrap());

substrate/core/service/src/builder.rs

@@ -17,7 +17,7 @@
 use crate::{Service, NetworkStatus, NetworkState, error::{self, Error}, DEFAULT_PROTOCOL_ID};
 use crate::{SpawnTaskHandle, start_rpc_servers, build_network_future, TransactionPoolAdapter};
 use crate::status_sinks;
-use crate::config::Configuration;
+use crate::config::{Configuration, DatabaseConfig};
 use client::{
 	BlockchainEvents, Client, runtime_api,
 	backend::RemoteBackend, light::blockchain::RemoteBlockchain,
@@ -157,15 +157,6 @@ where TGen: RuntimeGenesis, TCSExt: Extension {
 	>, Error> {
 		let keystore = Keystore::open(config.keystore_path.clone(), config.keystore_password.clone())?;
 
-		let db_settings = client_db::DatabaseSettings {
-			cache_size: None,
-			state_cache_size: config.state_cache_size,
-			state_cache_child_ratio:
-				config.state_cache_child_ratio.map(|v| (v, 100)),
-			path: config.database_path.clone(),
-			pruning: config.pruning.clone(),
-		};
-
 		let executor = NativeExecutor::<TExecDisp>::new(
 			config.wasm_method,
 			config.default_heap_pages,
@@ -177,14 +168,32 @@ where TGen: RuntimeGenesis, TCSExt: Extension {
 			.cloned()
 			.unwrap_or_default();
 
-		let (client, backend) = client_db::new_client(
-			db_settings,
-			executor,
-			&config.chain_spec,
-			fork_blocks,
-			config.execution_strategies.clone(),
-			Some(keystore.clone()),
-		)?;
+		let (client, backend) = {
+			let db_config = client_db::DatabaseSettings {
+				state_cache_size: config.state_cache_size,
+				state_cache_child_ratio:
+					config.state_cache_child_ratio.map(|v| (v, 100)),
+				pruning: config.pruning.clone(),
+				source: match &config.database {
+					DatabaseConfig::Path { path, cache_size } =>
+						client_db::DatabaseSettingsSrc::Path {
+							path: path.clone(),
+							cache_size: cache_size.clone().map(|u| u as usize),
+						},
+					DatabaseConfig::Custom(db) =>
+						client_db::DatabaseSettingsSrc::Custom(db.clone()),
+				},
+			};
+			
+			client_db::new_client(
+				db_config,
+				executor,
+				&config.chain_spec,
+				fork_blocks,
+				config.execution_strategies.clone(),
+				Some(keystore.clone()),
+			)?
+		};
 
 		let client = Arc::new(client);
 
@@ -229,21 +238,29 @@ where TGen: RuntimeGenesis, TCSExt: Extension {
 	>, Error> {
 		let keystore = Keystore::open(config.keystore_path.clone(), config.keystore_password.clone())?;
 
-		let db_settings = client_db::DatabaseSettings {
-			cache_size: config.database_cache_size.map(|u| u as usize),
-			state_cache_size: config.state_cache_size,
-			state_cache_child_ratio:
-				config.state_cache_child_ratio.map(|v| (v, 100)),
-			path: config.database_path.clone(),
-			pruning: config.pruning.clone(),
-		};
-
 		let executor = NativeExecutor::<TExecDisp>::new(
 			config.wasm_method,
 			config.default_heap_pages,
 		);
 
-		let db_storage = client_db::light::LightStorage::new(db_settings)?;
+		let db_storage = {
+			let db_settings = client_db::DatabaseSettings {
+				state_cache_size: config.state_cache_size,
+				state_cache_child_ratio:
+					config.state_cache_child_ratio.map(|v| (v, 100)),
+				pruning: config.pruning.clone(),
+				source: match &config.database {
+					DatabaseConfig::Path { path, cache_size } =>
+						client_db::DatabaseSettingsSrc::Path {
+							path: path.clone(),
+							cache_size: cache_size.clone().map(|u| u as usize),
+						},
+					DatabaseConfig::Custom(db) =>
+						client_db::DatabaseSettingsSrc::Custom(db.clone()),
+				},
+			};
+			client_db::light::LightStorage::new(db_settings)?
+		};
 		let light_blockchain = client::light::new_light_blockchain(db_storage);
 		let fetch_checker = Arc::new(client::light::new_fetch_checker(light_blockchain.clone(), executor.clone()));
 		let fetcher = Arc::new(network::OnDemand::new(fetch_checker));

substrate/core/service/src/config.rs

@@ -17,11 +17,11 @@
 //! Service configuration.
 
 pub use client::ExecutionStrategies;
-pub use client_db::PruningMode;
+pub use client_db::{kvdb::KeyValueDB, PruningMode};
 pub use network::config::{ExtTransport, NetworkConfiguration, Roles};
 pub use substrate_executor::WasmExecutionMethod;
 
-use std::{path::PathBuf, net::SocketAddr};
+use std::{path::PathBuf, net::SocketAddr, sync::Arc};
 use transaction_pool;
 use chain_spec::{ChainSpec, RuntimeGenesis, Extension, NoExtension};
 use primitives::crypto::Protected;
@@ -45,10 +45,8 @@ pub struct Configuration<C, G, E = NoExtension> {
 	pub network: NetworkConfiguration,
 	/// Path to key files.
 	pub keystore_path: PathBuf,
-	/// Path to the database.
-	pub database_path: PathBuf,
-	/// Cache Size for internal database in MiB
-	pub database_cache_size: Option<u32>,
+	/// Configuration for the database.
+	pub database: DatabaseConfig,
 	/// Size of internal state cache in Bytes
 	pub state_cache_size: usize,
 	/// Size in percent of cache size dedicated to child tries
@@ -100,6 +98,21 @@ pub struct Configuration<C, G, E = NoExtension> {
 	pub dev_key_seed: Option<String>,
 }
 
+/// Configuration of the database of the client.
+#[derive(Clone)]
+pub enum DatabaseConfig {
+	/// Database file at a specific path. Recommended for most uses.
+	Path {
+		/// Path to the database.
+		path: PathBuf,
+		/// Cache Size for internal database in MiB
+		cache_size: Option<u32>,
+	},
+
+	/// A custom implementation of an already-open database.
+	Custom(Arc<dyn KeyValueDB>),
+}
+
 impl<C, G, E> Configuration<C, G, E> where
 	C: Default,
 	G: RuntimeGenesis,
@@ -117,8 +130,10 @@ impl<C, G, E> Configuration<C, G, E> where
 			transaction_pool: Default::default(),
 			network: Default::default(),
 			keystore_path: Default::default(),
-			database_path: Default::default(),
-			database_cache_size: Default::default(),
+			database: DatabaseConfig::Path {
+				path: Default::default(),
+				cache_size: Default::default(),
+			},
 			state_cache_size: Default::default(),
 			state_cache_child_ratio: Default::default(),
 			custom: Default::default(),

substrate/core/service/test/src/lib.rs

@@ -29,6 +29,7 @@ use service::{
 	AbstractService,
 	ChainSpec,
 	Configuration,
+	config::DatabaseConfig,
 	Roles,
 	Error,
 };
@@ -170,8 +171,10 @@ fn node_config<G, E: Clone> (
 		network: network_config,
 		keystore_path: root.join("key"),
 		keystore_password: None,
-		database_path: root.join("db"),
-		database_cache_size: None,
+		database: DatabaseConfig::Path {
+			path: root.join("db"),
+			cache_size: None
+		},
 		state_cache_size: 16777216,
 		state_cache_child_ratio: None,
 		pruning: Default::default(),