feat: opinionated writes for AstraDB destination (#391)

* make collection optional; wip create dest

* name formatter

* code

* test and clean

* bump version

* better none handling; fixes

* move import into functions

* move collection verification to callers; rewrite exists method

* tidy

Author: shreyanid

CHANGELOG.md

@@ -1,4 +1,8 @@
-## 0.5.9-dev0
+## 0.5.9-dev1
+
+### Features
+
+* Add auto create collection support for AstraDB destination
 
 ### Fixes
 

test/integration/connectors/test_astradb.py

@@ -56,12 +56,7 @@ def test_precheck_succeeds_indexer(connection_config: AstraDBConnectionConfig):
         connection_config=connection_config,
         index_config=AstraDBIndexerConfig(collection_name=EXISTENT_COLLECTION_NAME),
     )
-    uploader = AstraDBUploader(
-        connection_config=connection_config,
-        upload_config=AstraDBUploaderConfig(collection_name=EXISTENT_COLLECTION_NAME),
-    )
     indexer.precheck()
-    uploader.precheck()
 
 
 @pytest.mark.tags(CONNECTOR_TYPE, DESTINATION_TAG, VECTOR_DB_TAG)
@@ -73,6 +68,12 @@ def test_precheck_succeeds_uploader(connection_config: AstraDBConnectionConfig):
     )
     uploader.precheck()
 
+    uploader2 = AstraDBUploader(
+        connection_config=connection_config,
+        upload_config=AstraDBUploaderConfig(),
+    )
+    uploader2.precheck()
+
 
 @pytest.mark.tags(CONNECTOR_TYPE, SOURCE_TAG, VECTOR_DB_TAG)
 @requires_env("ASTRA_DB_APPLICATION_TOKEN", "ASTRA_DB_API_ENDPOINT")
@@ -216,6 +217,32 @@ async def test_astra_search_destination(
     )
 
 
+@pytest.mark.tags(CONNECTOR_TYPE, DESTINATION_TAG, VECTOR_DB_TAG)
+@requires_env("ASTRA_DB_API_ENDPOINT", "ASTRA_DB_APPLICATION_TOKEN")
+def test_astra_create_destination():
+    env_data = get_env_data()
+    connection_config = AstraDBConnectionConfig(
+        access_config=AstraDBAccessConfig(api_endpoint=env_data.api_endpoint, token=env_data.token),
+    )
+    uploader = AstraDBUploader(
+        connection_config=connection_config,
+        upload_config=AstraDBUploaderConfig(),
+    )
+    collection_name = "system_created-123"
+    formatted_collection_name = "system_created_123"
+    created = uploader.create_destination(destination_name=collection_name, vector_length=3072)
+    assert created
+    assert uploader.upload_config.collection_name == formatted_collection_name
+
+    created = uploader.create_destination(destination_name=collection_name, vector_length=3072)
+    assert not created
+
+    # cleanup
+    client = AstraDBClient()
+    db = client.get_database(api_endpoint=env_data.api_endpoint, token=env_data.token)
+    db.drop_collection(formatted_collection_name)
+
+
 @pytest.mark.tags(CONNECTOR_TYPE, DESTINATION_TAG, VECTOR_DB_TAG)
 @pytest.mark.parametrize("upload_file_str", ["upload_file_ndjson", "upload_file"])
 def test_astra_stager(

unstructured_ingest/__version__.py

@@ -1 +1 @@
-__version__ = "0.5.9-dev0"  # pragma: no cover
+__version__ = "0.5.9-dev1"  # pragma: no cover

unstructured_ingest/v2/processes/connectors/astradb.py

@@ -1,5 +1,6 @@
 import csv
 import hashlib
+import re
 from dataclasses import dataclass, field
 from pathlib import Path
 from time import time
@@ -48,6 +49,7 @@
     from astrapy import AsyncCollection as AstraDBAsyncCollection
     from astrapy import Collection as AstraDBCollection
     from astrapy import DataAPIClient as AstraDBClient
+    from astrapy import Database as AstraDB
 
 
 CONNECTOR_TYPE = "astradb"
@@ -85,11 +87,10 @@ def get_client(self) -> "AstraDBClient":
         )
 
 
