Merge pull request #1164 from jawad-khan/jawad/sharding-support

Add sharding support

Author: Strift

meilisearch/models/task.py

@@ -21,6 +21,7 @@ class Task(CamelBase):
     enqueued_at: datetime
     started_at: Optional[datetime] = None
     finished_at: Optional[datetime] = None
+    network: Optional[Dict[str, Any]] = None
 
     if is_pydantic_2():
 

tests/client/test_client_multi_search_meilisearch.py

@@ -2,6 +2,7 @@
 
 from meilisearch.errors import MeilisearchApiError
 from tests.common import INDEX_UID, REMOTE_MS_1, REMOTE_MS_2
+from tests.test_utils import disable_sharding
 
 
 def test_basic_multi_search(client, empty_index):
@@ -84,14 +85,17 @@ def test_multi_search_with_network(client, index_with_documents):
     resp = client.add_or_update_networks(
         {
             "self": REMOTE_MS_1,
+            "sharding": True,
             "remotes": {
                 REMOTE_MS_1: {
                     "url": "http://ms-1235.example.meilisearch.io",
                     "searchApiKey": "xxxxxxxx",
+                    "writeApiKey": "xxxxxxxx",
                 },
                 REMOTE_MS_2: {
                     "url": "http://ms-1255.example.meilisearch.io",
                     "searchApiKey": "xxxxxxxx",
+                    "writeApiKey": "xxxxxxxx",
                 },
             },
         }
@@ -108,3 +112,4 @@ def test_multi_search_with_network(client, index_with_documents):
     assert response["hits"][0]["_federation"]["indexUid"] == INDEX_UID
     assert response["hits"][0]["_federation"]["remote"] == REMOTE_MS_1
     assert response["remoteErrors"] == {}
+    disable_sharding(client)

tests/client/test_client_network.py

@@ -1,6 +1,7 @@
 import pytest
 
 from tests.common import REMOTE_MS_1, REMOTE_MS_2
+from tests.test_utils import disable_sharding
 
 
 @pytest.mark.usefixtures("enable_network_options")
@@ -16,15 +17,27 @@ def test_add_or_update_networks(client):
     """Tests upsert network remote instance."""
     body = {
         "self": REMOTE_MS_1,
+        "sharding": True,
         "remotes": {
-            REMOTE_MS_1: {"url": "http://localhost:7700", "searchApiKey": "xxxxxxxxxxxxxx"},
-            REMOTE_MS_2: {"url": "http://localhost:7720", "searchApiKey": "xxxxxxxxxxxxxxx"},
+            REMOTE_MS_1: {
+                "url": "http://localhost:7700",
+                "searchApiKey": "xxxxxxxxxxxxxx",
+                "writeApiKey": "xxxxxxxxx",
+            },
+            REMOTE_MS_2: {
+                "url": "http://localhost:7720",
+                "searchApiKey": "xxxxxxxxxxxxxxx",
+                "writeApiKey": "xxxxxxxx",
+            },
         },
     }
     response = client.add_or_update_networks(body=body)
 
     assert isinstance(response, dict)
     assert response["self"] == REMOTE_MS_1
+    assert response["sharding"] is True
     assert len(response["remotes"]) >= 2
     assert REMOTE_MS_2 in response["remotes"]
     assert REMOTE_MS_1 in response["remotes"]
+
+    disable_sharding(client)

tests/client/test_client_sharding.py

@@ -0,0 +1,37 @@
+import pytest
+
+from tests.common import BASE_URL, REMOTE_MS_1
+from tests.test_utils import disable_sharding
+
+
+@pytest.mark.usefixtures("enable_network_options")
+def test_update_and_get_network_settings(client):
+    """Test updating and getting network settings."""
+    instance_name = REMOTE_MS_1
+    options = {
+        "self": instance_name,
+        "remotes": {
+            instance_name: {
+                "url": BASE_URL,
+                "searchApiKey": "search-key-1",
+                "writeApiKey": "write-key-1",
+            }
+        },
+        "sharding": True,
+    }
+
+    client.add_or_update_networks(options)
+    response = client.get_all_networks()
+
+    assert response["self"] == options["self"]
+    assert response["remotes"][instance_name]["url"] == options["remotes"][instance_name]["url"]
+    assert (
+        response["remotes"][instance_name]["searchApiKey"]
+        == options["remotes"][instance_name]["searchApiKey"]
+    )
+    assert (
+        response["remotes"][instance_name]["writeApiKey"]
+        == options["remotes"][instance_name]["writeApiKey"]
+    )
+    assert response["sharding"] == options["sharding"]
+    disable_sharding(client)

tests/test_utils.py

@@ -31,3 +31,8 @@ def test_iso_to_date_time(iso_date, expected):
 def test_iso_to_date_time_invalid_format():
     with pytest.raises(ValueError):
         iso_to_date_time("2023-07-13T23:37:20Z")
+
+
+# Refactor to use the unified API to toggle experimental features
+def disable_sharding(client):
+    client.add_or_update_networks(body={"sharding": False})