Allow creating dirs if they don't exist

Signed-off-by: Victor Porof <[email protected]>

src/backend/impl_lmdb/environment.rs

@@ -8,7 +8,10 @@
 // CONDITIONS OF ANY KIND, either express or implied. See the License for the
 // specific language governing permissions and limitations under the License.
 
-use std::path::Path;
+use std::{
+    fs,
+    path::Path,
+};
 
 use super::{
     DatabaseFlagsImpl,
@@ -26,42 +29,59 @@ use crate::backend::traits::{
 };
 
 #[derive(Debug, PartialEq, Eq, Copy, Clone)]
-pub struct EnvironmentBuilderImpl(lmdb::EnvironmentBuilder);
+pub struct EnvironmentBuilderImpl {
+    builder: lmdb::EnvironmentBuilder,
+    make_dir: bool,
+}
 
 impl<'b> BackendEnvironmentBuilder<'b> for EnvironmentBuilderImpl {
     type Environment = EnvironmentImpl;
     type Error = ErrorImpl;
     type Flags = EnvironmentFlagsImpl;
 
     fn new() -> EnvironmentBuilderImpl {
-        EnvironmentBuilderImpl(lmdb::Environment::new())
+        EnvironmentBuilderImpl {
+            builder: lmdb::Environment::new(),
+            make_dir: false,
+        }
     }
 
     fn set_flags<T>(&mut self, flags: T) -> &mut Self
     where
         T: Into<Self::Flags>,
     {
-        self.0.set_flags(flags.into().0);
+        self.builder.set_flags(flags.into().0);
         self
     }
 
     fn set_max_readers(&mut self, max_readers: u32) -> &mut Self {
-        self.0.set_max_readers(max_readers);
+        self.builder.set_max_readers(max_readers);
         self
     }
 
     fn set_max_dbs(&mut self, max_dbs: u32) -> &mut Self {
-        self.0.set_max_dbs(max_dbs);
+        self.builder.set_max_dbs(max_dbs);
         self
     }
 
     fn set_map_size(&mut self, size: usize) -> &mut Self {
-        self.0.set_map_size(size);
+        self.builder.set_map_size(size);
+        self
+    }
+
+    fn set_make_dir_if_needed(&mut self, make_dir: bool) -> &mut Self {
+        self.make_dir = make_dir;
         self
     }
 
     fn open(&self, path: &Path) -> Result<Self::Environment, Self::Error> {
-        self.0.open(path).map(EnvironmentImpl).map_err(ErrorImpl)
+        if !path.is_dir() {
+            if !self.make_dir {
+                return Err(ErrorImpl::DirectoryDoesNotExistError(path.into()));
+            }
+            fs::create_dir_all(path).map_err(ErrorImpl::IoError)?;
+        }
+        self.builder.open(path).map(EnvironmentImpl).map_err(ErrorImpl::LmdbError)
     }
 }
 
@@ -78,38 +98,38 @@ impl<'e> BackendEnvironment<'e> for EnvironmentImpl {
     type Stat = StatImpl;
 
     fn open_db(&self, name: Option<&str>) -> Result<Self::Database, Self::Error> {
-        self.0.open_db(name).map(DatabaseImpl).map_err(ErrorImpl)
+        self.0.open_db(name).map(DatabaseImpl).map_err(ErrorImpl::LmdbError)
     }
 
     fn create_db(&self, name: Option<&str>, flags: Self::Flags) -> Result<Self::Database, Self::Error> {
-        self.0.create_db(name, flags.0).map(DatabaseImpl).map_err(ErrorImpl)
+        self.0.create_db(name, flags.0).map(DatabaseImpl).map_err(ErrorImpl::LmdbError)
     }
 
     fn begin_ro_txn(&'e self) -> Result<Self::RoTransaction, Self::Error> {
-        self.0.begin_ro_txn().map(RoTransactionImpl).map_err(ErrorImpl)
+        self.0.begin_ro_txn().map(RoTransactionImpl).map_err(ErrorImpl::LmdbError)
     }
 
     fn begin_rw_txn(&'e self) -> Result<Self::RwTransaction, Self::Error> {
-        self.0.begin_rw_txn().map(RwTransactionImpl).map_err(ErrorImpl)
+        self.0.begin_rw_txn().map(RwTransactionImpl).map_err(ErrorImpl::LmdbError)
     }
 
     fn sync(&self, force: bool) -> Result<(), Self::Error> {
-        self.0.sync(force).map_err(ErrorImpl)
+        self.0.sync(force).map_err(ErrorImpl::LmdbError)
     }
 
     fn stat(&self) -> Result<Self::Stat, Self::Error> {
-        self.0.stat().map(StatImpl).map_err(ErrorImpl)
+        self.0.stat().map(StatImpl).map_err(ErrorImpl::LmdbError)
     }
 
     fn info(&self) -> Result<Self::Info, Self::Error> {
-        self.0.info().map(InfoImpl).map_err(ErrorImpl)
+        self.0.info().map(InfoImpl).map_err(ErrorImpl::LmdbError)
     }
 
     fn freelist(&self) -> Result<usize, Self::Error> {
-        self.0.freelist().map_err(ErrorImpl)
+        self.0.freelist().map_err(ErrorImpl::LmdbError)
     }
 
     fn set_map_size(&self, size: usize) -> Result<(), Self::Error> {
-        self.0.set_map_size(size).map_err(ErrorImpl)
+        self.0.set_map_size(size).map_err(ErrorImpl::LmdbError)
     }
 }

src/backend/impl_lmdb/error.rs

@@ -8,34 +8,48 @@
 // CONDITIONS OF ANY KIND, either express or implied. See the License for the
 // specific language governing permissions and limitations under the License.
 
-use std::fmt;
+use std::{
+    fmt,
+    io,
+    path::PathBuf,
+};
 
 use crate::{
     backend::traits::BackendError,
     error::StoreError,
 };
 
 #[derive(Debug)]
-pub struct ErrorImpl(pub(crate) lmdb::Error);
+pub enum ErrorImpl {
+    LmdbError(lmdb::Error),
+    DirectoryDoesNotExistError(PathBuf),
+    IoError(io::Error),
+}
 
 impl BackendError for ErrorImpl {}
 
 impl fmt::Display for ErrorImpl {
     fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
-        self.0.fmt(fmt)
+        match self {
+            ErrorImpl::LmdbError(e) => e.fmt(fmt),
+            ErrorImpl::DirectoryDoesNotExistError(_) => write!(fmt, "DirectoryDoesNotExistError"),
+            ErrorImpl::IoError(e) => e.fmt(fmt),
+        }
     }
 }
 
 impl Into<StoreError> for ErrorImpl {
     fn into(self) -> StoreError {
-        match self.0 {
-            lmdb::Error::NotFound => StoreError::KeyValuePairNotFound,
-            lmdb::Error::BadValSize => StoreError::KeyValuePairBadSize,
-            lmdb::Error::Invalid => StoreError::FileInvalid,
-            lmdb::Error::MapFull => StoreError::MapFull,
-            lmdb::Error::DbsFull => StoreError::DbsFull,
-            lmdb::Error::ReadersFull => StoreError::ReadersFull,
-            _ => StoreError::LmdbError(self.0),
+        match self {
+            ErrorImpl::LmdbError(lmdb::Error::NotFound) => StoreError::KeyValuePairNotFound,
+            ErrorImpl::LmdbError(lmdb::Error::BadValSize) => StoreError::KeyValuePairBadSize,
+            ErrorImpl::LmdbError(lmdb::Error::Invalid) => StoreError::FileInvalid,
+            ErrorImpl::LmdbError(lmdb::Error::MapFull) => StoreError::MapFull,
+            ErrorImpl::LmdbError(lmdb::Error::DbsFull) => StoreError::DbsFull,
+            ErrorImpl::LmdbError(lmdb::Error::ReadersFull) => StoreError::ReadersFull,
+            ErrorImpl::LmdbError(error) => StoreError::LmdbError(error),
+            ErrorImpl::DirectoryDoesNotExistError(path) => StoreError::DirectoryDoesNotExistError(path),
+            ErrorImpl::IoError(error) => StoreError::IoError(error),
         }
     }
 }

src/backend/impl_lmdb/iter.rs

@@ -36,6 +36,6 @@ impl<'i, C> BackendIter<'i> for IterImpl<'i, C> {
 
     #[allow(clippy::type_complexity)]
     fn next(&mut self) -> Option<Result<(&'i [u8], &'i [u8]), Self::Error>> {
-        self.iter.next().map(|e| e.map_err(ErrorImpl))
+        self.iter.next().map(|e| e.map_err(ErrorImpl::LmdbError))
     }
 }

