From 5f1bfb73eeeb1da17a1d2598d01916b08022aa8c Mon Sep 17 00:00:00 2001
From: Kerollmops <clement@meilisearch.com>
Date: Tue, 21 Jun 2022 10:45:25 +0200
Subject: [PATCH] Extract the primary key name and make it accessible

---
 milli/src/documents/enriched.rs              | 28 +++++++++++++++++---
 milli/src/update/index_documents/validate.rs |  6 ++++-
 2 files changed, 29 insertions(+), 5 deletions(-)

diff --git a/milli/src/documents/enriched.rs b/milli/src/documents/enriched.rs
index 8645e06c4..918b47c95 100644
--- a/milli/src/documents/enriched.rs
+++ b/milli/src/documents/enriched.rs
@@ -17,17 +17,20 @@ use crate::FieldId;
 /// `FieldId`. The mapping between the field ids and the field names is done thanks to the index.
 pub struct EnrichedDocumentsBatchReader<R> {
     documents: DocumentsBatchReader<R>,
+    primary_key: String,
     external_ids: grenad::ReaderCursor<File>,
 }
 
 impl<R: io::Read + io::Seek> EnrichedDocumentsBatchReader<R> {
     pub fn new(
         documents: DocumentsBatchReader<R>,
+        primary_key: String,
         external_ids: grenad::Reader<File>,
     ) -> Result<Self, Error> {
         if documents.documents_count() as u64 == external_ids.len() {
             Ok(EnrichedDocumentsBatchReader {
                 documents,
+                primary_key,
                 external_ids: external_ids.into_cursor()?,
             })
         } else {
@@ -39,6 +42,10 @@ impl<R: io::Read + io::Seek> EnrichedDocumentsBatchReader<R> {
         self.documents.documents_count()
     }
 
+    pub fn primary_key(&self) -> &str {
+        &self.primary_key
+    }
+
     pub fn is_empty(&self) -> bool {
         self.documents.is_empty()
     }
@@ -49,9 +56,13 @@ impl<R: io::Read + io::Seek> EnrichedDocumentsBatchReader<R> {
 
     /// This method returns a forward cursor over the enriched documents.
     pub fn into_cursor(self) -> EnrichedDocumentsBatchCursor<R> {
-        let EnrichedDocumentsBatchReader { documents, mut external_ids } = self;
+        let EnrichedDocumentsBatchReader { documents, primary_key, mut external_ids } = self;
         external_ids.reset();
-        EnrichedDocumentsBatchCursor { documents: documents.into_cursor(), external_ids }
+        EnrichedDocumentsBatchCursor {
+            documents: documents.into_cursor(),
+            primary_key,
+            external_ids,
+        }
     }
 }
 
@@ -63,13 +74,22 @@ pub struct EnrichedDocument<'a> {
 
 pub struct EnrichedDocumentsBatchCursor<R> {
     documents: DocumentsBatchCursor<R>,
+    primary_key: String,
     external_ids: grenad::ReaderCursor<File>,
 }
 
 impl<R> EnrichedDocumentsBatchCursor<R> {
     pub fn into_reader(self) -> EnrichedDocumentsBatchReader<R> {
-        let EnrichedDocumentsBatchCursor { documents, external_ids } = self;
-        EnrichedDocumentsBatchReader { documents: documents.into_reader(), external_ids }
+        let EnrichedDocumentsBatchCursor { documents, primary_key, external_ids } = self;
+        EnrichedDocumentsBatchReader {
+            documents: documents.into_reader(),
+            primary_key,
+            external_ids,
+        }
+    }
+
+    pub fn primary_key(&self) -> &str {
+        &self.primary_key
     }
 
     pub fn documents_batch_index(&self) -> &DocumentsBatchIndex {
diff --git a/milli/src/update/index_documents/validate.rs b/milli/src/update/index_documents/validate.rs
index 8b68532cb..83d7ef38e 100644
--- a/milli/src/update/index_documents/validate.rs
+++ b/milli/src/update/index_documents/validate.rs
@@ -94,7 +94,11 @@ pub fn validate_and_enrich_documents_batch<R: Read + Seek>(
     }
 
     let external_ids = writer_into_reader(external_ids)?;
-    let reader = EnrichedDocumentsBatchReader::new(cursor.into_reader(), external_ids)?;
+    let reader = EnrichedDocumentsBatchReader::new(
+        cursor.into_reader(),
+        primary_key.primary_key().to_string(),
+        external_ids,
+    )?;
 
     Ok(Ok(reader))
 }