Add more batching tests, export wrapping types from more logical locations

Author: tsmith023

integration/test_batch_v4.py

@@ -1,3 +1,4 @@
+import concurrent.futures
 import uuid
 from dataclasses import dataclass
 from typing import Generator, List, Optional, Protocol, Tuple, Callable
@@ -6,6 +7,7 @@
 from _pytest.fixtures import SubRequest
 
 import weaviate
+from weaviate import BatchClient, ClientBatchingContextManager
 from integration.conftest import _sanitize_collection_name
 from weaviate.collections.classes.batch import Shard
 from weaviate.collections.classes.config import (
@@ -20,7 +22,6 @@
     ReferenceToMulti,
 )
 from weaviate.collections.classes.tenants import Tenant
-from weaviate.outputs.batch import ClientBatchingContextManager
 from weaviate.types import UUID, VECTORS
 
 UUID1 = uuid.UUID("806827e0-2b31-43ca-9269-24fa95a221f9")
@@ -575,6 +576,29 @@ def test_add_one_object_and_a_self_reference(
     assert obj.references["test"].objects[0].uuid == uuid
 
 
+def test_multi_threaded_batching(
+    client_factory: ClientFactory,
+) -> None:
+    client, name = client_factory()
+    nr_objects = 1000
+    nr_threads = 10
+
+    def batch_insert(batch: BatchClient) -> None:
+        for i in range(nr_objects):
+            batch.add_object(
+                collection=name,
+                properties={"name": "test" + str(i)},
+            )
+
+    with concurrent.futures.ThreadPoolExecutor() as executor:
+        with client.batch.dynamic() as batch:
+            futures = [executor.submit(batch_insert, batch) for _ in range(nr_threads)]
+    for future in concurrent.futures.as_completed(futures):
+        future.result()
+    objs = client.collections.get(name).query.fetch_objects(limit=nr_objects * nr_threads).objects
+    assert len(objs) == nr_objects * nr_threads
+
+
 def test_error_reset(client_factory: ClientFactory) -> None:
     client, name = client_factory()
     with client.batch.dynamic() as batch:

weaviate/__init__.py

@@ -12,6 +12,7 @@
     __version__ = "unknown version"
 
 from .client import Client, WeaviateAsyncClient, WeaviateClient
+from .collections.batch.client import BatchClient, ClientBatchingContextManager
 from .connect.helpers import (
     connect_to_custom,
     connect_to_embedded,
@@ -45,6 +46,8 @@
 from .warnings import _Warnings
 
 __all__ = [
+    "BatchClient",
+    "ClientBatchingContextManager",
     "Client",
     "WeaviateClient",
     "WeaviateAsyncClient",

weaviate/collections/__init__.py

@@ -1,3 +1,4 @@
+from weaviate.collections.batch.collection import BatchCollection, CollectionBatchingContextManager
 from weaviate.collections.collection import Collection, CollectionAsync
 
-__all__ = ["Collection", "CollectionAsync"]
+__all__ = ["BatchCollection", "Collection", "CollectionAsync", "CollectionBatchingContextManager"]

weaviate/collections/batch/client.py

@@ -115,7 +115,8 @@ def add_reference(
         )
 
 
-ClientBatchingContextManager = _ContextManagerWrapper[_BatchClient]
+BatchClient = _BatchClient
+ClientBatchingContextManager = _ContextManagerWrapper[BatchClient]
 
 
 class _BatchClientWrapper(_BatchWrapper):

weaviate/collections/batch/collection.py

@@ -114,7 +114,8 @@ def add_reference(
         )
 
 
-CollectionBatchingContextManager = _ContextManagerWrapper[_BatchCollection[Properties]]
+BatchCollection = _BatchCollection[Properties]
+CollectionBatchingContextManager = _ContextManagerWrapper[BatchCollection[Properties]]
 
 
 class _BatchCollectionWrapper(Generic[Properties], _BatchWrapper):

weaviate/outputs/batch.py

@@ -5,15 +5,11 @@
     ErrorObject,
     ErrorReference,
 )
-from weaviate.collections.batch.client import ClientBatchingContextManager
-from weaviate.collections.batch.collection import CollectionBatchingContextManager
 
 __all__ = [
     "BatchObjectReturn",
     "BatchReferenceReturn",
     "BatchResult",
     "ErrorObject",
     "ErrorReference",
-    "ClientBatchingContextManager",
-    "CollectionBatchingContextManager",
 ]