Add optional exception for failed tasks in wait_for_task

meilisearch/client.py

@@ -995,6 +995,8 @@ def wait_for_task(
         uid: int,
         timeout_in_ms: int = 5000,
         interval_in_ms: int = 50,
+        *,
+        raise_on_failure: bool = False,
     ) -> Task:
         """Wait until Meilisearch processes a task until it fails or succeeds.
 
@@ -1006,6 +1008,9 @@ def wait_for_task(
             Time the method should wait before raising a MeilisearchTimeoutError
         interval_in_ms (optional):
             Time interval the method should wait (sleep) between requests
+        raise_on_failure (optional):
+            If True, raise a MeilisearchTaskFailedError when the task reaches the
+            ``failed`` status. Defaults to False.
 
         Returns
         -------
@@ -1016,7 +1021,17 @@ def wait_for_task(
         ------
         MeilisearchTimeoutError
             An error containing details about why Meilisearch can't process your request. Meilisearch error codes are described here: https://www.meilisearch.com/docs/reference/errors/error_codes#meilisearch-errors
+        MeilisearchTaskFailedError
+            Raised when ``raise_on_failure`` is True and the task reaches the
+            ``failed`` status. The exception retains the failed Task and its error mapping.
         """
+        if raise_on_failure:
+            return self.task_handler.wait_for_task(
+                uid,
+                timeout_in_ms,
+                interval_in_ms,
+                raise_on_failure=True,
+            )
         return self.task_handler.wait_for_task(uid, timeout_in_ms, interval_in_ms)
 
     def get_batches(self, parameters: MutableMapping[str, Any] | None = None) -> BatchResults:

meilisearch/errors.py

@@ -3,10 +3,13 @@
 import json
 from collections.abc import Callable
 from functools import wraps
-from typing import Any, TypeVar
+from typing import TYPE_CHECKING, Any, TypeVar
 
 from requests import Response
 
+if TYPE_CHECKING:
+    from meilisearch.models.task import Task
+
 T = TypeVar("T")
 
 
@@ -67,6 +70,20 @@ def __str__(self) -> str:  # pragma: no cover
         return f"MeilisearchTimeoutError, {self.message}"
 
 
+class MeilisearchTaskFailedError(MeilisearchError):
+    """Raised when a task reaches the ``failed`` status and raising is enabled."""
+
+    def __init__(self, task: "Task") -> None:
+        self.task = task
+        self.uid = task.uid
+        self.error = task.error
+        message = (task.error or {}).get("message") or f"Task {task.uid} failed"
+        super().__init__(message)
+
+    def __str__(self) -> str:  # pragma: no cover
+        return f"MeilisearchTaskFailedError. {self.message}"
+
+
 def version_error_hint_message(func: Callable[..., T]) -> Callable[..., T]:
     @wraps(func)
     def wrapper(*args: Any, **kwargs: Any) -> Any:

meilisearch/index.py

@@ -278,6 +278,8 @@ def wait_for_task(
         uid: int,
         timeout_in_ms: int = 5000,
         interval_in_ms: int = 50,
+        *,
+        raise_on_failure: bool = False,
     ) -> Task:
         """Wait until Meilisearch processes a task until it fails or succeeds.
 
@@ -289,6 +291,9 @@ def wait_for_task(
             time the method should wait before raising a MeilisearchTimeoutError.
         interval_in_ms (optional):
             time interval the method should wait (sleep) between requests.
+        raise_on_failure (optional):
+            If True, raise a MeilisearchTaskFailedError when the task reaches the
+            ``failed`` status. Defaults to False.
 
         Returns
         -------
@@ -299,7 +304,17 @@ def wait_for_task(
         ------
         MeilisearchTimeoutError
             An error containing details about why Meilisearch can't process your request. Meilisearch error codes are described here: https://www.meilisearch.com/docs/reference/errors/error_codes#meilisearch-errors
+        MeilisearchTaskFailedError
+            Raised when ``raise_on_failure`` is True and the task reaches the
+            ``failed`` status. The exception retains the failed Task and its error mapping.
         """