src/backend/impl_lmdb/transaction.rs

@@ -31,7 +31,7 @@ impl<'t> BackendRoTransaction for RoTransactionImpl<'t> {
     type Error = ErrorImpl;
 
     fn get(&self, db: &Self::Database, key: &[u8]) -> Result<&[u8], Self::Error> {
-        self.0.get(db.0, &key).map_err(ErrorImpl)
+        self.0.get(db.0, &key).map_err(ErrorImpl::LmdbError)
     }
 
     fn abort(self) {
@@ -43,7 +43,7 @@ impl<'t> BackendRoCursorTransaction<'t> for RoTransactionImpl<'t> {
     type RoCursor = RoCursorImpl<'t>;
 
     fn open_ro_cursor(&'t self, db: &Self::Database) -> Result<Self::RoCursor, Self::Error> {
-        self.0.open_ro_cursor(db.0).map(RoCursorImpl).map_err(ErrorImpl)
+        self.0.open_ro_cursor(db.0).map(RoCursorImpl).map_err(ErrorImpl::LmdbError)
     }
 }
 
@@ -56,29 +56,29 @@ impl<'t> BackendRwTransaction for RwTransactionImpl<'t> {
     type Flags = WriteFlagsImpl;
 
     fn get(&self, db: &Self::Database, key: &[u8]) -> Result<&[u8], Self::Error> {
-        self.0.get(db.0, &key).map_err(ErrorImpl)
+        self.0.get(db.0, &key).map_err(ErrorImpl::LmdbError)
     }
 
     fn put(&mut self, db: &Self::Database, key: &[u8], value: &[u8], flags: Self::Flags) -> Result<(), Self::Error> {
-        self.0.put(db.0, &key, &value, flags.0).map_err(ErrorImpl)
+        self.0.put(db.0, &key, &value, flags.0).map_err(ErrorImpl::LmdbError)
     }
 
     #[cfg(not(feature = "db-dup-sort"))]
     fn del(&mut self, db: &Self::Database, key: &[u8]) -> Result<(), Self::Error> {
-        self.0.del(db.0, &key, None).map_err(ErrorImpl)
+        self.0.del(db.0, &key, None).map_err(ErrorImpl::LmdbError)
     }
 
     #[cfg(feature = "db-dup-sort")]
     fn del(&mut self, db: &Self::Database, key: &[u8], value: Option<&[u8]>) -> Result<(), Self::Error> {
-        self.0.del(db.0, &key, value).map_err(ErrorImpl)
+        self.0.del(db.0, &key, value).map_err(ErrorImpl::LmdbError)
     }
 
     fn clear_db(&mut self, db: &Self::Database) -> Result<(), Self::Error> {
-        self.0.clear_db(db.0).map_err(ErrorImpl)
+        self.0.clear_db(db.0).map_err(ErrorImpl::LmdbError)
     }
 
     fn commit(self) -> Result<(), Self::Error> {
-        self.0.commit().map_err(ErrorImpl)
+        self.0.commit().map_err(ErrorImpl::LmdbError)
     }
 
     fn abort(self) {
@@ -90,6 +90,6 @@ impl<'t> BackendRwCursorTransaction<'t> for RwTransactionImpl<'t> {
     type RoCursor = RoCursorImpl<'t>;
 
     fn open_ro_cursor(&'t self, db: &Self::Database) -> Result<Self::RoCursor, Self::Error> {
-        self.0.open_ro_cursor(db.0).map(RoCursorImpl).map_err(ErrorImpl)
+        self.0.open_ro_cursor(db.0).map(RoCursorImpl).map_err(ErrorImpl::LmdbError)
     }
 }

src/backend/impl_safe/environment.rs

@@ -54,6 +54,7 @@ pub struct EnvironmentBuilderImpl {
     max_readers: Option<usize>,
     max_dbs: Option<usize>,
     map_size: Option<usize>,
+    make_dir: bool,
 }
 
 impl<'b> BackendEnvironmentBuilder<'b> for EnvironmentBuilderImpl {
@@ -67,6 +68,7 @@ impl<'b> BackendEnvironmentBuilder<'b> for EnvironmentBuilderImpl {
             max_readers: None,
             max_dbs: None,
             map_size: None,
+            make_dir: false,
         }
     }
 
@@ -93,7 +95,18 @@ impl<'b> BackendEnvironmentBuilder<'b> for EnvironmentBuilderImpl {
         self
     }
 
+    fn set_make_dir_if_needed(&mut self, make_dir: bool) -> &mut Self {
+        self.make_dir = make_dir;
+        self
+    }
+
     fn open(&self, path: &Path) -> Result<Self::Environment, Self::Error> {
+        if !path.is_dir() {
+            if !self.make_dir {
+                return Err(ErrorImpl::DirectoryDoesNotExistError(path.into()));
+            }
+            fs::create_dir_all(path)?;
+        }
         let mut env = EnvironmentImpl::new(path, self.flags, self.max_readers, self.max_dbs, self.map_size)?;
         env.read_from_disk()?;
         Ok(env)

src/backend/impl_safe/error.rs

@@ -11,6 +11,7 @@
 use std::{
     fmt,
     io,
+    path::PathBuf,
 };
 
 use bincode::Error as BincodeError;
@@ -28,6 +29,7 @@ pub enum ErrorImpl {
     DbsIllegalOpen,
     DbNotFoundError,
     DbIsForeignError,
+    DirectoryDoesNotExistError(PathBuf),
     IoError(io::Error),
     BincodeError(BincodeError),
 }
@@ -43,6 +45,7 @@ impl fmt::Display for ErrorImpl {
             ErrorImpl::DbsIllegalOpen => write!(fmt, "DbIllegalOpen (safe mode)"),
             ErrorImpl::DbNotFoundError => write!(fmt, "DbNotFoundError (safe mode)"),
             ErrorImpl::DbIsForeignError => write!(fmt, "DbIsForeignError (safe mode)"),
+            ErrorImpl::DirectoryDoesNotExistError(_) => write!(fmt, "DirectoryDoesNotExistError (safe mode)"),
             ErrorImpl::IoError(e) => e.fmt(fmt),
             ErrorImpl::BincodeError(e) => e.fmt(fmt),
         }
@@ -59,6 +62,8 @@ impl Into<StoreError> for ErrorImpl {
             ErrorImpl::KeyValuePairNotFound => StoreError::KeyValuePairNotFound,
             ErrorImpl::BincodeError(_) => StoreError::FileInvalid,
             ErrorImpl::DbsFull => StoreError::DbsFull,
+            ErrorImpl::DirectoryDoesNotExistError(path) => StoreError::DirectoryDoesNotExistError(path),
+            ErrorImpl::IoError(error) => StoreError::IoError(error),
             _ => StoreError::SafeModeError(self),
         }
     }

src/backend/traits.rs

@@ -88,6 +88,8 @@ pub trait BackendEnvironmentBuilder<'b>: Debug + Eq + PartialEq + Copy + Clone {
 
     fn set_map_size(&mut self, size: usize) -> &mut Self;
 
+    fn set_make_dir_if_needed(&mut self, make_dir: bool) -> &mut Self;
+
     fn open(&self, path: &Path) -> Result<Self::Environment, Self::Error>;
 }
 

src/env.rs

@@ -88,10 +88,6 @@ where
     where
         B: BackendEnvironmentBuilder<'e, Environment = E>,
     {
-        if !path.is_dir() {
-            return Err(StoreError::DirectoryDoesNotExistError(path.into()));
-        }
-
         let mut builder = B::new();
         builder.set_max_dbs(max_dbs);
 
@@ -104,10 +100,6 @@ where
     where
         B: BackendEnvironmentBuilder<'e, Environment = E>,
     {
-        if !path.is_dir() {
-            return Err(StoreError::DirectoryDoesNotExistError(path.into()));
-        }
-
         Ok(Rkv {
             path: path.into(),
             env: builder.open(path).map_err(|e| {

tests/env-lmdb.rs

@@ -14,6 +14,7 @@
 
 use std::{
     fs,
+    path::Path,
     str,
     sync::{
         Arc,
@@ -103,6 +104,33 @@ fn test_open_from_builder() {
     check_rkv(&k);
 }
 
+#[test]
+fn test_open_from_builder_with_dir_1() {
+    let root = Builder::new().prefix("test_open_from_builder").tempdir().expect("tempdir");
+    println!("Root path: {:?}", root.path());
+
+    let mut builder = Rkv::environment_builder::<Lmdb>();
+    builder.set_max_dbs(2);
+    builder.set_make_dir_if_needed(true);
+
+    let k = Rkv::from_builder(root.path(), builder).expect("rkv");
+    check_rkv(&k);
+}
+
+#[test]
+#[should_panic(expected = "rkv: DirectoryDoesNotExistError(\"bogus\")")]
+fn test_open_from_builder_with_dir_2() {
+    let root = Path::new("bogus");
+    println!("Root path: {:?}", root);
+    assert!(!root.is_dir());
+
+    let mut builder = Rkv::environment_builder::<Lmdb>();
+    builder.set_max_dbs(2);
+
+    let k = Rkv::from_builder(root, builder).expect("rkv");
+    check_rkv(&k);
+}
+
 #[test]
 #[should_panic(expected = "opened: DbsFull")]
 fn test_open_with_capacity() {