Merge branch 'dev/1.37' into feat/tokenizer-endpoint

Author: amourao

.github/workflows/main.yaml

@@ -28,7 +28,7 @@ env:
   WEAVIATE_134: 1.34.19
   WEAVIATE_135: 1.35.16-efdedfa
   WEAVIATE_136: 1.36.9-d905e6c
-  WEAVIATE_137: 1.37.0-rc.1-578c4eb
+  WEAVIATE_137: 1.37.0-rc.1-bc3891e
 
 jobs:
   lint-and-format:

ci/docker-compose.yml

@@ -31,6 +31,8 @@ services:
       DISABLE_LAZY_LOAD_SHARDS: 'true'
       GRPC_MAX_MESSAGE_SIZE: 100000000  # 100mb
       OBJECTS_TTL_DELETE_SCHEDULE: "@every 12h" # for objectTTL tests to work
+      EXPORT_ENABLED: 'true'
+      EXPORT_DEFAULT_PATH: "/var/lib/weaviate/exports"
 
   contextionary:
     environment:

integration/test_export.py

@@ -0,0 +1,222 @@
+import time
+import uuid
+from typing import Generator, List, Union
+
+import pytest
+from _pytest.fixtures import SubRequest
+
+import weaviate
+from weaviate.collections.classes.config import DataType, Property
+from weaviate.exceptions import UnexpectedStatusCodeException
+from weaviate.export.export import (
+    ExportFileFormat,
+    ExportStatus,
+    ExportStorage,
+)
+
+from .conftest import _sanitize_collection_name
+
+pytestmark = pytest.mark.xdist_group(name="export")
+
+BACKEND = ExportStorage.FILESYSTEM
+
+COLLECTION_NAME = "ExportTestCollection"
+
+OBJECT_PROPS = [{"title": f"object {i}", "count": i} for i in range(5)]
+
+OBJECT_IDS = [
+    "fd34ccf4-1a2a-47ad-8446-231839366c3f",
+    "2653442b-05d8-4fa3-b46a-d4a152eb63bc",
+    "55374edb-17de-487f-86cb-9a9fbc30823f",
+    "124ff6aa-597f-44d0-8c13-62fbb1e66888",
+    "f787386e-7d1c-481f-b8c3-3dbfd8bbad85",
+]
+
+
+@pytest.fixture(scope="module")
+def client() -> Generator[weaviate.WeaviateClient, None, None]:
+    client = weaviate.connect_to_local()
+    if client._connection._weaviate_version.is_lower_than(1, 37, 0):
+        client.close()
+        pytest.skip("Collection export is not supported in versions lower than 1.37.0")
+    client.collections.delete(COLLECTION_NAME)
+
+    col = client.collections.create(
+        name=COLLECTION_NAME,
+        properties=[
+            Property(name="title", data_type=DataType.TEXT),
+            Property(name="count", data_type=DataType.INT),
+        ],
+    )
+    for i, props in enumerate(OBJECT_PROPS):
+        col.data.insert(properties=props, uuid=OBJECT_IDS[i])
+
+    yield client
+    client.collections.delete(COLLECTION_NAME)
+    client.close()
+
+
+def unique_export_id(name: str) -> str:
+    """Generate a unique export ID based on the test name."""
+    name = _sanitize_collection_name(name)
+    random_part = str(uuid.uuid4()).replace("-", "")[:12]
+    return name + random_part
+
+
+def test_create_export_with_waiting(client: weaviate.WeaviateClient, request: SubRequest) -> None:
+    """Create an export with wait_for_completion=True."""
+    export_id = unique_export_id(request.node.name)
+
+    resp = client.export.create(
+        export_id=export_id,
+        backend=BACKEND,
+        file_format=ExportFileFormat.PARQUET,
+        include_collections=[COLLECTION_NAME],
+        wait_for_completion=True,
+    )
+    assert resp.status == ExportStatus.SUCCESS
+    assert COLLECTION_NAME in resp.collections
+
+
+def test_create_export_without_waiting(
+    client: weaviate.WeaviateClient, request: SubRequest
+) -> None:
+    """Create an export without waiting, then poll status."""
+    export_id = unique_export_id(request.node.name)
+
+    resp = client.export.create(
+        export_id=export_id,
+        backend=BACKEND,
+        file_format=ExportFileFormat.PARQUET,
+        include_collections=[COLLECTION_NAME],
+    )
+    assert resp.status in [ExportStatus.STARTED, ExportStatus.TRANSFERRING, ExportStatus.SUCCESS]
+
+    # poll until done
+    while True:
+        status = client.export.get_status(export_id=export_id, backend=BACKEND)
+        assert status.status in [
+            ExportStatus.STARTED,
+            ExportStatus.TRANSFERRING,
+            ExportStatus.SUCCESS,
+        ]
+        if status.status == ExportStatus.SUCCESS:
+            break
+        time.sleep(0.1)
+
+    assert status.export_id.lower() == export_id.lower()
+
+
+def test_get_export_status(client: weaviate.WeaviateClient, request: SubRequest) -> None:
+    """Check status of a completed export."""
+    export_id = unique_export_id(request.node.name)
+
+    client.export.create(
+        export_id=export_id,
+        backend=BACKEND,
+        file_format=ExportFileFormat.PARQUET,
+        include_collections=[COLLECTION_NAME],
+        wait_for_completion=True,
+    )
+
+    status = client.export.get_status(export_id=export_id, backend=BACKEND)
+    assert status.status == ExportStatus.SUCCESS
+    assert status.export_id.lower() == export_id.lower()
+    assert status.backend == BACKEND.value
+
+
+def test_create_export_with_parquet_format(
+    client: weaviate.WeaviateClient, request: SubRequest
+) -> None:
+    """Create an export explicitly specifying parquet format."""
+    export_id = unique_export_id(request.node.name)
+
+    resp = client.export.create(
+        export_id=export_id,
+        backend=BACKEND,
+        file_format=ExportFileFormat.PARQUET,
+        include_collections=[COLLECTION_NAME],
+        wait_for_completion=True,
+    )
+    assert resp.status == ExportStatus.SUCCESS
+
+
+@pytest.mark.parametrize("include", [[COLLECTION_NAME], COLLECTION_NAME])
+def test_create_export_include_as_str_and_list(
+    client: weaviate.WeaviateClient, include: Union[str, List[str]], request: SubRequest
+) -> None:
+    """Verify include_collections accepts both str and list."""
+    export_id = unique_export_id(request.node.name)
+
+    resp = client.export.create(
+        export_id=export_id,
+        backend=BACKEND,
+        file_format=ExportFileFormat.PARQUET,
+        include_collections=include,
+        wait_for_completion=True,
+    )
+    assert resp.status == ExportStatus.SUCCESS
+    assert COLLECTION_NAME in resp.collections
+
+
+def test_cancel_export(client: weaviate.WeaviateClient, request: SubRequest) -> None:
+    """Cancel a running export."""
+    export_id = unique_export_id(request.node.name)
+
+    resp = client.export.create(
+        export_id=export_id,
+        backend=BACKEND,
+        file_format=ExportFileFormat.PARQUET,
+        include_collections=[COLLECTION_NAME],
+    )
+    assert resp.status in [ExportStatus.STARTED, ExportStatus.TRANSFERRING, ExportStatus.SUCCESS]
+
+    client.export.cancel(export_id=export_id, backend=BACKEND)
+
+    # verify it's cancelled or already completed (race condition)
+    start = time.time()
+    while time.time() - start < 5:
+        status = client.export.get_status(export_id=export_id, backend=BACKEND)
+        if status.status in [ExportStatus.CANCELED, ExportStatus.SUCCESS]:
+            break
+        time.sleep(0.1)
+    assert status.status in [ExportStatus.CANCELED, ExportStatus.SUCCESS]
+
+
+def test_fail_on_non_existing_collection(
+    client: weaviate.WeaviateClient, request: SubRequest
+) -> None:
+    """Fail export on non-existing collection."""
+    export_id = unique_export_id(request.node.name)
+    with pytest.raises(UnexpectedStatusCodeException):
+        client.export.create(
+            export_id=export_id,
+            backend=BACKEND,
+            file_format=ExportFileFormat.PARQUET,
+            include_collections=["NonExistingCollection"],
+            wait_for_completion=True,
+        )
+
+
+def test_fail_on_both_include_and_exclude(
+    client: weaviate.WeaviateClient, request: SubRequest
+) -> None:
+    """Fail when both include and exclude collections are set."""
+    export_id = unique_export_id(request.node.name)
+    with pytest.raises(ValueError):
+        client.export.create(
+            export_id=export_id,
+            backend=BACKEND,
+            file_format=ExportFileFormat.PARQUET,
+            include_collections=COLLECTION_NAME,
+            exclude_collections="SomeOther",
+        )
+
+
+def test_fail_status_for_non_existing_export(
+    client: weaviate.WeaviateClient, request: SubRequest
+) -> None:
+    """Fail checking status for non-existing export."""
+    export_id = unique_export_id(request.node.name)
+    with pytest.raises(UnexpectedStatusCodeException):
+        client.export.get_status(export_id=export_id, backend=BACKEND)

