MeiliSearch/meilisearch-http/src/index/dump.rs

135 lines
4.1 KiB
Rust
Raw Normal View History

2021-05-31 16:40:59 +02:00
use std::fs::{create_dir_all, File};
use std::io::{BufRead, BufReader, Write};
use std::path::Path;
use std::sync::Arc;
2021-06-23 14:48:33 +02:00
use anyhow::{bail, Context};
2021-05-26 22:52:06 +02:00
use heed::RoTxn;
use indexmap::IndexMap;
use milli::update::{IndexDocumentsMethod, UpdateFormat::JsonStream};
use serde::{Deserialize, Serialize};
use crate::option::IndexerOpts;
2021-06-17 14:36:32 +02:00
use super::error::Result;
2021-06-15 17:39:07 +02:00
use super::{update_handler::UpdateHandler, Index, Settings, Unchecked};
2021-05-26 22:52:06 +02:00
#[derive(Serialize, Deserialize)]
struct DumpMeta {
2021-05-27 14:30:20 +02:00
settings: Settings<Unchecked>,
2021-05-26 22:52:06 +02:00
primary_key: Option<String>,
}
2021-05-31 16:03:39 +02:00
const META_FILE_NAME: &str = "meta.json";
const DATA_FILE_NAME: &str = "documents.jsonl";
2021-05-26 22:52:06 +02:00
impl Index {
pub fn dump(&self, path: impl AsRef<Path>) -> Result<()> {
// acquire write txn make sure any ongoing write is finished before we start.
2021-05-26 22:52:06 +02:00
let txn = self.env.write_txn()?;
self.dump_documents(&txn, &path)?;
self.dump_meta(&txn, &path)?;
Ok(())
}
fn dump_documents(&self, txn: &RoTxn, path: impl AsRef<Path>) -> Result<()> {
2021-05-26 22:52:06 +02:00
let document_file_path = path.as_ref().join(DATA_FILE_NAME);
let mut document_file = File::create(&document_file_path)?;
2021-06-17 14:36:32 +02:00
let documents = self.all_documents(txn)?;
2021-05-26 22:52:06 +02:00
let fields_ids_map = self.fields_ids_map(txn)?;
// dump documents
let mut json_map = IndexMap::new();
for document in documents {
let (_, reader) = document?;
for (fid, bytes) in reader.iter() {
if let Some(name) = fields_ids_map.name(fid) {
json_map.insert(name, serde_json::from_slice::<serde_json::Value>(bytes)?);
}
}
serde_json::to_writer(&mut document_file, &json_map)?;
2021-05-31 16:40:59 +02:00
document_file.write_all(b"\n")?;
2021-05-26 22:52:06 +02:00
json_map.clear();
}
Ok(())
}
fn dump_meta(&self, txn: &RoTxn, path: impl AsRef<Path>) -> Result<()> {
2021-05-26 22:52:06 +02:00
let meta_file_path = path.as_ref().join(META_FILE_NAME);
let mut meta_file = File::create(&meta_file_path)?;
2021-05-27 14:30:20 +02:00
let settings = self.settings_txn(txn)?.into_unchecked();
2021-05-26 22:52:06 +02:00
let primary_key = self.primary_key(txn)?.map(String::from);
2021-05-31 10:58:51 +02:00
let meta = DumpMeta {
settings,
primary_key,
};
2021-05-26 22:52:06 +02:00
serde_json::to_writer(&mut meta_file, &meta)?;
Ok(())
}
pub fn load_dump(
src: impl AsRef<Path>,
dst: impl AsRef<Path>,
2021-05-31 16:40:59 +02:00
size: usize,
2021-05-26 22:52:06 +02:00
indexing_options: &IndexerOpts,
2021-06-15 17:39:07 +02:00
) -> anyhow::Result<()> {
2021-05-26 22:52:06 +02:00
let dir_name = src
.as_ref()
.file_name()
2021-06-15 17:39:07 +02:00
.with_context(|| format!("invalid dump index: {}", src.as_ref().display()))?;
2021-05-27 14:30:20 +02:00
let dst_dir_path = dst.as_ref().join("indexes").join(dir_name);
2021-05-26 22:52:06 +02:00
create_dir_all(&dst_dir_path)?;
let meta_path = src.as_ref().join(META_FILE_NAME);
let mut meta_file = File::open(meta_path)?;
2021-05-31 10:58:51 +02:00
let DumpMeta {
settings,
primary_key,
} = serde_json::from_reader(&mut meta_file)?;
2021-05-27 14:30:20 +02:00
let settings = settings.check();
2021-05-31 16:40:59 +02:00
let index = Self::open(&dst_dir_path, size)?;
2021-05-26 22:52:06 +02:00
let mut txn = index.write_txn()?;
2021-07-29 18:14:36 +02:00
let handler = UpdateHandler::new(indexing_options)?;
2021-05-26 22:52:06 +02:00
index.update_settings_txn(&mut txn, &settings, handler.update_builder(0))?;
let document_file_path = src.as_ref().join(DATA_FILE_NAME);
2021-05-31 10:58:51 +02:00
let reader = File::open(&document_file_path)?;
let mut reader = BufReader::new(reader);
reader.fill_buf()?;
// If the document file is empty, we don't perform the document addition, to prevent
// a primary key error to be thrown.
if !reader.buffer().is_empty() {
index.update_documents_txn(
&mut txn,
JsonStream,
IndexDocumentsMethod::UpdateDocuments,
Some(reader),
handler.update_builder(0),
primary_key.as_deref(),
)?;
}
2021-05-26 22:52:06 +02:00
txn.commit()?;
match Arc::try_unwrap(index.0) {
Ok(inner) => inner.prepare_for_closing().wait(),
2021-06-21 13:57:32 +02:00
Err(_) => bail!("Could not close index properly."),
2021-05-26 22:52:06 +02:00
}
Ok(())
}
}