Merge pull request #95 from meilisearch/helpers-crate

Introduce an helpers crate that export the database to stdout

Cargo.lock

@@ -859,6 +859,18 @@ dependencies = [
  "zerocopy",
 ]
 
+[[package]]
+name = "helpers"
+version = "0.1.0"
+dependencies = [
+ "anyhow",
+ "byte-unit",
+ "heed",
+ "milli",
+ "stderrlog",
+ "structopt",
+]
+
 [[package]]
 name = "hermit-abi"
 version = "0.1.18"

Cargo.toml

@@ -1,5 +1,5 @@
 [workspace]
-members = ["milli", "http-ui", "infos", "search"]
+members = ["milli", "http-ui", "infos", "helpers", "search"]
 default-members = ["milli"]
 
 [profile.release]

helpers/Cargo.toml

@@ -0,0 +1,14 @@
+[package]
+name = "helpers"
+version = "0.1.0"
+authors = ["Clément Renault <clement@meilisearch.com>"]
+edition = "2018"
+
+[dependencies]
+anyhow = "1.0.38"
+byte-unit = { version = "4.0.9", default-features = false, features = ["std"] }
+heed = "0.10.6"
+jemallocator = "0.3.2"
+milli = { path = "../milli" }
+stderrlog = "0.5.1"
+structopt = { version = "0.3.21", default-features = false }

helpers/src/main.rs

@@ -0,0 +1,86 @@
+use std::path::PathBuf;
+
+use byte_unit::Byte;
+use heed::{Env, EnvOpenOptions, CompactionOption};
+use structopt::StructOpt;
+
+use Command::*;
+
+#[cfg(target_os = "linux")]
+#[global_allocator]
+static ALLOC: jemallocator::Jemalloc = jemallocator::Jemalloc;
+
+#[derive(Debug, StructOpt)]
+/// Some helpers commands for milli.
+pub struct Opt {
+    /// The database path where the database is located.
+    /// It is created if it doesn't already exist.
+    #[structopt(long = "db", parse(from_os_str))]
+    database: PathBuf,
+
+    /// The maximum size the database can take on disk. It is recommended to specify
+    /// the whole disk space (value must be a multiple of a page size).
+    #[structopt(long = "db-size", default_value = "100 GiB")]
+    database_size: Byte,
+
+    /// Verbose mode (-v, -vv, -vvv, etc.)
+    #[structopt(short, long, parse(from_occurrences))]
+    verbose: usize,
+
+    #[structopt(subcommand)]
+    command: Command,
+}
+
+#[derive(Debug, StructOpt)]
+enum Command {
+    /// Outputs the main LMDB database to stdout.
+    CopyMainDatabase {
+        /// Wether to enable or not the compaction of the database.
+        #[structopt(long, short = "c")]
+        enable_compaction: bool,
+    },
+}
+
+fn main() -> anyhow::Result<()> {
+    let opt = Opt::from_args();
+
+    stderrlog::new()
+        .verbosity(opt.verbose)
+        .show_level(false)
+        .timestamp(stderrlog::Timestamp::Off)
+        .init()?;
+
+    let mut options = EnvOpenOptions::new();
+    options.map_size(opt.database_size.get_bytes() as usize);
+
+    // Return an error if the database does not exist.
+    if !opt.database.exists() {
+        anyhow::bail!("The database ({}) does not exist.", opt.database.display());
+    }
+
+    let env = options.open(opt.database)?;
+
+    match opt.command {
+        CopyMainDatabase { enable_compaction } => {
+            use CompactionOption::*;
+            let compaction = if enable_compaction { Enabled } else { Disabled };
+            copy_main_database_to_stdout(env, compaction)
+        },
+    }
+}
+
+#[cfg(target_family = "unix")]
+fn copy_main_database_to_stdout(env: Env, compaction: CompactionOption) -> anyhow::Result<()> {
+    use std::os::unix::io::AsRawFd;
+
+    let stdout = std::io::stdout().as_raw_fd();
+    unsafe { env.copy_to_fd(stdout, compaction).map_err(Into::into) }
+}
+
+#[cfg(target_family = "windows")]
+fn copy_main_database_to_stdout(env: Env, compaction: CompactionOption) -> anyhow::Result<()> {
+    use std::os::windows::io::AsRawHandle;
+
+    let stdout = std::io::stdout().as_raw_handle();
+    unsafe { env.copy_to_fd(stdout, compaction).map_err(Into::into) }
+}