From 017757004eed25d1b6a94d6310e07193a5991577 Mon Sep 17 00:00:00 2001
From: Louis Dureuil <louis@meilisearch.com>
Date: Wed, 16 Oct 2024 09:26:18 +0200
Subject: [PATCH] Add PrimaryKey::new_or_insert

---
 milli/src/documents/primary_key.rs | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/milli/src/documents/primary_key.rs b/milli/src/documents/primary_key.rs
index 904109033..79fd07048 100644
--- a/milli/src/documents/primary_key.rs
+++ b/milli/src/documents/primary_key.rs
@@ -65,6 +65,18 @@ impl<'a> PrimaryKey<'a> {
         })
     }
 
+    pub fn new_or_insert(
+        path: &'a str,
+        fields: &mut impl MutFieldIdMapper,
+    ) -> StdResult<Self, UserError> {
+        Ok(if path.contains(PRIMARY_KEY_SPLIT_SYMBOL) {
+            Self::Nested { name: path }
+        } else {
+            let field_id = fields.insert(path).ok_or(UserError::AttributeLimitReached)?;
+            Self::Flat { name: path, field_id }
+        })
+    }
+
     pub fn name(&self) -> &'a str {
         match self {
             PrimaryKey::Flat { name, .. } => name,