From 91d8198d17e6801a24a7b18a24c16651001aec7e Mon Sep 17 00:00:00 2001
From: mpostma <postma.marin@protonmail.com>
Date: Wed, 30 Dec 2020 18:43:50 +0100
Subject: [PATCH] return documents number on addition

---
 src/update/index_documents/mod.rs | 15 ++++++++++++---
 src/update/mod.rs                 |  2 +-
 2 files changed, 13 insertions(+), 4 deletions(-)

diff --git a/src/update/index_documents/mod.rs b/src/update/index_documents/mod.rs
index 90bc5ef3d..2c5d34092 100644
--- a/src/update/index_documents/mod.rs
+++ b/src/update/index_documents/mod.rs
@@ -12,8 +12,9 @@ use grenad::{Writer, Sorter, Merger, Reader, FileFuse, CompressionType};
 use heed::types::ByteSlice;
 use log::{debug, info, error};
 use memmap::Mmap;
-use rayon::prelude::*;
 use rayon::ThreadPool;
+use rayon::prelude::*;
+use serde::{Serialize, Deserialize};
 
 use crate::index::Index;
 use crate::update::{Facets, UpdateIndexingStep};
@@ -32,6 +33,11 @@ mod merge_function;
 mod store;
 mod transform;
 
+#[derive(Debug, Serialize, Deserialize, Clone)]
+pub struct DocumentAdditionResult {
+    nb_documents: usize,
+}
+
 #[derive(Debug, Copy, Clone)]
 pub enum WriteMethod {
     Append,
@@ -253,7 +259,7 @@ impl<'t, 'u, 'i, 'a> IndexDocuments<'t, 'u, 'i, 'a> {
         self.autogenerate_docids = false;
     }
 
-    pub fn execute<R, F>(self, reader: R, progress_callback: F) -> anyhow::Result<()>
+    pub fn execute<R, F>(self, reader: R, progress_callback: F) -> anyhow::Result<DocumentAdditionResult>
     where
         R: io::Read,
         F: Fn(UpdateIndexingStep) + Sync,
@@ -279,9 +285,12 @@ impl<'t, 'u, 'i, 'a> IndexDocuments<'t, 'u, 'i, 'a> {
             UpdateFormat::JsonStream => transform.output_from_json_stream(reader, &progress_callback)?,
         };
 
+        let nb_documents = output.documents_count;
+
         info!("Update transformed in {:.02?}", before_transform.elapsed());
 
-        self.execute_raw(output, progress_callback)
+        self.execute_raw(output, progress_callback)?;
+        Ok(DocumentAdditionResult { nb_documents })
     }
 
     pub fn execute_raw<F>(self, output: TransformOutput, progress_callback: F) -> anyhow::Result<()>
diff --git a/src/update/mod.rs b/src/update/mod.rs
index 407d9f498..2cd532c83 100644
--- a/src/update/mod.rs
+++ b/src/update/mod.rs
@@ -10,7 +10,7 @@ mod update_step;
 pub use self::available_documents_ids::AvailableDocumentsIds;
 pub use self::clear_documents::ClearDocuments;
 pub use self::delete_documents::DeleteDocuments;
-pub use self::index_documents::{IndexDocuments, IndexDocumentsMethod, UpdateFormat};
+pub use self::index_documents::{IndexDocuments, IndexDocumentsMethod, UpdateFormat, DocumentAdditionResult};
 pub use self::facets::Facets;
 pub use self::settings::Settings;
 pub use self::update_builder::UpdateBuilder;