From e067d796b318af321b08fb49d5736d995f884ae3 Mon Sep 17 00:00:00 2001
From: Kerollmops <clement@meilisearch.com>
Date: Wed, 2 Apr 2025 15:56:56 +0200
Subject: [PATCH] Improve the primary key stop reasons error messages

---
 .../src/scheduler/autobatcher.rs              | 27 ++++++++++++++----
 crates/meilisearch-types/src/tasks.rs         | 28 ++++++++++++++++---
 2 files changed, 46 insertions(+), 9 deletions(-)

diff --git a/crates/index-scheduler/src/scheduler/autobatcher.rs b/crates/index-scheduler/src/scheduler/autobatcher.rs
index 42a508cb0..a37a849b0 100644
--- a/crates/index-scheduler/src/scheduler/autobatcher.rs
+++ b/crates/index-scheduler/src/scheduler/autobatcher.rs
@@ -7,7 +7,7 @@ The main function of the autobatcher is [`next_autobatch`].
 
 use std::ops::ControlFlow::{self, Break, Continue};
 
-use meilisearch_types::tasks::{BatchStopReason, TaskId};
+use meilisearch_types::tasks::{BatchStopReason, PrimaryKeyMismatchReason, TaskId};
 
 use crate::KindWithContent;
 
@@ -241,24 +241,41 @@ impl BatchKind {
         let pk: Option<String> = match (self.primary_key(), kind.primary_key(), primary_key) {
             // 1. If both task don't interact with primary key -> we can continue
             (batch_pk, None | Some(None), _) => {
-                batch_pk.flatten().to_owned()
+                batch_pk.flatten().map(ToOwned::to_owned)
             },
             // 2.1 If we already have a primary-key ->
             // 2.1.1 If the task we're trying to accumulate have a pk it must be equal to our primary key
             (batch_pk, Some(Some(task_pk)), Some(index_pk)) => if task_pk == index_pk {
                 Some(task_pk.to_owned())
             } else {
-                return Break((this, BatchStopReason::PrimaryKeyMismatch { id, batch_pk: todo!(), task_pk: todo!() }))
+                return Break((self, BatchStopReason::PrimaryKeyMismatch {
+                    id,
+                    reason: PrimaryKeyMismatchReason::TaskPrimaryKeyDifferFromIndexPrimaryKey {
+                        task_pk: task_pk.to_owned(),
+                        index_pk: index_pk.to_owned(),
+                    },
+                }))
             },
             // 2.2 If we don't have a primary-key ->
             // 2.2.2 If the batch is set to Some(None), the task should be too
             (Some(None), Some(None), None) => None,
-            (Some(None), Some(Some(_)), None) => return Break((this, BatchStopReason::PrimaryKeyMismatch { id, batch_pk: todo!(), task_pk: todo!() })),
+            (Some(None), Some(Some(task_pk)), None) => return Break((self, BatchStopReason::PrimaryKeyMismatch {
+                id,
+                reason: PrimaryKeyMismatchReason::CannotInterfereWithPrimaryKeyGuessing {
+                    task_pk: task_pk.to_owned(),
+                },
+            })),
             (Some(Some(batch_pk)), Some(None), None) => Some(batch_pk.to_owned()),
             (Some(Some(batch_pk)), Some(Some(task_pk)), None) => if task_pk == batch_pk {
                 Some(task_pk.to_owned())
             } else {
-                return Break((this, BatchStopReason::PrimaryKeyMismatch { id, batch_pk: todo!(), task_pk: todo!() }))
+                return Break((self, BatchStopReason::PrimaryKeyMismatch {
+                    id,
+                    reason: PrimaryKeyMismatchReason::TaskPrimaryKeyDifferFromCurrentBatchPrimaryKey {
+                        batch_pk: batch_pk.to_owned(),
+                        task_pk: task_pk.to_owned(),
+                    },
+                }))
             },
             (None, Some(Some(task_pk)), None) => Some(task_pk.to_owned())
         };
diff --git a/crates/meilisearch-types/src/tasks.rs b/crates/meilisearch-types/src/tasks.rs
index cbed6df78..2b079fa35 100644
--- a/crates/meilisearch-types/src/tasks.rs
+++ b/crates/meilisearch-types/src/tasks.rs
@@ -704,8 +704,7 @@ pub enum BatchStopReason {
     },
     PrimaryKeyMismatch {
         id: TaskId,
-        batch_pk: Option<String>,
-        task_pk: Option<String>,
+        reason: PrimaryKeyMismatchReason,
     },
     IndexDeletion {
         id: TaskId,
@@ -732,7 +731,12 @@ impl BatchStopReason {
     }
 }
 
-pub enum PrimaryKeyMismatchReason {}
+#[derive(Debug, Clone)]
+pub enum PrimaryKeyMismatchReason {
+    TaskPrimaryKeyDifferFromIndexPrimaryKey { task_pk: String, index_pk: String },
+    TaskPrimaryKeyDifferFromCurrentBatchPrimaryKey { task_pk: String, batch_pk: String },
+    CannotInterfereWithPrimaryKeyGuessing { task_pk: String },
+}
 
 impl Display for BatchStopReason {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
@@ -758,7 +762,23 @@ impl Display for BatchStopReason {
             BatchStopReason::IndexCreationMismatch { id } => {
                 write!(f, "task with id {id} has different index creation rules as in the batch")
             }
-            BatchStopReason::PrimaryKeyMismatch { id, batch_pk, task_pk } => {}
+            BatchStopReason::PrimaryKeyMismatch { reason, id } => match reason {
+                PrimaryKeyMismatchReason::TaskPrimaryKeyDifferFromIndexPrimaryKey {
+                    task_pk,
+                    index_pk,
+                } => {
+                    write!(f, "primary key `{task_pk}` in task with id {id} is different from the primary key of the index `{index_pk}`")
+                }
+                PrimaryKeyMismatchReason::TaskPrimaryKeyDifferFromCurrentBatchPrimaryKey {
+                    task_pk,
+                    batch_pk,
+                } => {
+                    write!(f, "primary key `{task_pk}` in task with id {id} is different from `{batch_pk}`: the primary key of the batch being build")
+                }
+                PrimaryKeyMismatchReason::CannotInterfereWithPrimaryKeyGuessing { task_pk } => {
+                    write!(f, "Task {id} is setting the `{task_pk}` primary key but cannot interfere with primary key guessing of the previous batch")
+                }
+            },
             BatchStopReason::IndexDeletion { id } => {
                 write!(f, "task with id {id} deletes the index")
             }