Merge pull request #165 from weaviate/jose/fix-parallel-mt-return-collection

Fix parallel MT ingestion returning base collection without tenant context

Author: jfrancoa

.github/workflows/main.yaml

@@ -86,7 +86,7 @@ jobs:
     - name: Run integration tests with pytest
       run: |
         pip install pytest-html
-        pytest test/integration/test_integration.py --html=test-report-${{ matrix.version }}.html --self-contained-html
+        pytest test/integration/test_integration.py test/integration/test_data_integration.py test/integration/test_create_data_return_collection.py --html=test-report-${{ matrix.version }}.html --self-contained-html
   integration-auth-tests:
     needs: [unit-tests, get-latest-weaviate-version]
     env:

test/integration/test_create_data_return_collection.py

@@ -0,0 +1,178 @@
+"""Integration tests verifying that DataManager.create_data returns a usable
+Collection object for both single-tenant and multi-tenant collections.
+
+Regression guard for a bug where parallel multi-tenant ingestion returned the
+base collection (without tenant context), causing ``len(collection)`` and
+``batch.wait_for_vector_indexing()`` to fail with:
+
+    "class X has multi-tenancy enabled, but request was without tenant"
+"""
+
+import pytest
+import weaviate
+from weaviate_cli.managers.collection_manager import CollectionManager
+from weaviate_cli.managers.config_manager import ConfigManager
+from weaviate_cli.managers.data_manager import DataManager
+from weaviate_cli.managers.tenant_manager import TenantManager
+
+NUM_OBJECTS = 20
+NUM_TENANTS = 3
+
+
+@pytest.fixture
+def client() -> weaviate.WeaviateClient:
+    config = ConfigManager()
+    return config.get_client()
+
+
+@pytest.fixture
+def collection_manager(client: weaviate.WeaviateClient) -> CollectionManager:
+    return CollectionManager(client)
+
+
+@pytest.fixture
+def data_manager(client: weaviate.WeaviateClient) -> DataManager:
+    return DataManager(client)
+
+
+@pytest.fixture
+def tenant_manager(client: weaviate.WeaviateClient) -> TenantManager:
+    return TenantManager(client)
+
+
+def test_create_data_returns_usable_collection_single_tenant(
+    client: weaviate.WeaviateClient,
+    collection_manager: CollectionManager,
+    data_manager: DataManager,
+):
+    """For a non-MT collection, the returned collection must support len() and
+    batch.wait_for_vector_indexing() without errors."""
+    collection_name = "TestReturnCollSingleTenant"
+
+    try:
+        if client.collections.exists(collection_name):
+            client.collections.delete(collection_name)
+
+        collection_manager.create_collection(
+            collection=collection_name,
+            vectorizer="none",
+            replication_factor=1,
+            async_enabled=True,
+        )
+
+        returned_col = data_manager.create_data(
+            collection=collection_name,
+            limit=NUM_OBJECTS,
+            consistency_level="one",
+            randomize=True,
+            skip_seed=True,
+            vector_dimensions=128,
+        )
+
+        # The returned collection must be usable for these operations
+        returned_col.batch.wait_for_vector_indexing()
+        assert (
+            len(returned_col) == NUM_OBJECTS
+        ), f"Expected {NUM_OBJECTS} objects via returned collection, got {len(returned_col)}"
+    finally:
+        if client.collections.exists(collection_name):
+            client.collections.delete(collection_name)
+
+
+def test_create_data_returns_usable_collection_multitenant_sequential(
+    client: weaviate.WeaviateClient,
+    collection_manager: CollectionManager,
+    data_manager: DataManager,
+    tenant_manager: TenantManager,
+):
+    """For a multi-tenant collection with sequential ingestion (parallel_workers=1),
+    the returned collection must carry tenant context so len() works."""
+    collection_name = "TestReturnCollMTSeq"
+
+    try:
+        if client.collections.exists(collection_name):
+            client.collections.delete(collection_name)
+
+        collection_manager.create_collection(
+            collection=collection_name,
+            vectorizer="none",
+            replication_factor=1,
+            async_enabled=True,
+            multitenant=True,
+        )
+
+        tenant_manager.create_tenants(
+            collection=collection_name,
+            number_tenants=NUM_TENANTS,
+        )
+
+        returned_col = data_manager.create_data(
+            collection=collection_name,
+            limit=NUM_OBJECTS,
+            consistency_level="one",
+            randomize=True,
+            skip_seed=True,
+            vector_dimensions=128,
+            parallel_workers=1,
+        )
+
+        # Must not raise "multi-tenancy enabled, but request was without tenant"
+        returned_col.batch.wait_for_vector_indexing()
+        count = len(returned_col)
+        assert (
+            count == NUM_OBJECTS
+        ), f"Expected {NUM_OBJECTS} objects via returned collection, got {count}"
+    finally:
+        if client.collections.exists(collection_name):
+            client.collections.delete(collection_name)
+
+
+def test_create_data_returns_usable_collection_multitenant_parallel(
+    client: weaviate.WeaviateClient,
+    collection_manager: CollectionManager,
+    data_manager: DataManager,
+    tenant_manager: TenantManager,
+):
+    """For a multi-tenant collection with parallel ingestion (parallel_workers>1),
+    the returned collection must carry tenant context so len() works.
+
+    This is the exact scenario that was broken: parallel mode never assigned the
+    tenant-scoped collection back to the return variable."""
+    collection_name = "TestReturnCollMTParallel"
+
+    try:
+        if client.collections.exists(collection_name):
+            client.collections.delete(collection_name)
+
+        collection_manager.create_collection(
+            collection=collection_name,
+            vectorizer="none",
+            replication_factor=1,
+            async_enabled=True,
+            multitenant=True,
+        )
+
+        tenant_manager.create_tenants(
+            collection=collection_name,
+            number_tenants=NUM_TENANTS,
+        )
+
+        returned_col = data_manager.create_data(
+            collection=collection_name,
+            limit=NUM_OBJECTS,
+            consistency_level="one",
+            randomize=True,
+            skip_seed=True,
+            vector_dimensions=128,
+            parallel_workers=NUM_TENANTS,
+        )
+
+        # Must not raise "multi-tenancy enabled, but request was without tenant"
+        returned_col.batch.wait_for_vector_indexing()
+        count = len(returned_col)
+        assert (
+            count == NUM_OBJECTS
+        ), f"Expected {NUM_OBJECTS} objects via returned collection, got {count}"
+    finally:
+        if client.collections.exists(collection_name):
+            client.collections.delete(collection_name)