weaviate/classes/__init__.py

@@ -5,6 +5,7 @@
     batch,
     config,
     data,
+    export,
     generate,
     generics,
     init,
@@ -22,6 +23,7 @@
     "config",
     "ConsistencyLevel",
     "data",
+    "export",
     "generate",
     "generics",
     "init",

weaviate/classes/export.py

@@ -0,0 +1,9 @@
+from weaviate.export.export import (
+    ExportFileFormat,
+    ExportStorage,
+)
+
+__all__ = [
+    "ExportFileFormat",
+    "ExportStorage",
+]

weaviate/client.py

@@ -20,6 +20,7 @@
 from .connect.v4 import ConnectionAsync, ConnectionSync
 from .debug import _Debug, _DebugAsync
 from .embedded import EmbeddedOptions
+from .export import _Export, _ExportAsync
 from .groups import _Groups, _GroupsAsync
 from .rbac import _Roles, _RolesAsync
 from .tokenization import _Tokenization, _TokenizationAsync
@@ -77,6 +78,7 @@ def __init__(
         )
         self.alias = _AliasAsync(self._connection)
         self.backup = _BackupAsync(self._connection)
+        self.export = _ExportAsync(self._connection)
         self.batch = _BatchClientWrapperAsync(self._connection)
         self.cluster = _ClusterAsync(self._connection)
         self.collections = _CollectionsAsync(self._connection)
