Fix: Milvus error handling (#292)

Add test and improve logging messages

---------

Co-authored-by: Roman Isecke <[email protected]>
Co-authored-by: Filip Knefel <[email protected]>
Co-authored-by: mpolomdeepsense <[email protected]>

Author: 4 people

CHANGELOG.md

@@ -3,6 +3,7 @@
 ### Enhancements
 
 * **Migrate Vectara Destination Connector to v2**
+* **Improved Milvus error handling**
 
 ### Fixes
 * **Register Neo4j Upload Stager**

pyproject.toml

@@ -74,3 +74,6 @@ omit=[
     "unstructured_ingest/ingest_backoff/*",
     "unstructured_ingest/enhanced_dataclass/*"
 ]
+
+[tool.pytest.ini_options]
+asyncio_default_fixture_loop_scope="function"

test/integration/connectors/test_milvus.py

@@ -174,6 +174,19 @@ def test_precheck_fails_on_nonexistent_collection(collection: str):
         uploader.precheck()
 
 
+@pytest.mark.tags(CONNECTOR_TYPE, DESTINATION_TAG)
+def test_precheck_fails_on_nonexisting_db(collection: str):
+    uploader = MilvusUploader(
+        connection_config=MilvusConnectionConfig(uri=DB_URI),
+        upload_config=MilvusUploaderConfig(db_name="nonexisting_db", collection_name=collection),
+    )
+    with pytest.raises(
+        DestinationConnectionError,
+        match="database not found",
+    ):
+        uploader.precheck()
+
+
 @pytest.mark.parametrize("upload_file_str", ["upload_file_ndjson", "upload_file"])
 def test_milvus_stager(
     request: TopRequest,

unstructured_ingest/v2/processes/connectors/chroma.py

@@ -138,7 +138,6 @@ def precheck(self) -> None:
 
     @DestinationConnectionError.wrap
     def upsert_batch(self, collection, batch):
-
         try:
             # Chroma wants lists even if there is only one element
             # Upserting to prevent duplicates

unstructured_ingest/v2/processes/connectors/milvus.py

@@ -156,11 +156,18 @@ class MilvusUploader(Uploader):
 
     @DestinationConnectionError.wrap
     def precheck(self):
-        with self.get_client() as client:
-            if not client.has_collection(self.upload_config.collection_name):
-                raise DestinationConnectionError(
-                    f"Collection '{self.upload_config.collection_name}' does not exist"
-                )
+        from pymilvus import MilvusException
+
+        try:
+            with self.get_client() as client:
+                if not client.has_collection(self.upload_config.collection_name):
+                    raise DestinationConnectionError(
+                        f"Collection '{self.upload_config.collection_name}' does not exist"
+                    )
+        except MilvusException as milvus_exception:
+            raise DestinationConnectionError(
+                f"failed to precheck Milvus: {str(milvus_exception.message)}"
+            ) from milvus_exception
 
     @contextmanager
     def get_client(self) -> Generator["MilvusClient", None, None]:
@@ -197,7 +204,9 @@ def insert_results(self, data: Union[dict, list[dict]]):
             try:
                 res = client.insert(collection_name=self.upload_config.collection_name, data=data)
             except MilvusException as milvus_exception:
-                raise WriteError("failed to upload records to milvus") from milvus_exception
+                raise WriteError(
+                    f"failed to upload records to Milvus: {str(milvus_exception.message)}"
+                ) from milvus_exception
             if "err_count" in res and isinstance(res["err_count"], int) and res["err_count"] > 0:
                 err_count = res["err_count"]
                 raise WriteError(f"failed to upload {err_count} docs")

unstructured_ingest/v2/processes/connectors/qdrant/qdrant.py

@@ -128,7 +128,6 @@ async def run_data_async(
         file_data: FileData,
         **kwargs: Any,
     ) -> None:
-
         batches = list(batch_generator(data, batch_size=self.upload_config.batch_size))
         logger.debug(
             "Elements split into %i batches of size %i.",