From 9423310816a4473b0f6d998954525ad85dee7fbb Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Cl=C3=A9ment=20Renault?= <clement@meilisearch.com>
Date: Sun, 21 Feb 2021 17:21:13 +0100
Subject: [PATCH] Introduce an helpers crate that export the database to stdout

---
 Cargo.lock          | 12 +++++++
 Cargo.toml          |  2 +-
 helpers/Cargo.toml  | 14 ++++++++
 helpers/src/main.rs | 86 +++++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 113 insertions(+), 1 deletion(-)
 create mode 100644 helpers/Cargo.toml
 create mode 100644 helpers/src/main.rs

diff --git a/Cargo.lock b/Cargo.lock
index 112b98690..679eed2c0 100644
--- a/Cargo.lock
+++ b/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"
diff --git a/Cargo.toml b/Cargo.toml
index 16a5ab8d5..a60c293e3 100644
--- a/Cargo.toml
+++ b/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]
diff --git a/helpers/Cargo.toml b/helpers/Cargo.toml
new file mode 100644
index 000000000..ea0e1ddc8
--- /dev/null
+++ b/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 }
diff --git a/helpers/src/main.rs b/helpers/src/main.rs
new file mode 100644
index 000000000..c916d0448
--- /dev/null
+++ b/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) }
+}