+        if raise_on_failure:
+            return self.task_handler.wait_for_task(
+                uid,
+                timeout_in_ms,
+                interval_in_ms,
+                raise_on_failure=True,
+            )
         return self.task_handler.wait_for_task(uid, timeout_in_ms, interval_in_ms)
 
     def get_stats(

meilisearch/task.py

@@ -8,7 +8,7 @@
 
 from meilisearch._httprequests import HttpRequests
 from meilisearch.config import Config
-from meilisearch.errors import MeilisearchTimeoutError
+from meilisearch.errors import MeilisearchTaskFailedError, MeilisearchTimeoutError
 from meilisearch.models.task import Batch, BatchResults, Task, TaskInfo, TaskResults
 
 
@@ -189,6 +189,8 @@ def wait_for_task(
         uid: int,
         timeout_in_ms: int = 5000,
         interval_in_ms: int = 50,
+        *,
+        raise_on_failure: bool = False,
     ) -> Task:
         """Wait until the task fails or succeeds in Meilisearch.
 
@@ -200,6 +202,9 @@ def wait_for_task(
             Time the method should wait before raising a MeilisearchTimeoutError.
         interval_in_ms (optional):
             Time interval the method should wait (sleep) between requests.
+        raise_on_failure (optional):
+            If True, raise a MeilisearchTaskFailedError when the task reaches the
+            ``failed`` status. Defaults to False.
 
         Returns
         -------
@@ -210,12 +215,17 @@ def wait_for_task(
         ------
         MeilisearchTimeoutError
             An error containing details about why Meilisearch can't process your request. Meilisearch error codes are described here: https://www.meilisearch.com/docs/reference/errors/error_codes#meilisearch-errors
+        MeilisearchTaskFailedError
+            Raised when ``raise_on_failure`` is True and the task reaches the
+            ``failed`` status. The exception retains the failed Task and its error mapping.
         """
         start_time = datetime.now()
         elapsed_time = 0.0
         while elapsed_time < timeout_in_ms:
             task = self.get_task(uid)
             if task.status not in ("enqueued", "processing"):
+                if raise_on_failure and task.status == "failed":
+                    raise MeilisearchTaskFailedError(task)
                 return task
             sleep(interval_in_ms / 1000)
             time_delta = datetime.now() - start_time

tests/test_task_failure.py

@@ -0,0 +1,140 @@
+from datetime import datetime, timezone
+from typing import Any, cast
+from unittest.mock import Mock
+
+import pytest
+
+from meilisearch.client import Client
+from meilisearch.errors import MeilisearchTaskFailedError, MeilisearchTimeoutError
+from meilisearch.index import Index
+from meilisearch.models.task import Task
+from meilisearch.task import TaskHandler
+
+
+def make_task(status: str, error: dict | None = None, uid: int = 42) -> Task:
+    timestamp = datetime.now(timezone.utc)
+    return Task(
+        uid=uid,
+        status=status,
+        type="documentAddition",
+        error=error,
+        enqueuedAt=timestamp,
+        startedAt=timestamp,
+        finishedAt=timestamp if status not in ("enqueued", "processing") else None,
+    )
+
+
+def handler_with_tasks(*tasks: Task) -> TaskHandler:
+    handler = cast(Any, TaskHandler.__new__(TaskHandler))
+    handler.get_task = Mock(side_effect=tasks)
+    return cast(TaskHandler, handler)
+
+
+def test_wait_for_task_failed_preserves_default_behavior():
+    task = make_task("failed", {"message": "bad document", "code": "invalid"})
+    handler = handler_with_tasks(task)
+
+    result = handler.wait_for_task(task.uid, interval_in_ms=0)
+
+    assert result is task
+
+
+def test_wait_for_task_failed_raises_and_retains_task():
+    task = make_task("failed", {"message": "bad document", "code": "invalid"})
+    handler = handler_with_tasks(task)
+
+    with pytest.raises(MeilisearchTaskFailedError) as raised:
+        handler.wait_for_task(task.uid, interval_in_ms=0, raise_on_failure=True)
+
+    assert raised.value.task is task
+    assert raised.value.uid == task.uid
+    assert raised.value.error == task.error
+    assert raised.value.message == "bad document"
+
+
+def test_wait_for_task_failed_without_error_raises_useful_exception():
+    task = make_task("failed", None)
+    handler = handler_with_tasks(task)
+
+    with pytest.raises(MeilisearchTaskFailedError, match="Task 42 failed"):
+        handler.wait_for_task(task.uid, interval_in_ms=0, raise_on_failure=True)
+
+
+def test_wait_for_task_failed_without_message_uses_fallback():
+    task = make_task("failed", {"code": "invalid"})
+    handler = handler_with_tasks(task)
+
+    with pytest.raises(MeilisearchTaskFailedError) as raised:
+        handler.wait_for_task(task.uid, interval_in_ms=0, raise_on_failure=True)
+
+    assert raised.value.message == "Task 42 failed"
+
+
+def test_wait_for_task_succeeded_with_raise_enabled_returns_task():
+    task = make_task("succeeded")
+    handler = handler_with_tasks(task)
+
+    assert handler.wait_for_task(task.uid, interval_in_ms=0, raise_on_failure=True) is task
+
+
+def test_wait_for_task_timeout_is_unchanged():
+    handler = handler_with_tasks()
+
+    with pytest.raises(MeilisearchTimeoutError):
+        handler.wait_for_task(42, timeout_in_ms=0, raise_on_failure=True)
+
+    handler.get_task.assert_not_called()
+
+
+def test_wait_for_task_polls_until_failed_then_raises(monkeypatch):
+    import meilisearch.task as task_module
+
+    monkeypatch.setattr(task_module, "sleep", lambda _: None)
+    processing = make_task("processing")
+    failed = make_task("failed", {"code": "invalid"})
+    handler = handler_with_tasks(processing, failed)
+
+    with pytest.raises(MeilisearchTaskFailedError):
+        handler.wait_for_task(42, timeout_in_ms=5000, raise_on_failure=True)
+
+    assert handler.get_task.call_count == 2
+
+
+def test_client_wait_for_task_forwards_raise_on_failure():
+    client = Client.__new__(Client)
+    client.task_handler = Mock()
+
+    client.wait_for_task(42, 5000, 50, raise_on_failure=True)
+
+    client.task_handler.wait_for_task.assert_called_once_with(
+        42, 5000, 50, raise_on_failure=True
+    )
+
+
+def test_client_wait_for_task_preserves_legacy_default_forwarding():
+    client = Client.__new__(Client)
+    client.task_handler = Mock()
+
+    client.wait_for_task(42, 5000, 50)
+
+    client.task_handler.wait_for_task.assert_called_once_with(42, 5000, 50)
+
+
+def test_index_wait_for_task_forwards_raise_on_failure():
+    index = Index.__new__(Index)
+    index.task_handler = Mock()
+
+    index.wait_for_task(42, 5000, 50, raise_on_failure=True)
+
+    index.task_handler.wait_for_task.assert_called_once_with(
+        42, 5000, 50, raise_on_failure=True
+    )
+
+
+def test_index_wait_for_task_preserves_legacy_default_forwarding():
+    index = Index.__new__(Index)
+    index.task_handler = Mock()
+
+    index.wait_for_task(42, 5000, 50)
+
+    index.task_handler.wait_for_task.assert_called_once_with(42, 5000, 50)