feat/use latest astradb client version (#457)

* use latest astradb client version

* update _collection_exists method

* fix everything that broke with version update

* revert bae chang

* revert change

Author: rbiseck3

CHANGELOG.md

@@ -1,3 +1,9 @@
+## 1.0.2
+
+### Features
+
+* **Update astra source connector to use new astrapy client**
+
 ## 1.0.1
 
 ### Features

pyproject.toml

@@ -172,7 +172,6 @@ ci = [
     "lancedb<=0.15.0",
     # TODO: versions higher than this are missing the macos wheel
     "pykx==2.5.3",
-    "astrapy<2.0.0"
 ]
 
 [project.scripts]

requirements/connectors/astradb.txt

@@ -1 +1 @@
-astrapy
+astrapy>2.0.0

test/integration/connectors/databricks/test_volumes_native.py

@@ -260,10 +260,7 @@ async def test_volumes_native_destination(upload_file: Path):
             ),
         )
         uploader.precheck()
-        if uploader.is_async():
-            await uploader.run_async(path=upload_file, file_data=file_data)
-        else:
-            uploader.run(path=upload_file, file_data=file_data)
+        uploader.run(path=upload_file, file_data=file_data)
 
         validate_upload(
             client=workspace_client,

test/integration/connectors/test_astradb.py

@@ -9,6 +9,7 @@
 from _pytest.fixtures import TopRequest
 from astrapy import Collection
 from astrapy import DataAPIClient as AstraDBClient
+from astrapy.info import CollectionDefinition
 
 from test.integration.connectors.utils.constants import DESTINATION_TAG, SOURCE_TAG, VECTOR_DB_TAG
 from test.integration.connectors.utils.validation.destination import (
@@ -127,11 +128,16 @@ def collection(upload_file: Path) -> Collection:
         api_endpoint=env_data.api_endpoint,
         token=env_data.token,
     )
-    collection = astra_db.create_collection(collection_name, dimension=embedding_dimension)
+    collection = astra_db.create_collection(
+        collection_name,
+        definition=CollectionDefinition.builder()
+        .set_vector_dimension(dimension=embedding_dimension)
+        .build(),
+    )
     try:
         yield collection
     finally:
-        astra_db.drop_collection(collection)
+        astra_db.drop_collection(collection.name)
 
 
 @pytest.mark.asyncio
@@ -198,7 +204,7 @@ async def test_astra_search_destination(
         output_filename=upload_file.name,
     )
     uploader.precheck()
-    uploader.run(path=staged_filepath, file_data=file_data)
+    await uploader.run_async(path=staged_filepath, file_data=file_data)
 
     # Run validation
     with staged_filepath.open() as f:
@@ -211,7 +217,7 @@ async def test_astra_search_destination(
     )
 
     # Rerun and make sure the same documents get updated
-    uploader.run(path=staged_filepath, file_data=file_data)
+    await uploader.run_async(path=staged_filepath, file_data=file_data)
     current_count = collection.count_documents(filter={}, upper_bound=expected_count * 2)
     assert current_count == expected_count, (
         f"Expected count ({expected_count}) doesn't match how "

unstructured_ingest/__version__.py

@@ -1 +1 @@
-__version__ = "1.0.1"  # pragma: no cover
+__version__ = "1.0.2"  # pragma: no cover

unstructured_ingest/processes/connectors/astradb.py

@@ -1,5 +1,7 @@
+import asyncio
 import csv
 import hashlib
+import os
 import re
 from dataclasses import dataclass, field
 from pathlib import Path
@@ -8,7 +10,6 @@
 
 from pydantic import BaseModel, Field, Secret
 
-from unstructured_ingest import __name__ as integration_name
 from unstructured_ingest.__version__ import __version__ as integration_version
 from unstructured_ingest.data_types.file_data import (
     BatchFileData,
@@ -83,10 +84,8 @@ def get_client(self) -> "AstraDBClient":
 
         # Create a client object to interact with the Astra DB
         # caller_name/version for Astra DB tracking
-        return AstraDBClient(
-            caller_name=integration_name,
-            caller_version=integration_version,
-        )
+        user_agent = os.getenv("UNSTRUCTURED_USER_AGENT", "unstructuredio_oss")
+        return AstraDBClient(callers=[(user_agent, integration_version)])
 
 
 def get_astra_db(
@@ -141,7 +140,7 @@ async def get_async_astra_collection(
     )
 
     # Get async collection from AsyncDatabase
-    async_astra_db_collection = await async_astra_db.get_collection(name=collection_name)
+    async_astra_db_collection = async_astra_db.get_collection(name=collection_name)
     return async_astra_db_collection
 
 
@@ -360,13 +359,22 @@ class AstraDBUploader(Uploader):
     upload_config: AstraDBUploaderConfig
     connector_type: str = CONNECTOR_TYPE
 
+    def is_async(self) -> bool:
+        return True
+
     def init(self, **kwargs: Any) -> None:
         self.create_destination(**kwargs)
 
+    @requires_dependencies(["astrapy"], extras="astradb")
     def precheck(self) -> None:
         try:
             if self.upload_config.collection_name:
-                self.get_collection(collection_name=self.upload_config.collection_name).options()
+                collection = get_astra_collection(
+                    connection_config=self.connection_config,
+                    collection_name=self.upload_config.collection_name,
+                    keyspace=self.upload_config.keyspace,
+                )
+                collection.options()
             else:
                 # check for db connection only if collection name is not provided
                 get_astra_db(
@@ -377,17 +385,7 @@ def precheck(self) -> None:
             logger.error(f"Failed to validate connection {e}", exc_info=True)
             raise DestinationConnectionError(f"failed to validate connection: {e}")
 
-    @requires_dependencies(["astrapy"], extras="astradb")
-    def get_collection(self, collection_name: Optional[str] = None) -> "AstraDBCollection":
-        return get_astra_collection(
-            connection_config=self.connection_config,
-            collection_name=collection_name or self.upload_config.collection_name,
-            keyspace=self.upload_config.keyspace,
-        )
-
     def _collection_exists(self, collection_name: str):
-        from astrapy.exceptions import CollectionNotFoundException
-
         collection = get_astra_collection(
             connection_config=self.connection_config,
             collection_name=collection_name,
@@ -397,8 +395,10 @@ def _collection_exists(self, collection_name: str):
         try:
             collection.options()
             return True
-        except CollectionNotFoundException:
-            return False
+        except RuntimeError as e:
+            if "not found" in str(e):
+                return False
+            raise DestinationConnectionError(f"failed to check if astra collection exists : {e}")
         except Exception as e:
             logger.error(f"failed to check if astra collection exists : {e}")
             raise DestinationConnectionError(f"failed to check if astra collection exists : {e}")
@@ -422,51 +422,65 @@ def create_destination(
         self.upload_config.collection_name = collection_name
 
         if not self._collection_exists(collection_name):
+            from astrapy.info import CollectionDefinition
+
             astra_db = get_astra_db(
                 connection_config=self.connection_config, keyspace=self.upload_config.keyspace
             )
             logger.info(
                 f"creating default astra collection '{collection_name}' with dimension "
                 f"{vector_length} and metric {similarity_metric}"
             )
-            astra_db.create_collection(
-                collection_name,
-                dimension=vector_length,
-                metric=similarity_metric,
+            definition = (
+                CollectionDefinition.builder()
+                .set_vector_dimension(dimension=vector_length)
+                .set_vector_metric(similarity_metric)
+                .build()
             )
+            (astra_db.create_collection(collection_name, definition=definition),)
             return True
         logger.debug(f"collection with name '{collection_name}' already exists, skipping creation")
         return False
 
-    def delete_by_record_id(self, collection: "AstraDBCollection", file_data: FileData):
+    async def delete_by_record_id(self, collection: "AstraDBAsyncCollection", file_data: FileData):
         logger.debug(
             f"deleting records from collection {collection.name} "
             f"with {self.upload_config.record_id_key} "
             f"set to {file_data.identifier}"
         )
         delete_filter = {self.upload_config.record_id_key: {"$eq": file_data.identifier}}
-        delete_resp = collection.delete_many(filter=delete_filter)
+        delete_resp = await collection.delete_many(filter=delete_filter)
         logger.debug(
             f"deleted {delete_resp.deleted_count} records from collection {collection.name}"
         )
 
-    def run_data(self, data: list[dict], file_data: FileData, **kwargs: Any) -> None:
+    async def run_data(self, data: list[dict], file_data: FileData, **kwargs: Any) -> None:
         logger.info(
             f"writing {len(data)} objects to destination "
             f"collection {self.upload_config.collection_name}"
         )
 
         astra_db_batch_size = self.upload_config.batch_size
-        collection = self.get_collection()
-
-        self.delete_by_record_id(collection=collection, file_data=file_data)
+        async_astra_collection = await get_async_astra_collection(
+            connection_config=self.connection_config,
+            collection_name=self.upload_config.collection_name,
+            keyspace=self.upload_config.keyspace,
+        )
 
-        for chunk in batch_generator(data, astra_db_batch_size):
-            collection.insert_many(chunk)
+        await self.delete_by_record_id(collection=async_astra_collection, file_data=file_data)
+        await asyncio.gather(
+            *[
+                async_astra_collection.insert_many(chunk)
+                for chunk in batch_generator(data, astra_db_batch_size)
+            ]
+        )
 
-    def run(self, path: Path, file_data: FileData, **kwargs: Any) -> None:
+    async def run_async(self, path: Path, file_data: FileData, **kwargs: Any) -> None:
         data = get_json_data(path=path)
-        self.run_data(data=data, file_data=file_data, **kwargs)
+        await self.run_data(data=data, file_data=file_data)
+
+    def run(self, **kwargs: Any) -> Any:
+        raise NotImplementedError("Use astradb run_async instead")
 
 
 astra_db_source_entry = SourceRegistryEntry(