unstructured-ingest 0.0.2.dev0__py3-none-any.whl → 0.0.4__py3-none-any.whl

unstructured_ingest/v2/processes/connectors/couchbase.py

@@ -0,0 +1,151 @@
+import json
+from dataclasses import dataclass, field
+from datetime import timedelta
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+from pydantic import Field, Secret
+
+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 (
+    AccessConfig,
+    ConnectionConfig,
+    UploadContent,
+    Uploader,
+    UploaderConfig,
+    UploadStager,
+    UploadStagerConfig,
+)
+from unstructured_ingest.v2.logger import logger
+from unstructured_ingest.v2.processes.connector_registry import (
+    DestinationRegistryEntry,
+)
+
+if TYPE_CHECKING:
+    from couchbase.cluster import Cluster
+
+CONNECTOR_TYPE = "couchbase"
+SERVER_API_VERSION = "1"
+
+
+class CouchbaseAccessConfig(AccessConfig):
+    password: str = Field(description="The password for the Couchbase server")
+
+
+class CouchbaseConnectionConfig(ConnectionConfig):
+    username: str = Field(description="The username for the Couchbase server")
+    bucket: str = Field(description="The bucket to connect to on the Couchbase server")
+    connection_string: str = Field(
+        default="couchbase://localhost", description="The connection string of the Couchbase server"
+    )
+    scope: str = Field(
+        default="_default", description="The scope to connect to on the Couchbase server"
+    )
+    collection: str = Field(
+        default="_default", description="The collection to connect to on the Couchbase server"
+    )
+    connector_type: str = Field(default=CONNECTOR_TYPE, init=False)
+    access_config: Secret[CouchbaseAccessConfig]
+
+
+class CouchbaseUploadStagerConfig(UploadStagerConfig):
+    pass
+
+
+@dataclass
+class CouchbaseUploadStager(UploadStager):
+    upload_stager_config: CouchbaseUploadStagerConfig = field(
+        default_factory=lambda: CouchbaseUploadStagerConfig()
+    )
+
+    def run(
+        self,
+        elements_filepath: Path,
+        output_dir: Path,
+        output_filename: str,
+        **kwargs: Any,
+    ) -> Path:
+        with open(elements_filepath) as elements_file:
+            elements_contents = json.load(elements_file)
+
+        output_elements = []
+        for element in elements_contents:
+            new_doc = {
+                element["element_id"]: {
+                    "embedding": element.get("embeddings", None),
+                    "text": element.get("text", None),
+                    "metadata": element.get("metadata", None),
+                    "type": element.get("type", None),
+                }
+            }
+            output_elements.append(new_doc)
+
+        output_path = Path(output_dir) / Path(f"{output_filename}.json")
+        with open(output_path, "w") as output_file:
+            json.dump(output_elements, output_file)
+        return output_path
+
+
+class CouchbaseUploaderConfig(UploaderConfig):
+    batch_size: int = Field(default=50, description="Number of documents to upload per batch")
+
+
+@dataclass
+class CouchbaseUploader(Uploader):
+    connection_config: CouchbaseConnectionConfig
+    upload_config: CouchbaseUploaderConfig
+    connector_type: str = CONNECTOR_TYPE
+
+    @requires_dependencies(["couchbase"], extras="couchbase")
+    def connect_to_couchbase(self) -> "Cluster":
+        from couchbase.auth import PasswordAuthenticator
+        from couchbase.cluster import Cluster
+        from couchbase.options import ClusterOptions
+
+        connection_string = self.connection_config.connection_string
+        username = self.connection_config.username
+        password = self.connection_config.access_config.get_secret_value().password
+
+        auth = PasswordAuthenticator(username, password)
+        options = ClusterOptions(auth)
+        options.apply_profile("wan_development")
+        cluster = Cluster(connection_string, options)
+        cluster.wait_until_ready(timedelta(seconds=5))
+        return cluster
+
+    def precheck(self) -> None:
+        try:
+            self.connect_to_couchbase()
+        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:
+        elements = []
+        for content in contents:
+            with open(content.path) as elements_file:
+                elements.extend(json.load(elements_file))
+
+        logger.info(
+            f"writing {len(elements)} objects to destination "
+            f"bucket, {self.connection_config.bucket} "
+            f"at {self.connection_config.connection_string}",
+        )
+        cluster = self.connect_to_couchbase()
+        bucket = cluster.bucket(self.connection_config.bucket)
+        scope = bucket.scope(self.connection_config.scope)
+        collection = scope.collection(self.connection_config.collection)
+
+        for chunk in batch_generator(elements, self.upload_config.batch_size):
+            collection.upsert_multi({doc_id: doc for doc in chunk for doc_id, doc in doc.items()})
+
+
+couchbase_destination_entry = DestinationRegistryEntry(
+    connection_config=CouchbaseConnectionConfig,
+    uploader=CouchbaseUploader,
+    uploader_config=CouchbaseUploaderConfig,
+    upload_stager=CouchbaseUploadStager,
+    upload_stager_config=CouchbaseUploadStagerConfig,
+)

