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/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.