use std::ops::Deref; use async_compression::tokio_02::write::GzipEncoder; use futures_util::stream::StreamExt; use tokio::io::AsyncWriteExt; use milli::update::{IndexDocumentsMethod, UpdateFormat}; use milli::update_store::UpdateStatus; use super::Data; use crate::updates::{UpdateMeta, UpdateResult, UpdateStatusResponse, Settings}; impl Data { pub async fn add_documents( &self, _index: S, method: IndexDocumentsMethod, format: UpdateFormat, mut stream: impl futures::Stream> + Unpin, ) -> anyhow::Result where B: Deref, E: std::error::Error + Send + Sync + 'static, S: AsRef, { let file = tokio::task::spawn_blocking(tempfile::tempfile).await?; let file = tokio::fs::File::from_std(file?); let mut encoder = GzipEncoder::new(file); while let Some(result) = stream.next().await { let bytes = &*result?; encoder.write_all(&bytes[..]).await?; } encoder.shutdown().await?; let mut file = encoder.into_inner(); file.sync_all().await?; let file = file.into_std().await; let mmap = unsafe { memmap::Mmap::map(&file)? }; let meta = UpdateMeta::DocumentsAddition { method, format }; let queue = self.update_queue.clone(); let update = tokio::task::spawn_blocking(move || queue.register_update(meta, &mmap[..])).await??; Ok(update.into()) } pub async fn update_settings>( &self, _index: S, settings: Settings ) -> anyhow::Result { let meta = UpdateMeta::Settings(settings); let queue = self.update_queue.clone(); let update = tokio::task::spawn_blocking(move || queue.register_update(meta, &[])).await??; Ok(update.into()) } #[inline] pub fn get_update_status(&self, _index: &str, uid: u64) -> anyhow::Result>> { self.update_queue.get_update_status(uid) } pub fn get_updates_status(&self, _index: &str) -> anyhow::Result>> { let result = self.update_queue.iter_metas(|processing, processed, pending, aborted, failed| { let mut metas = processing .map(UpdateStatus::from) .into_iter() .chain(processed.filter_map(|i| Some(i.ok()?.1)).map(UpdateStatus::from)) .chain(pending.filter_map(|i| Some(i.ok()?.1)).map(UpdateStatus::from)) .chain(aborted.filter_map(|i| Some(i.ok()?.1)).map(UpdateStatus::from)) .chain(failed.filter_map(|i| Some(i.ok()?.1)).map(UpdateStatus::from)) .collect::>(); metas.sort_by(|a, b| a.id().cmp(&b.id())); Ok(metas) })?; Ok(result) } }