feat: add precheck to all connectors (#4)

* Support check connection via precheck method

* Add back in precheck for all connectors

* Add entry to changelog

* Fix sql precheck

Author: rbiseck3

.github/workflows/release.yml

@@ -4,8 +4,6 @@ on:
   release:
     types:
       - published
-      - created
-      - released
 
 env:
   PYTHON_VERSION: "3.10"

CHANGELOG.md

@@ -1,8 +1,9 @@
-## 0.0.2-dev0
+## 0.0.2-dev1
 
 ### Enhancements
 
 * **Use uuid for s3 identifiers** Update unique id to use uuid derived from file path rather than the filepath itself.
+* **V2 connectors precheck support** All steps in the v2 pipeline support an optional precheck call, which encompasses the previous check connection functionality. 
 
 ## 0.0.1
 

unstructured_ingest/__version__.py

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

unstructured_ingest/v2/examples/example_sql.py

@@ -13,8 +13,8 @@
 )
 from unstructured_ingest.v2.processes.connectors.sql import (
     DatabaseType,
-    SimpleSqlConfig,
     SQLAccessConfig,
+    SQLConnectionConfig,
     SQLUploaderConfig,
     SQLUploadStagerConfig,
 )
@@ -67,7 +67,7 @@
 
     # sqlite test first
     Pipeline.from_configs(
-        destination_connection_config=SimpleSqlConfig(
+        destination_connection_config=SQLConnectionConfig(
             db_type=DatabaseType.SQLITE,
             database=SQLITE_DB,
             access_config=SQLAccessConfig(),
@@ -77,7 +77,7 @@
 
     # now, pg with pgvector
     Pipeline.from_configs(
-        destination_connection_config=SimpleSqlConfig(
+        destination_connection_config=SQLConnectionConfig(
             db_type=DatabaseType.POSTGRESQL,
             database="elements",
             host="localhost",

unstructured_ingest/v2/interfaces/process.py

@@ -8,6 +8,9 @@ class BaseProcess(ABC):
     def is_async(self) -> bool:
         return False
 
+    def precheck(self) -> None:
+        pass
+
     @abstractmethod
     def run(self, **kwargs: Any) -> Any:
         pass

unstructured_ingest/v2/pipeline/pipeline.py

@@ -109,6 +109,7 @@ def log_statuses(self):
     def run(self):
         try:
             start_time = time()
+            self._run_prechecks()
             self._run()
             logger.info(f"Finished ingest process in {time() - start_time}s")
         finally:
@@ -130,6 +131,27 @@ def clean_results(self, results: Optional[list[Union[Any, list[Any]]]]) -> Optio
         final = [f for f in flat if f]
         return final or None
 
+    def _run_prechecks(self):
+        steps = [self.indexer_step, self.downloader_step, self.partitioner_step, self.uploader_step]
+        if self.chunker_step:
+            steps.append(self.chunker_step)
+        if self.embedder_step:
+            steps.append(self.embedder_step)
+        if self.uncompress_step:
+            steps.append(self.uncompress_step)
+        if self.stager_step:
+            steps.append(self.stager_step)
+        failures = {}
+        for step in steps:
+            try:
+                step.process.precheck()
+            except Exception as e:
+                failures[step.process.__class__.__name__] = f"[{type(e).__name__}] {e}"
+        if failures:
+            for k, v in failures.items():
+                logger.error(f"Step precheck failure: {k}: {v}")
+            raise PipelineError("Precheck failed")
+
     def _run(self):
         logger.info(
             f"Running local pipline: {self} with configs: "

unstructured_ingest/v2/processes/connectors/astra.py

@@ -7,6 +7,7 @@
 from unstructured.__version__ import __version__ as integration_version
 
 from unstructured_ingest.enhanced_dataclass import enhanced_field
+from unstructured_ingest.error import DestinationConnectionError
 from unstructured_ingest.utils.data_prep import batch_generator
 from unstructured_ingest.utils.dep_check import requires_dependencies
 from unstructured_ingest.v2.interfaces import (
@@ -94,6 +95,13 @@ class AstraUploader(Uploader):
     upload_config: AstraUploaderConfig
     connector_type: str = CONNECTOR_TYPE
 
+    def precheck(self) -> None:
+        try:
+            self.get_collection()
+        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="astra")
     def get_collection(self) -> "AstraDBCollection":
         from astrapy.db import AstraDB

unstructured_ingest/v2/processes/connectors/azure_cognitive_search.py

@@ -175,6 +175,14 @@ def write_dict(self, *args, elements_dict: t.List[t.Dict[str, t.Any]], **kwargs)
                 ),
             )
 
+    def precheck(self) -> None:
+        try:
+            client = self.connection_config.generate_client()
+            client.get_document_count()
+        except Exception as e:
+            logger.error(f"failed to validate connection: {e}", exc_info=True)
+            raise DestinationConnectionError(f"failed to validate connection: {e}")
+
     def write_dict_wrapper(self, elements_dict):
         return self.write_dict(elements_dict=elements_dict)
 

unstructured_ingest/v2/processes/connectors/chroma.py

@@ -111,10 +111,13 @@ class ChromaUploader(Uploader):
     connector_type: str = CONNECTOR_TYPE
     upload_config: ChromaUploaderConfig
     connection_config: ChromaConnectionConfig
-    client: Optional["Client"] = field(init=False)
 
-    def __post_init__(self):
-        self.client = self.create_client()
+    def precheck(self) -> None:
+        try:
+            self.create_client()
+        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(["chromadb"], extras="chroma")
     def create_client(self) -> "Client":
@@ -187,10 +190,9 @@ def run(self, contents: list[UploadContent], **kwargs: Any) -> None:
             f"collection {self.connection_config.collection_name} "
             f"at {self.connection_config.host}",
         )
+        client = self.create_client()
 
-        collection = self.client.get_or_create_collection(
-            name=self.connection_config.collection_name
-        )
+        collection = client.get_or_create_collection(name=self.connection_config.collection_name)
         for chunk in batch_generator(elements_dict, self.upload_config.batch_size):
             self.upsert_batch(collection, self.prepare_chroma_list(chunk))
 

unstructured_ingest/v2/processes/connectors/databricks_volumes.py

@@ -3,6 +3,7 @@
 from typing import TYPE_CHECKING, Any, Optional
 
 from unstructured_ingest.enhanced_dataclass import enhanced_field
+from unstructured_ingest.error import DestinationConnectionError
 from unstructured_ingest.utils.dep_check import requires_dependencies
 from unstructured_ingest.v2.interfaces import (
     AccessConfig,
@@ -11,6 +12,7 @@
     Uploader,
     UploaderConfig,
 )
+from unstructured_ingest.v2.logger import logger
 from unstructured_ingest.v2.processes.connector_registry import DestinationRegistryEntry
 
 if TYPE_CHECKING:
@@ -78,6 +80,13 @@ def __post_init__(self) -> "WorkspaceClient":
             host=self.connection_config.host, **self.connection_config.access_config.to_dict()
         )
 
+    def precheck(self) -> None:
+        try:
+            assert self.client.current_user.me().active
+        except Exception as e:
+            logger.error(f"failed to validate connection: {e}", exc_info=True)
+            raise DestinationConnectionError(f"failed to validate connection: {e}")
+
     def run(self, contents: list[UploadContent], **kwargs: Any) -> None:
         for content in contents:
             with open(content.path, "rb") as elements_file: