From 9e76062bef65db9c567e08299c6a85c14c4ce9db Mon Sep 17 00:00:00 2001 From: Sabari Iyyappan Date: Mon, 13 Jul 2026 07:13:02 -0700 Subject: [PATCH] fix(store): honor _createCollection() return value before indexing DocumentStoreBase.createCollection() discarded the bool returned by the provider-specific _createCollection(), so a driver that reported failure by returning False was ignored: the base proceeded to addChunks() against a collection that was never created, crashing cryptically (e.g. Chroma's "'NoneType' object has no attribute 'delete'") far from the real cause. Capture the return value and raise a clear error when it is exactly False, before indexing. `is False` is deliberate so drivers that return None on success (weaviate, vectordb_postgres) are unaffected; drivers that raise on failure (qdrant, pinecone) propagate their own error before the check. Milvus's _createCollection had inverted semantics (returned False on success, True on failure), which honoring the return value would otherwise break; it is corrected to the documented contract (return True on success, let real failures raise) in the same change. Fixes #1495 --- nodes/src/nodes/milvus/milvus.py | 7 +- packages/ai/src/ai/common/store.py | 14 +++- packages/ai/tests/ai/common/test_store.py | 97 +++++++++++++++++++++++ 3 files changed, 111 insertions(+), 7 deletions(-) create mode 100644 packages/ai/tests/ai/common/test_store.py diff --git a/nodes/src/nodes/milvus/milvus.py b/nodes/src/nodes/milvus/milvus.py index 5927c37da..65b28a364 100644 --- a/nodes/src/nodes/milvus/milvus.py +++ b/nodes/src/nodes/milvus/milvus.py @@ -186,11 +186,8 @@ def _createCollection(self, vectorSize: int) -> bool: # this method is called from the base class createCollection() which already # holds collectionLock. Using doesCollectionExist() would cause a deadlock. if not self._doesCollectionExist(): - try: - self.client.create_collection(collection_name=self.collection, schema=schema, index_params=index_params) - except Exception: - return True - return False + self.client.create_collection(collection_name=self.collection, schema=schema, index_params=index_params) + return True def _convertFilter(self, docFilter: DocFilter) -> str: """ diff --git a/packages/ai/src/ai/common/store.py b/packages/ai/src/ai/common/store.py index f3b7aaebf..c08899f55 100644 --- a/packages/ai/src/ai/common/store.py +++ b/packages/ai/src/ai/common/store.py @@ -39,7 +39,10 @@ def _createCollection() -> bool: """ Create the collection. - This the abstract method that the driver must implement + This the abstract method that the driver must implement. Implementations + must either raise on failure or explicitly return False; the caller + (createCollection) treats a return value of exactly False as failure and + aborts before indexing any documents. """ @abstractmethod @@ -515,7 +518,14 @@ def createCollection(self, documents: List[Doc]): doc.embedding = [0] * vectorSize # List of zeros for validation purposes # Create the actual collection with the specified vector size - self._createCollection(vectorSize) + created = self._createCollection(vectorSize) + + # Some drivers signal failure by returning False instead of raising; + # honor that instead of silently indexing into a collection that was + # never created (do not treat None as failure - a few drivers return + # nothing on success). + if created is False: + raise Exception(f'{type(self).__name__} failed to create the vector collection') # Add the "bogus" document to the collection self.addChunks([doc], checkCollection=False) diff --git a/packages/ai/tests/ai/common/test_store.py b/packages/ai/tests/ai/common/test_store.py new file mode 100644 index 000000000..346f10fa0 --- /dev/null +++ b/packages/ai/tests/ai/common/test_store.py @@ -0,0 +1,97 @@ +""" +Unit tests for ai.common.store.DocumentStoreBase.createCollection. + +Regression coverage for issue #1495: createCollection() must not silently proceed +to index documents when the driver's _createCollection() reports failure via a +bare `return False` (as opposed to raising). +""" + +from unittest.mock import MagicMock + +import pytest + +from ai.common.schema import Doc, DocMetadata +from ai.common.store import DocumentStoreBase + + +class _FakeStore(DocumentStoreBase): + """Minimal concrete DocumentStoreBase for exercising createCollection() in isolation.""" + + def __init__(self, create_result): + super().__init__('fake', {}, {}) + self._create_result = create_result + # Shadow the concrete addChunks below with a mock so tests can assert on + # the calls. The class-level definition is what satisfies the ABC. + self.addChunks = MagicMock() + + def _doesCollectionExist(self): + return False + + def _createCollection(self, vectorSize): + if isinstance(self._create_result, Exception): + raise self._create_result + return self._create_result + + def addChunks(self, chunks, checkCollection=True): + # Concrete implementation so the ABC is instantiable; each instance + # replaces this with a MagicMock in __init__. + return None + + def count_documents(self): + return 0 + + def searchKeyword(self, query, docFilter): + return [] + + def searchSemantic(self, query, docFilter): + return [] + + def get(self, docFilter, checkCollection=True): + return [] + + def getPaths(self, parent=None, offset=0, limit=1000): + return {} + + def remove(self, objectIds): + return None + + def markDeleted(self, objectIds): + return None + + def markActive(self, objectIds): + return None + + def render(self, objectId, callback): + return None + + +def _make_docs(): + metadata = DocMetadata(objectId='obj', chunkId=0) + doc = Doc(page_content='hello', metadata=metadata, embedding=[0.1, 0.2, 0.3], score=0.0) + doc.embedding_model = 'test-model' + return [doc] + + +def test_create_collection_raises_when_driver_returns_false(): + store = _FakeStore(create_result=False) + + with pytest.raises(Exception): + store.createCollection(_make_docs()) + + store.addChunks.assert_not_called() + + +def test_create_collection_succeeds_when_driver_returns_true(): + store = _FakeStore(create_result=True) + + assert store.createCollection(_make_docs()) is True + store.addChunks.assert_called_once() + + +def test_create_collection_propagates_driver_exception(): + store = _FakeStore(create_result=RuntimeError('boom')) + + with pytest.raises(RuntimeError, match='boom'): + store.createCollection(_make_docs()) + + store.addChunks.assert_not_called()