unstructured_ingest/v2/processes/connectors/databricks_volumes.py

@@ -1,8 +1,10 @@
 import os
-from dataclasses import dataclass, field
+from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, Optional
 
-from unstructured_ingest.enhanced_dataclass import enhanced_field
+from pydantic import Field, Secret
+
+from unstructured_ingest.error import DestinationConnectionError
 from unstructured_ingest.utils.dep_check import requires_dependencies
 from unstructured_ingest.v2.interfaces import (
     AccessConfig,
@@ -11,6 +13,7 @@ from unstructured_ingest.v2.interfaces import (
     Uploader,
     UploaderConfig,
 )
+from unstructured_ingest.v2.logger import logger
 from unstructured_ingest.v2.processes.connector_registry import DestinationRegistryEntry
 
 if TYPE_CHECKING:
@@ -19,45 +22,99 @@ if TYPE_CHECKING:
 CONNECTOR_TYPE = "databricks_volumes"
 
 
-@dataclass
 class DatabricksVolumesAccessConfig(AccessConfig):
-    account_id: Optional[str] = None
-    username: Optional[str] = None
-    password: Optional[str] = None
-    client_id: Optional[str] = None
-    client_secret: Optional[str] = None
-    token: Optional[str] = None
+    account_id: Optional[str] = Field(
+        default=None,
+        description="The Databricks account ID for the Databricks "
+        "accounts endpoint. Only has effect when Host is "
+        "either https://accounts.cloud.databricks.com/ (AWS), "
+        "https://accounts.azuredatabricks.net/ (Azure), "
+        "or https://accounts.gcp.databricks.com/ (GCP).",
+    )
+    username: Optional[str] = Field(
+        default=None,
+        description="The Databricks username part of basic authentication. "
+        "Only possible when Host is *.cloud.databricks.com (AWS).",
+    )
+    password: Optional[str] = Field(
+        default=None,
+        description="The Databricks password part of basic authentication. "
+        "Only possible when Host is *.cloud.databricks.com (AWS).",
+    )
+    client_id: Optional[str] = Field(default=None)
+    client_secret: Optional[str] = Field(default=None)
+    token: Optional[str] = Field(
+        default=None,
+        description="The Databricks personal access token (PAT) (AWS, Azure, and GCP) or "
+        "Azure Active Directory (Azure AD) token (Azure).",
+    )
     profile: Optional[str] = None
-    azure_workspace_resource_id: Optional[str] = None
-    azure_client_secret: Optional[str] = None
-    azure_client_id: Optional[str] = None
-    azure_tenant_id: Optional[str] = None
-    azure_environment: Optional[str] = None
-    auth_type: Optional[str] = None
+    azure_workspace_resource_id: Optional[str] = Field(
+        default=None,
+        description="The Azure Resource Manager ID for the Azure Databricks workspace, "
+        "which is exchanged for a Databricks host URL.",
+    )
+    azure_client_secret: Optional[str] = Field(
+        default=None, description="The Azure AD service principal’s client secret."
+    )
+    azure_client_id: Optional[str] = Field(
+        default=None, description="The Azure AD service principal’s application ID."
+    )
+    azure_tenant_id: Optional[str] = Field(
+        default=None, description="The Azure AD service principal’s tenant ID."
+    )
+    azure_environment: Optional[str] = Field(
+        default=None,
+        description="The Azure environment type for a " "specific set of API endpoints",
+        examples=["Public", "UsGov", "China", "Germany"],
+    )
+    auth_type: Optional[str] = Field(
+        default=None,
+        description="When multiple auth attributes are available in the "
+        "environment, use the auth type specified by this "
+        "argument. This argument also holds the currently "
+        "selected auth.",
+    )
     cluster_id: Optional[str] = None
     google_credentials: Optional[str] = None
     google_service_account: Optional[str] = None
 
 
-@dataclass
+SecretDatabricksVolumesAccessConfig = Secret[DatabricksVolumesAccessConfig]
+
+
 class DatabricksVolumesConnectionConfig(ConnectionConfig):
-    access_config: DatabricksVolumesAccessConfig = enhanced_field(
-        default_factory=DatabricksVolumesAccessConfig, sensitive=True
+    access_config: SecretDatabricksVolumesAccessConfig = Field(
+        default_factory=lambda: SecretDatabricksVolumesAccessConfig(
+            secret_value=DatabricksVolumesAccessConfig()
+        )
+    )
+    host: Optional[str] = Field(
+        default=None,
+        description="The Databricks host URL for either the "
+        "Databricks workspace endpoint or the "
+        "Databricks accounts endpoint.",
     )
-    host: Optional[str] = None
 
 
-@dataclass
 class DatabricksVolumesUploaderConfig(UploaderConfig):
-    volume: str
-    catalog: str
-    volume_path: Optional[str] = None
-    overwrite: bool = False
-    schema: str = "default"
+    volume: str = Field(description="Name of volume in the Unity Catalog")
+    catalog: str = Field(description="Name of the catalog in the Databricks Unity Catalog service")
+    volume_path: Optional[str] = Field(
+        default=None, description="Optional path within the volume to write to"
+    )
+    overwrite: bool = Field(
+        default=False, description="If true, an existing file will be overwritten."
+    )
+    databricks_schema: str = Field(
+        default="default",
+        alias="schema",
+        description="Schema associated with the volume to write to in the Unity Catalog service",
+    )
 
     @property
     def path(self) -> str:
-        path = f"/Volumes/{self.catalog}/{self.schema}/{self.volume}"
+        path = f"/Volumes/{self.catalog}/{self.databricks_schema}/{self.volume}"
         if self.volume_path:
             path = f"{path}/{self.volume_path}"
         return path
@@ -68,21 +125,28 @@ class DatabricksVolumesUploader(Uploader):
     connector_type: str = CONNECTOR_TYPE
     upload_config: DatabricksVolumesUploaderConfig
     connection_config: DatabricksVolumesConnectionConfig
-    client: Optional["WorkspaceClient"] = field(init=False, default=None)
 
     @requires_dependencies(dependencies=["databricks.sdk"], extras="databricks-volumes")
-    def __post_init__(self) -> "WorkspaceClient":
+    def get_client(self) -> "WorkspaceClient":
         from databricks.sdk import WorkspaceClient
 
-        self.client = WorkspaceClient(
-            host=self.connection_config.host, **self.connection_config.access_config.to_dict()
+        return WorkspaceClient(
+            host=self.connection_config.host,
+            **self.connection_config.access_config.get_secret_value().dict(),
         )
 
+    def precheck(self) -> None:
+        try:
+            assert self.get_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:
                 output_path = os.path.join(self.upload_config.path, content.path.name)
-                self.client.files.upload(
+                self.get_client().files.upload(
                     file_path=output_path,
                     contents=elements_file,
                     overwrite=self.upload_config.overwrite,

unstructured_ingest/v2/processes/connectors/elasticsearch.py

@@ -5,12 +5,15 @@ import uuid
 from dataclasses import dataclass, field
 from pathlib import Path
 from time import time
-from typing import TYPE_CHECKING, Any, Generator, Optional
+from typing import TYPE_CHECKING, Any, Generator, Optional, Union
 
-from unstructured.documents.elements import DataSourceMetadata
+from pydantic import BaseModel, Field, Secret, SecretStr
 
-from unstructured_ingest.enhanced_dataclass import EnhancedDataClassJsonMixin, enhanced_field
-from unstructured_ingest.error import SourceConnectionError, SourceConnectionNetworkError
+from unstructured_ingest.error import (
+    DestinationConnectionError,
+    SourceConnectionError,
+    SourceConnectionNetworkError,
+)
 from unstructured_ingest.utils.data_prep import flatten_dict, generator_batching_wbytes
 from unstructured_ingest.utils.dep_check import requires_dependencies
 from unstructured_ingest.v2.interfaces import (
@@ -20,6 +23,7 @@ from unstructured_ingest.v2.interfaces import (
     DownloaderConfig,
     DownloadResponse,
     FileData,
+    FileDataSourceMetadata,
     Indexer,
     IndexerConfig,
     UploadContent,
@@ -41,57 +45,74 @@ if TYPE_CHECKING:
 CONNECTOR_TYPE = "elasticsearch"
 
 
-@dataclass
 class ElasticsearchAccessConfig(AccessConfig):
-    password: Optional[str] = None
-    api_key: Optional[str] = enhanced_field(default=None, overload_name="es_api_key")
-    bearer_auth: Optional[str] = None
-    ssl_assert_fingerprint: Optional[str] = None
-
-
-@dataclass
-class ElasticsearchClientInput(EnhancedDataClassJsonMixin):
+    password: Optional[str] = Field(
+        default=None, description="password when using basic auth or connecting to a cloud instance"
+    )
+    es_api_key: Optional[str] = Field(default=None, description="api key used for authentication")
+    bearer_auth: Optional[str] = Field(
+        default=None, description="bearer token used for HTTP bearer authentication"
+    )
+    ssl_assert_fingerprint: Optional[str] = Field(
+        default=None, description="SHA256 fingerprint value"
+    )
+
+
+class ElasticsearchClientInput(BaseModel):
     hosts: Optional[list[str]] = None
     cloud_id: Optional[str] = None
-    ca_certs: Optional[str] = None
-    basic_auth: Optional[tuple[str, str]] = enhanced_field(sensitive=True, default=None)
-    api_key: Optional[str] = enhanced_field(sensitive=True, default=None)
+    ca_certs: Optional[Path] = None
+    basic_auth: Optional[Secret[tuple[str, str]]] = None
+    api_key: Optional[Union[Secret[tuple[str, str]], SecretStr]] = None
 
 
-@dataclass
 class ElasticsearchConnectionConfig(ConnectionConfig):
-    hosts: Optional[list[str]] = None
-    username: Optional[str] = None
-    cloud_id: Optional[str] = None
-    api_key_id: Optional[str] = None
-    ca_certs: Optional[str] = None
-    access_config: ElasticsearchAccessConfig = enhanced_field(sensitive=True)
+    hosts: Optional[list[str]] = Field(
+        default=None,
+        description="list of the Elasticsearch hosts to connect to",
+        examples=["http://localhost:9200"],
+    )
+    username: Optional[str] = Field(default=None, description="username when using basic auth")
+    cloud_id: Optional[str] = Field(default=None, description="id used to connect to Elastic Cloud")
+    api_key_id: Optional[str] = Field(
+        default=None,
+        description="id associated with api key used for authentication: "
+        "https://www.elastic.co/guide/en/elasticsearch/reference/current/security-api-create-api-key.html",  # noqa: E501
+    )
+    ca_certs: Optional[Path] = None
+    access_config: Secret[ElasticsearchAccessConfig]
 
     def get_client_kwargs(self) -> dict:
         # Update auth related fields to conform to what the SDK expects based on the
         # supported methods:
         # https://www.elastic.co/guide/en/elasticsearch/client/python-api/current/connecting.html
-        client_input = ElasticsearchClientInput()
+        client_input_kwargs: dict[str, Any] = {}
+        access_config = self.access_config.get_secret_value()
         if self.hosts:
-            client_input.hosts = self.hosts
+            client_input_kwargs["hosts"] = self.hosts
         if self.cloud_id:
-            client_input.cloud_id = self.cloud_id
+            client_input_kwargs["cloud_id"] = self.cloud_id
         if self.ca_certs:
-            client_input.ca_certs = self.ca_certs
-        if self.access_config.password and (
-            self.cloud_id or self.ca_certs or self.access_config.ssl_assert_fingerprint
+            client_input_kwargs["ca_certs"] = self.ca_certs
+        if access_config.password and (
+            self.cloud_id or self.ca_certs or access_config.ssl_assert_fingerprint
         ):
-            client_input.basic_auth = ("elastic", self.access_config.password)
-        elif not self.cloud_id and self.username and self.access_config.password:
-            client_input.basic_auth = (self.username, self.access_config.password)
-        elif self.access_config.api_key and self.api_key_id:
-            client_input.api_key = (self.api_key_id, self.access_config.api_key)
-        elif self.access_config.api_key:
-            client_input.api_key = self.access_config.api_key
-        logger.debug(
-            f"Elasticsearch client inputs mapped to: {client_input.to_dict(redact_sensitive=True)}"
+            client_input_kwargs["basic_auth"] = ("elastic", access_config.password)
+        elif not self.cloud_id and self.username and access_config.password:
+            client_input_kwargs["basic_auth"] = (self.username, access_config.password)
+        elif access_config.es_api_key and self.api_key_id:
+            client_input_kwargs["api_key"] = (self.api_key_id, access_config.es_api_key)
+        elif access_config.es_api_key:
+            client_input_kwargs["api_key"] = access_config.es_api_key
+        client_input = ElasticsearchClientInput(**client_input_kwargs)
+        logger.debug(f"Elasticsearch client inputs mapped to: {client_input.dict()}")
+        client_kwargs = client_input.dict()
+        client_kwargs["basic_auth"] = (
+            client_input.basic_auth.get_secret_value() if client_input.basic_auth else None
+        )
+        client_kwargs["api_key"] = (
+            client_input.api_key.get_secret_value() if client_input.api_key else None
         )
-        client_kwargs = client_input.to_dict(redact_sensitive=False)
         client_kwargs = {k: v for k, v in client_kwargs.items() if v is not None}
         return client_kwargs
 
@@ -111,7 +132,6 @@ class ElasticsearchConnectionConfig(ConnectionConfig):
             raise SourceConnectionError(f"failed to validate connection: {e}")
 
 
-@dataclass
 class ElasticsearchIndexerConfig(IndexerConfig):
     index_name: str
     batch_size: int = 100
@@ -121,11 +141,14 @@ class ElasticsearchIndexerConfig(IndexerConfig):
 class ElasticsearchIndexer(Indexer):
     connection_config: ElasticsearchConnectionConfig
     index_config: ElasticsearchIndexerConfig
-    client: "ElasticsearchClient" = field(init=False)
     connector_type: str = CONNECTOR_TYPE
 
-    def __post_init__(self):
-        self.client = self.connection_config.get_client()
+    def precheck(self) -> None:
+        try:
+            self.connection_config.get_client()
+        except Exception as e:
+            logger.error(f"failed to validate connection: {e}", exc_info=True)
+            raise SourceConnectionError(f"failed to validate connection: {e}")
 
     @requires_dependencies(["elasticsearch"], extras="elasticsearch")
     def load_scan(self):
@@ -138,8 +161,9 @@ class ElasticsearchIndexer(Indexer):
         scan = self.load_scan()
 
         scan_query: dict = {"stored_fields": [], "query": {"match_all": {}}}
+        client = self.connection_config.get_client()
         hits = scan(
-            self.client,
+            client,
             query=scan_query,
             scroll="1m",
             index=self.index_config.index_name,
@@ -168,7 +192,7 @@ class ElasticsearchIndexer(Indexer):
             yield FileData(
                 identifier=identified,
                 connector_type=CONNECTOR_TYPE,
-                metadata=DataSourceMetadata(
+                metadata=FileDataSourceMetadata(
                     url=f"{self.connection_config.hosts[0]}/{self.index_config.index_name}",
                     date_processed=str(time()),
                 ),
@@ -179,7 +203,6 @@ class ElasticsearchIndexer(Indexer):
             )
 
 
-@dataclass
 class ElasticsearchDownloaderConfig(DownloaderConfig):
     fields: list[str] = field(default_factory=list)
 
@@ -234,7 +257,7 @@ class ElasticsearchDownloader(Downloader):
             file_data=FileData(
                 identifier=filename_id,
                 connector_type=CONNECTOR_TYPE,
-                metadata=DataSourceMetadata(
+                metadata=FileDataSourceMetadata(
                     version=str(result["_version"]) if "_version" in result else None,
                     date_processed=str(time()),
                     record_locator={
@@ -285,9 +308,10 @@ class ElasticsearchDownloader(Downloader):
         return download_responses
 
 
-@dataclass
 class ElasticsearchUploadStagerConfig(UploadStagerConfig):
-    index_name: str
+    index_name: str = Field(
+        description="Name of the Elasticsearch index to pull data from, or upload data to."
+    )
 
 
 @dataclass
@@ -326,11 +350,19 @@ class ElasticsearchUploadStager(UploadStager):
         return output_path
 
 
-@dataclass
 class ElasticsearchUploaderConfig(UploaderConfig):
-    index_name: str
-    batch_size_bytes: int = 15_000_000
-    num_threads: int = 4
+    index_name: str = Field(
+        description="Name of the Elasticsearch index to pull data from, or upload data to."
+    )
+    batch_size_bytes: int = Field(
+        default=15_000_000,
+        description="Size limit (in bytes) for each batch of items to be uploaded. Check"
+        " https://www.elastic.co/guide/en/elasticsearch/guide/current/bulk.html"
+        "#_how_big_is_too_big for more information.",
+    )
+    num_threads: int = Field(
+        default=4, description="Number of threads to be used while uploading content"
+    )
 
 
 @dataclass
@@ -339,6 +371,13 @@ class ElasticsearchUploader(Uploader):
     upload_config: ElasticsearchUploaderConfig
     connection_config: ElasticsearchConnectionConfig
 
+    def precheck(self) -> None:
+        try:
+            self.connection_config.get_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(["elasticsearch"], extras="elasticsearch")
     def load_parallel_bulk(self):
         from elasticsearch.helpers import parallel_bulk