Merge pull request #1206 from sanders41/get-batches

Add option to pass parameters to get_batches

Author: sanders41

meilisearch_python_sdk/_batch.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from datetime import datetime
 from typing import TYPE_CHECKING
 
 from meilisearch_python_sdk._utils import get_async_client, get_client
@@ -28,9 +29,38 @@ async def async_get_batch(
     return BatchResult(**response.json())
 
 
-async def async_get_batches(client: HttpxAsyncClient | AsyncClient) -> BatchStatus:
+async def async_get_batches(
+    client: HttpxAsyncClient | AsyncClient,
+    *,
+    uids: list[int] | None = None,
+    batch_uids: list[int] | None = None,
+    index_uids: list[int] | None = None,
+    statuses: list[str] | None = None,
+    types: list[str] | None = None,
+    limit: int = 20,
+    from_: str | None = None,
+    reverse: bool = False,
+    before_enqueued_at: datetime | None = None,
+    after_enqueued_at: datetime | None = None,
+    before_started_at: datetime | None = None,
+    after_finished_at: datetime | None = None,
+) -> BatchStatus:
     client_ = get_async_client(client)
-    response = await client_.get("batches")
+    params = _build_parameters(
+        uids=uids,
+        batch_uids=batch_uids,
+        index_uids=index_uids,
+        statuses=statuses,
+        types=types,
+        limit=limit,
+        from_=from_,
+        reverse=reverse,
+        before_enqueued_at=before_enqueued_at,
+        after_enqueued_at=after_enqueued_at,
+        before_started_at=before_started_at,
+        after_finished_at=after_finished_at,
+    )
+    response = await client_.get("batches", params=params)
 
     return BatchStatus(**response.json())
 
@@ -45,8 +75,92 @@ def get_batch(client: HttpxClient | Client, batch_uid: int) -> BatchResult | Non
     return BatchResult(**response.json())
 
 
-def get_batches(client: HttpxClient | Client) -> BatchStatus:
+def get_batches(
+    client: HttpxClient | Client,
+    *,
+    uids: list[int] | None = None,
+    batch_uids: list[int] | None = None,
+    index_uids: list[int] | None = None,
+    statuses: list[str] | None = None,
+    types: list[str] | None = None,
+    limit: int = 20,
+    from_: str | None = None,
+    reverse: bool = False,
+    before_enqueued_at: datetime | None = None,
+    after_enqueued_at: datetime | None = None,
+    before_started_at: datetime | None = None,
+    after_finished_at: datetime | None = None,
+) -> BatchStatus:
     client_ = get_client(client)
-    response = client_.get("batches")
+    params = _build_parameters(
+        uids=uids,
+        batch_uids=batch_uids,
+        index_uids=index_uids,
+        statuses=statuses,
+        types=types,
+        limit=limit,
+        from_=from_,
+        reverse=reverse,
+        before_enqueued_at=before_enqueued_at,
+        after_enqueued_at=after_enqueued_at,
+        before_started_at=before_started_at,
+        after_finished_at=after_finished_at,
+    )
+
+    response = client_.get("batches", params=params)
 
     return BatchStatus(**response.json())
+
+
+def _build_parameters(
+    *,
+    uids: list[int] | None = None,
+    batch_uids: list[int] | None = None,
+    index_uids: list[int] | None = None,
+    statuses: list[str] | None = None,
+    types: list[str] | None = None,
+    limit: int = 20,
+    from_: str | None = None,
+    reverse: bool = False,
+    before_enqueued_at: datetime | None = None,
+    after_enqueued_at: datetime | None = None,
+    before_started_at: datetime | None = None,
+    after_finished_at: datetime | None = None,
+) -> dict[str, str]:
+    params = {}
+
+    if uids:
+        params["uids"] = ",".join([str(uid) for uid in uids])
+
+    if batch_uids:  # pragma: no cover
+        params["batchUids"] = ",".join([str(uid) for uid in batch_uids])
+
+    if index_uids:  # pragma: no cover
+        params["indexUids"] = ",".join([str(uid) for uid in index_uids])
+
+    if statuses:  # pragma: no cover
+        params["statuses"] = ",".join(statuses)
+
+    if types:  # pragma: no cover
+        params["types"] = ",".join(types)
+
+    params["limit"] = str(limit)
+
+    if from_:  # pragma: no cover
+        params["from"] = from_
+
+    params["reverse"] = "true" if reverse else "false"
+
+    if before_enqueued_at:  # pragma: no cover
+        params["beforeEnqueuedAt"] = before_enqueued_at.isoformat()
+
+    if after_enqueued_at:  # pragma: no cover
+        params["afterEnqueuedAt"] = after_enqueued_at.isoformat()
+
+    if before_started_at:  # pragma: no cover
+        params["beforeStartedAt"] = before_started_at.isoformat()
+
+    if after_finished_at:  # pragma: no cover
+        params["afterFinishedAt"] = after_finished_at.isoformat()
+
+    return params

meilisearch_python_sdk/_client.py

@@ -9,7 +9,9 @@
 from httpx import Client as HttpxClient
 
 from meilisearch_python_sdk import _task
-from meilisearch_python_sdk._batch import async_get_batch, async_get_batches, get_batch, get_batches
+from meilisearch_python_sdk._batch import async_get_batch, async_get_batches
+from meilisearch_python_sdk._batch import get_batch as _get_batch
+from meilisearch_python_sdk._batch import get_batches as _get_batches
 from meilisearch_python_sdk._http_requests import AsyncHttpRequests, HttpRequests
 from meilisearch_python_sdk.errors import InvalidRestriction, MeilisearchApiError
 from meilisearch_python_sdk.index import AsyncIndex, Index
@@ -775,8 +777,37 @@ async def swap_indexes(self, indexes: list[tuple[str, str]]) -> TaskInfo:
     async def get_batch(self, batch_uid: int) -> BatchResult | None:
         return await async_get_batch(self, batch_uid)
 
-    async def get_batches(self) -> BatchStatus:
-        return await async_get_batches(self)
+    async def get_batches(
+        self,
+        *,
+        uids: list[int] | None = None,
+        batch_uids: list[int] | None = None,
+        index_uids: list[int] | None = None,
+        statuses: list[str] | None = None,
+        types: list[str] | None = None,
+        limit: int = 20,
+        from_: str | None = None,
+        reverse: bool = False,
+        before_enqueued_at: datetime | None = None,
+        after_enqueued_at: datetime | None = None,
+        before_started_at: datetime | None = None,
+        after_finished_at: datetime | None = None,
+    ) -> BatchStatus:
+        return await async_get_batches(
+            self,
+            uids=uids,
+            batch_uids=batch_uids,
+            index_uids=index_uids,
+            statuses=statuses,
+            types=types,
+            limit=limit,
+            from_=from_,
+            reverse=reverse,
+            before_enqueued_at=before_enqueued_at,
+            after_enqueued_at=after_enqueued_at,
+            before_started_at=before_started_at,
+            after_finished_at=after_finished_at,
+        )
 
     async def cancel_tasks(
         self,
@@ -1599,10 +1630,39 @@ def swap_indexes(self, indexes: list[tuple[str, str]]) -> TaskInfo:
         return TaskInfo(**response.json())
 
     def get_batch(self, batch_uid: int) -> BatchResult | None:
-        return get_batch(self, batch_uid)
+        return _get_batch(self, batch_uid)
 
-    def get_batches(self) -> BatchStatus:
-        return get_batches(self)
+    def get_batches(
+        self,
+        *,
+        uids: list[int] | None = None,
+        batch_uids: list[int] | None = None,
+        index_uids: list[int] | None = None,
+        statuses: list[str] | None = None,
+        types: list[str] | None = None,
+        limit: int = 20,
+        from_: str | None = None,
+        reverse: bool = False,
+        before_enqueued_at: datetime | None = None,
+        after_enqueued_at: datetime | None = None,
+        before_started_at: datetime | None = None,
+        after_finished_at: datetime | None = None,
+    ) -> BatchStatus:
+        return _get_batches(
+            self,
+            uids=uids,
+            batch_uids=batch_uids,
+            index_uids=index_uids,
+            statuses=statuses,
+            types=types,
+            limit=limit,
+            from_=from_,
+            reverse=reverse,
+            before_enqueued_at=before_enqueued_at,
+            after_enqueued_at=after_enqueued_at,
+            before_started_at=before_started_at,
+            after_finished_at=after_finished_at,
+        )
 
     def cancel_tasks(
         self,

tests/test_async_client.py

@@ -1021,6 +1021,17 @@ async def test_get_batches(async_client, async_empty_index, small_movies):
     assert len(result.results) > 0
 
 
+async def test_get_batches_with_params(async_client, async_empty_index, small_movies):
+    # Add documents to create batches
+    index = await async_empty_index()
+    tasks = await index.add_documents_in_batches(small_movies, batch_size=5)
+    waits = [async_client.wait_for_task(x.task_uid) for x in tasks]
+    await asyncio.gather(*waits)
+
+    result = await async_client.get_batches(uids=[task.task_uid for task in tasks])
+    assert len(result.results) > 0
+
+
 async def test_get_batch(async_client, async_empty_index, small_movies):
     # Add documents to create batches
     index = await async_empty_index()

tests/test_client.py

@@ -1002,6 +1002,17 @@ def test_get_batches(client, empty_index, small_movies):
     assert len(result.results) > 0
 
 
+def test_get_batches_with_params(client, empty_index, small_movies):
+    # Add documents to create batches
+    index = empty_index()
+    tasks = index.add_documents_in_batches(small_movies, batch_size=5)
+    for task in tasks:
+        client.wait_for_task(task.task_uid)
+
+    result = client.get_batches(uids=[task.task_uid for task in tasks])
+    assert len(result.results) > 0
+
+
 def test_get_batch(client, empty_index, small_movies):
     # Add documents to create batches
     index = empty_index()