test/integration/test_data_integration.py

@@ -93,7 +93,7 @@ def test_data_creation_with_different_configs(
             assert obj.vector is not None
 
             vector_dimensions = {
-                "transformers": 384,
+                "transformers": 256,
                 "contextionary": 300,
             }
 
@@ -193,11 +193,12 @@ def test_data_creation_with_named_vectors(
             named_vector = obj.vector[named_vector_name]
             assert named_vector is not None
 
-            # Check vector dimensions (should be 768 for transformers)
-            assert len(named_vector) == 384
+            # Check vector dimensions (256 for model2vec transformers, 384 for none/custom)
+            expected_dim = 256 if vectorizer == "transformers" else 384
+            assert len(named_vector) == expected_dim
 
             # Verify vector is not all zeros
-            assert not np.allclose(named_vector, np.zeros(384))
+            assert not np.allclose(named_vector, np.zeros(expected_dim))
 
             # Verify vector has finite values
             assert np.all(np.isfinite(named_vector))

test/unittests/test_managers/test_data_manager.py

@@ -779,6 +779,81 @@ def fake_ingest(collection, **kwargs):
         # Single "None" pseudo-tenant processed
         assert len(processed) == 1
 
+    def test_parallel_returns_tenant_scoped_collection(self, mock_client):
+        """Parallel mode must return a tenant-scoped collection, not the base one.
+
+        Regression test: parallel ingestion discarded the tenant-scoped collection
+        returned by __ingest_data, causing callers to get back the base collection
+        which fails with 'multi-tenancy enabled, but request was without tenant'.
+        """
+        manager = DataManager(mock_client)
+        tenants = ["Tenant-0", "Tenant-1", "Tenant-2"]
+        col = self._make_col(tenants)
+        _setup_mock_client_with_col(mock_client, col)
+
+        def fake_ingest(collection, **kwargs):
+            return collection
+
+        with patch.object(
+            manager, "_DataManager__ingest_data", side_effect=fake_ingest
+        ):
+            result = manager.create_data(
+                collection="TestCollection",
+                limit=5,
+                parallel_workers=4,
+            )
+
+        # The returned collection must NOT be the base col (which has no tenant)
+        assert (
+            result is not col
+        ), "Parallel mode returned the base collection instead of a tenant-scoped one"
+        # It should be one of the tenant-scoped collections
+        col.with_tenant.assert_called()
+
+    def test_sequential_returns_tenant_scoped_collection(self, mock_client):
+        """Sequential mode must also return a tenant-scoped collection."""
+        manager = DataManager(mock_client)
+        tenants = ["Tenant-0", "Tenant-1"]
+        col = self._make_col(tenants)
+        _setup_mock_client_with_col(mock_client, col)
+
+        def fake_ingest(collection, **kwargs):
+            return collection
+
+        with patch.object(
+            manager, "_DataManager__ingest_data", side_effect=fake_ingest
+        ):
+            result = manager.create_data(
+                collection="TestCollection",
+                limit=5,
+                parallel_workers=1,
+            )
+
+        assert (
+            result is not col
+        ), "Sequential mode returned the base collection instead of a tenant-scoped one"
+
+    def test_non_mt_returns_base_collection(self, mock_client):
+        """Non-MT collections should return the base collection (no tenant context)."""
+        manager = DataManager(mock_client)
+        col = _make_non_mt_col()
+        _setup_mock_client_with_col(mock_client, col)
+
+        def fake_ingest(collection, **kwargs):
+            return collection
+
+        with patch.object(
+            manager, "_DataManager__ingest_data", side_effect=fake_ingest
+        ):
+            result = manager.create_data(
+                collection="TestCollection",
+                limit=5,
+                parallel_workers=4,
+            )
+
+        # For non-MT, the base collection IS the correct return value
+        assert result is col
+
 
 # ---------------------------------------------------------------------------
 # create_data – concurrent_requests scaling with parallel_workers

weaviate_cli/managers/data_manager.py

@@ -961,6 +961,7 @@ def _ingest_one_tenant(tenant: str):
                         inserted, _coll = future.result()
                         with _lock:
                             total_inserted += inserted
+                            collection = _coll
                     except Exception as exc:
                         _errors.append(f"Tenant '{t}': {exc}")
             if _errors: