Add ability to list all created dbs

Signed-off-by: Victor Porof <victor.porof@gmail.com>

src/backend/impl_lmdb/environment.rs

@@ -29,12 +29,16 @@ use crate::backend::traits::{
     BackendEnvironment,
     BackendEnvironmentBuilder,
     BackendInfo,
+    BackendIter,
+    BackendRoCursor,
+    BackendRoCursorTransaction,
     BackendStat,
 };
 
 #[derive(Debug, PartialEq, Eq, Copy, Clone)]
 pub struct EnvironmentBuilderImpl {
     builder: lmdb::EnvironmentBuilder,
+    envtype: EnvironmentType,
     make_dir: bool,
 }
 
@@ -46,6 +50,7 @@ impl<'b> BackendEnvironmentBuilder<'b> for EnvironmentBuilderImpl {
     fn new() -> EnvironmentBuilderImpl {
         EnvironmentBuilderImpl {
             builder: lmdb::Environment::new(),
+            envtype: EnvironmentType::SingleDatabase,
             make_dir: false,
         }
     }
@@ -65,6 +70,9 @@ impl<'b> BackendEnvironmentBuilder<'b> for EnvironmentBuilderImpl {
 
     fn set_max_dbs(&mut self, max_dbs: u32) -> &mut Self {
         self.builder.set_max_dbs(max_dbs);
+        if max_dbs > 0 {
+            self.envtype = EnvironmentType::MultipleNamedDatabases
+        }
         self
     }
 
@@ -85,12 +93,18 @@ impl<'b> BackendEnvironmentBuilder<'b> for EnvironmentBuilderImpl {
             }
             fs::create_dir_all(path).map_err(ErrorImpl::IoError)?;
         }
-        self.builder.open(path).map(EnvironmentImpl).map_err(ErrorImpl::LmdbError)
+        self.builder.open(path).map(|env| EnvironmentImpl(env, self.envtype)).map_err(ErrorImpl::LmdbError)
     }
 }
 
+#[derive(Debug, PartialEq, Eq, Copy, Clone)]
+enum EnvironmentType {
+    SingleDatabase,
+    MultipleNamedDatabases,
+}
+
 #[derive(Debug)]