@@ -154,6 +156,7 @@ def __init__(
             consistency_level=None,
         )
         self.backup = _Backup(self._connection)
+        self.export = _Export(self._connection)
         self.cluster = _Cluster(self._connection)
         self.collections = collections
         self.debug = _Debug(self._connection)

weaviate/client.pyi

@@ -20,6 +20,7 @@ from .backup import _Backup, _BackupAsync
 from .cluster import _Cluster, _ClusterAsync
 from .collections.batch.client import _BatchClientWrapper, _BatchClientWrapperAsync
 from .debug import _Debug, _DebugAsync
+from .export import _Export, _ExportAsync
 from .rbac import _Roles, _RolesAsync
 from .tokenization import _Tokenization, _TokenizationAsync
 from .types import NUMBER
@@ -30,6 +31,7 @@ class WeaviateAsyncClient(_WeaviateClientExecutor[ConnectionAsync]):
     _connection: ConnectionAsync
     alias: _AliasAsync
     backup: _BackupAsync
+    export: _ExportAsync
     batch: _BatchClientWrapperAsync
     collections: _CollectionsAsync
     cluster: _ClusterAsync
@@ -54,6 +56,7 @@ class WeaviateClient(_WeaviateClientExecutor[ConnectionSync]):
     _connection: ConnectionSync
     alias: _Alias
     backup: _Backup
+    export: _Export
     batch: _BatchClientWrapper
     collections: _Collections
     cluster: _Cluster

weaviate/collections/collections/executor.py

@@ -44,13 +44,13 @@
     _check_references_generic,
 )
 from weaviate.collections.collection import Collection, CollectionAsync
+from weaviate.collections.config.executor import _any_property_has_text_analyzer
 from weaviate.connect import executor
 from weaviate.connect.v4 import (
     ConnectionAsync,
     ConnectionType,
     _ExpectedStatusCodes,
 )
-from weaviate.collections.config.executor import _any_property_has_text_analyzer
 from weaviate.exceptions import WeaviateInvalidInputError, WeaviateUnsupportedFeatureError
 from weaviate.util import _capitalize_first_letter, _decode_json_response_dict
 from weaviate.validator import _validate_input, _ValidateArgument

weaviate/collections/config/executor.py

@@ -5,8 +5,8 @@
     Generic,
     List,
     Literal,
-    Sequence,
     Optional,
+    Sequence,
     Tuple,
     Union,
     cast,

weaviate/exceptions.py

@@ -141,6 +141,14 @@ class BackupCanceledError(WeaviateBaseError):
     """Backup canceled Exception."""
 
 
+class ExportFailedError(WeaviateBaseError):
+    """Export Failed Exception."""
+
+
+class ExportCanceledError(WeaviateBaseError):
+    """Export Canceled Exception."""
+
+
 class EmptyResponseError(WeaviateBaseError):
     """Occurs when an HTTP request unexpectedly returns an empty response."""