Allow passing a custom database when creating the Service (paritytech…

…#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

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)
 }

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" }

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,43 +790,22 @@ 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)
 	}
 
 	/// Create new memory-backed client backend for tests.
 	#[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())

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 {

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.

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());