-pub struct EnvironmentImpl(lmdb::Environment);
+pub struct EnvironmentImpl(lmdb::Environment, EnvironmentType);
 
 impl<'e> BackendEnvironment<'e> for EnvironmentImpl {
     type Database = DatabaseImpl;
@@ -101,6 +115,23 @@ impl<'e> BackendEnvironment<'e> for EnvironmentImpl {
     type RwTransaction = RwTransactionImpl<'e>;
     type Stat = StatImpl;
 
+    fn get_dbs(&self) -> Result<Vec<Option<String>>, Self::Error> {
+        if self.1 == EnvironmentType::SingleDatabase {
+            return Ok(vec![None]);
+        }
+        let db = self.0.open_db(None).map(DatabaseImpl).map_err(ErrorImpl::LmdbError)?;
+        let reader = self.begin_ro_txn()?;
+        let cursor = reader.open_ro_cursor(&db)?;
+        let mut iter = cursor.into_iter();
+        let mut store = vec![];
+        while let Some(result) = iter.next() {
+            let (key, _) = result?;
+            let name = String::from_utf8(key.to_owned()).map_err(|_| ErrorImpl::LmdbError(lmdb::Error::Corrupted))?;
+            store.push(Some(name));
+        }
+        Ok(store)
+    }
+
     fn open_db(&self, name: Option<&str>) -> Result<Self::Database, Self::Error> {
         self.0.open_db(name).map(DatabaseImpl).map_err(ErrorImpl::LmdbError)
     }

src/backend/impl_safe/environment.rs

@@ -211,6 +211,11 @@ impl<'e> BackendEnvironment<'e> for EnvironmentImpl {
     type RwTransaction = RwTransactionImpl<'e>;
     type Stat = StatImpl;
 
+    fn get_dbs(&self) -> Result<Vec<Option<String>>, Self::Error> {
+        let dbs = self.dbs.read().map_err(|_| ErrorImpl::EnvPoisonError)?;
+        Ok(dbs.keys().map(|key| key.to_owned()).collect())
+    }
+
     fn open_db(&self, name: Option<&str>) -> Result<Self::Database, Self::Error> {
         if Arc::strong_count(&self.ro_txns) > 1 {
             return Err(ErrorImpl::DbsIllegalOpen);

src/backend/traits.rs

@@ -102,6 +102,8 @@ pub trait BackendEnvironment<'e>: Debug {
     type RoTransaction: BackendRoCursorTransaction<'e, Database = Self::Database>;
     type RwTransaction: BackendRwCursorTransaction<'e, Database = Self::Database>;
 
+    fn get_dbs(&self) -> Result<Vec<Option<String>>, Self::Error>;
+
     fn open_db(&self, name: Option<&str>) -> Result<Self::Database, Self::Error>;
 
     fn create_db(&self, name: Option<&str>, flags: Self::Flags) -> Result<Self::Database, Self::Error>;

src/env.rs

@@ -110,6 +110,11 @@ impl<'e, E> Rkv<E>
 where
     E: BackendEnvironment<'e>,
 {
+    /// Return all created databases.
+    pub fn get_dbs(&self) -> Result<Vec<Option<String>>, StoreError> {
+        self.env.get_dbs().map_err(|e| e.into())
+    }
+
     /// Create or Open an existing database in (&[u8] -> Single Value) mode.
     /// Note: that create=true cannot be called concurrently with other operations so if
     /// you are sure that the database exists, call this with create=false.

tests/env-lmdb.rs

@@ -194,6 +194,51 @@ fn test_open_with_capacity_2() {
     let _zzz = k.open_single(None, StoreOptions::default()).expect("opened");
 }
 
+#[test]
+fn test_list_dbs_1() {
+    let root = Builder::new().prefix("test_list_dbs").tempdir().expect("tempdir");
+    println!("Root path: {:?}", root.path());
+    fs::create_dir_all(root.path()).expect("dir created");
+    assert!(root.path().is_dir());
+
+    let k = Rkv::with_capacity::<Lmdb>(root.path(), 1).expect("rkv");
+    check_rkv(&k);
+
+    let dbs = k.get_dbs().unwrap();
+    assert_eq!(dbs, vec![Some("s".to_owned())]);
+}
+
+#[test]
+fn test_list_dbs_2() {
+    let root = Builder::new().prefix("test_list_dbs").tempdir().expect("tempdir");
+    println!("Root path: {:?}", root.path());
+    fs::create_dir_all(root.path()).expect("dir created");
+    assert!(root.path().is_dir());
+
+    let k = Rkv::with_capacity::<Lmdb>(root.path(), 2).expect("rkv");
+    check_rkv(&k);
+
+    let _ = k.open_single("zzz", StoreOptions::create()).expect("opened");
+
+    let dbs = k.get_dbs().unwrap();
+    assert_eq!(dbs, vec![Some("s".to_owned()), Some("zzz".to_owned())]);
+}
+
+#[test]
+fn test_list_dbs_3() {
+    let root = Builder::new().prefix("test_list_dbs").tempdir().expect("tempdir");
+    println!("Root path: {:?}", root.path());
+    fs::create_dir_all(root.path()).expect("dir created");
+    assert!(root.path().is_dir());
+
+    let k = Rkv::with_capacity::<Lmdb>(root.path(), 0).expect("rkv");
+
+    let _ = k.open_single(None, StoreOptions::create()).expect("opened");
+
+    let dbs = k.get_dbs().unwrap();
+    assert_eq!(dbs, vec![None]);
+}
+
 fn get_larger_than_default_map_size_value() -> usize {
     // The LMDB C library and lmdb Rust crate docs for setting the map size
     // <http://www.lmdb.tech/doc/group__mdb.html#gaa2506ec8dab3d969b0e609cd82e619e5>

tests/env-safe.rs

@@ -185,6 +185,54 @@ fn test_open_with_capacity_safe_2() {
     let _zzz = k.open_single(None, StoreOptions::default()).expect("opened");
 }
 
+#[test]
+fn test_list_dbs_safe_1() {
+    let root = Builder::new().prefix("test_list_dbs_safe").tempdir().expect("tempdir");
+    println!("Root path: {:?}", root.path());
+    fs::create_dir_all(root.path()).expect("dir created");
+    assert!(root.path().is_dir());
+
+    let k = Rkv::with_capacity::<SafeMode>(root.path(), 1).expect("rkv");
+    check_rkv(&k);
+
+    let mut dbs = k.get_dbs().unwrap();
+    dbs.sort();
+    assert_eq!(dbs, vec![None, Some("s".to_owned())]);
+}
+
+#[test]
+fn test_list_dbs_safe_2() {
+    let root = Builder::new().prefix("test_list_dbs_safe").tempdir().expect("tempdir");
+    println!("Root path: {:?}", root.path());
+    fs::create_dir_all(root.path()).expect("dir created");
+    assert!(root.path().is_dir());
+
+    let k = Rkv::with_capacity::<SafeMode>(root.path(), 2).expect("rkv");
+    check_rkv(&k);
+
+    let _ = k.open_single("zzz", StoreOptions::create()).expect("opened");
+
+    let mut dbs = k.get_dbs().unwrap();
+    dbs.sort();
+    assert_eq!(dbs, vec![None, Some("s".to_owned()), Some("zzz".to_owned())]);
+}
+
+#[test]
+fn test_list_dbs_safe_3() {
+    let root = Builder::new().prefix("test_list_dbs_safe").tempdir().expect("tempdir");
+    println!("Root path: {:?}", root.path());
+    fs::create_dir_all(root.path()).expect("dir created");
+    assert!(root.path().is_dir());
+
+    let k = Rkv::with_capacity::<SafeMode>(root.path(), 0).expect("rkv");
+
+    let _ = k.open_single(None, StoreOptions::create()).expect("opened");
+
+    let mut dbs = k.get_dbs().unwrap();
+    dbs.sort();
+    assert_eq!(dbs, vec![None]);
+}
+
 #[test]
 fn test_round_trip_and_transactions_safe() {
     let root = Builder::new().prefix("test_round_trip_and_transactions_safe").tempdir().expect("tempdir");