-def get_astra_collection(
+def get_astra_db(
     connection_config: AstraDBConnectionConfig,
-    collection_name: str,
     keyspace: str,
-) -> "AstraDBCollection":
+) -> "AstraDB":
     # Build the Astra DB object.
     access_configs = connection_config.access_config.get_secret_value()
 
@@ -103,9 +104,20 @@ def get_astra_collection(
         token=access_configs.token,
         keyspace=keyspace,
     )
+    return astra_db
+
 
-    # Connect to the collection
+def get_astra_collection(
+    connection_config: AstraDBConnectionConfig,
+    collection_name: str,
+    keyspace: str,
+) -> "AstraDBCollection":
+
+    astra_db = get_astra_db(connection_config=connection_config, keyspace=keyspace)
+
+    # astradb will return a collection object in all cases (even if it doesn't exist)
     astra_db_collection = astra_db.get_collection(name=collection_name)
+
     return astra_db_collection
 
 
@@ -151,10 +163,11 @@ class AstraDBDownloaderConfig(DownloaderConfig):
 
 
 class AstraDBUploaderConfig(UploaderConfig):
-    collection_name: str = Field(
+    collection_name: Optional[str] = Field(
         description="The name of the Astra DB collection. "
         "Note that the collection name must only include letters, "
-        "numbers, and underscores."
+        "numbers, and underscores.",
+        default=None,
     )
     keyspace: Optional[str] = Field(default=None, description="The Astra DB connection keyspace.")
     requested_indexing_policy: Optional[dict[str, Any]] = Field(
@@ -337,25 +350,84 @@ class AstraDBUploader(Uploader):
     upload_config: AstraDBUploaderConfig
     connector_type: str = CONNECTOR_TYPE
 
+    def init(self, **kwargs: Any) -> None:
+        self.create_destination(**kwargs)
+
     def precheck(self) -> None:
         try:
-            get_astra_collection(
-                connection_config=self.connection_config,
-                collection_name=self.upload_config.collection_name,
-                keyspace=self.upload_config.keyspace,
-            ).options()
+            if self.upload_config.collection_name:
+                self.get_collection(collection_name=self.upload_config.collection_name).options()
+            else:
+                # check for db connection only if collection name is not provided
+                get_astra_db(
+                    connection_config=self.connection_config,
+                    keyspace=self.upload_config.keyspace,
+                )
         except Exception as e:
             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) -> "AstraDBCollection":
+    def get_collection(self, collection_name: Optional[str] = None) -> "AstraDBCollection":
         return get_astra_collection(
             connection_config=self.connection_config,
-            collection_name=self.upload_config.collection_name,
+            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,
+            keyspace=self.upload_config.keyspace,
+        )
+
+        try:
+            collection.options()
+            return True
+        except CollectionNotFoundException:
+            return False
+        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}")
+
+    def format_destination_name(self, destination_name: str) -> str:
+        # AstraDB collection naming requirements:
+        # must be below 50 characters
+        # must be lowercase alphanumeric and underscores only
+        formatted = re.sub(r"[^a-z0-9]", "_", destination_name.lower())
+        return formatted
+
+    def create_destination(
+        self,
+        vector_length: int,
+        destination_name: str = "unstructuredautocreated",
+        similarity_metric: Optional[str] = "cosine",
+        **kwargs: Any,
+    ) -> bool:
+        destination_name = self.format_destination_name(destination_name)
+        collection_name = self.upload_config.collection_name or destination_name
+        self.upload_config.collection_name = collection_name
+
+        if not self._collection_exists(collection_name):
+            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,
+            )
+            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):
         logger.debug(
             f"deleting records from collection {collection.name} "