unstructured-ingest 0.3.0__py3-none-any.whl → 0.3.1__py3-none-any.whl

unstructured_ingest/v2/processes/connectors/{elasticsearch.py → elasticsearch/elasticsearch.py}

@@ -2,6 +2,7 @@ import hashlib
 import json
 import sys
 import uuid
+from contextlib import contextmanager
 from dataclasses import dataclass, field
 from pathlib import Path
 from time import time
@@ -13,9 +14,11 @@ from unstructured_ingest.error import (
     DestinationConnectionError,
     SourceConnectionError,
     SourceConnectionNetworkError,
+    WriteError,
 )
 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.constants import RECORD_ID_LABEL
 from unstructured_ingest.v2.interfaces import (
     AccessConfig,
     ConnectionConfig,
@@ -26,6 +29,7 @@ from unstructured_ingest.v2.interfaces import (
     FileDataSourceMetadata,
     Indexer,
     IndexerConfig,
+    SourceIdentifiers,
     Uploader,
     UploaderConfig,
     UploadStager,
@@ -116,19 +120,12 @@ class ElasticsearchConnectionConfig(ConnectionConfig):
         return client_kwargs
 
     @requires_dependencies(["elasticsearch"], extras="elasticsearch")
-    def get_client(self) -> "ElasticsearchClient":
+    @contextmanager
+    def get_client(self) -> Generator["ElasticsearchClient", None, None]:
         from elasticsearch import Elasticsearch as ElasticsearchClient
 
-        client = ElasticsearchClient(**self.get_client_kwargs())
-        self.check_connection(client=client)
-        return client
-
-    def check_connection(self, client: "ElasticsearchClient"):
-        try:
-            client.perform_request("HEAD", "/", headers={"accept": "application/json"})
-        except Exception as e:
-            logger.error(f"failed to validate connection: {e}", exc_info=True)
-            raise SourceConnectionError(f"failed to validate connection: {e}")
+        with ElasticsearchClient(**self.get_client_kwargs()) as client:
+            yield client
 
 
 class ElasticsearchIndexerConfig(IndexerConfig):
@@ -144,7 +141,16 @@ class ElasticsearchIndexer(Indexer):
 
     def precheck(self) -> None:
         try:
-            self.connection_config.get_client()
+            with self.connection_config.get_client() as client:
+                if not client.ping():
+                    raise SourceConnectionError("cluster not detected")
+                indices = client.indices.get_alias(index="*")
+                if self.index_config.index_name not in indices:
+                    raise SourceConnectionError(
+                        "index {} not found: {}".format(
+                            self.index_config.index_name, ", ".join(indices.keys())
+                        )
+                    )
         except Exception as e:
             logger.error(f"failed to validate connection: {e}", exc_info=True)
             raise SourceConnectionError(f"failed to validate connection: {e}")
@@ -160,15 +166,15 @@ class ElasticsearchIndexer(Indexer):
         scan = self.load_scan()
 
         scan_query: dict = {"stored_fields": [], "query": {"match_all": {}}}
-        client = self.connection_config.get_client()
-        hits = scan(
-            client,
-            query=scan_query,
-            scroll="1m",
-            index=self.index_config.index_name,
-        )
+        with self.connection_config.get_client() as client:
+            hits = scan(
+                client,
+                query=scan_query,
+                scroll="1m",
+                index=self.index_config.index_name,
+            )
 
-        return {hit["_id"] for hit in hits}
+            return {hit["_id"] for hit in hits}
 
     def run(self, **kwargs: Any) -> Generator[FileData, None, None]:
         all_ids = self._get_doc_ids()
@@ -257,6 +263,7 @@ class ElasticsearchDownloader(Downloader):
             file_data=FileData(
                 identifier=filename_id,
                 connector_type=CONNECTOR_TYPE,
+                source_identifiers=SourceIdentifiers(filename=filename, fullpath=filename),
                 metadata=FileDataSourceMetadata(
                     version=str(result["_version"]) if "_version" in result else None,
                     date_processed=str(time()),
@@ -318,7 +325,7 @@ class ElasticsearchUploadStagerConfig(UploadStagerConfig):
 class ElasticsearchUploadStager(UploadStager):
     upload_stager_config: ElasticsearchUploadStagerConfig
 
-    def conform_dict(self, data: dict) -> dict:
+    def conform_dict(self, data: dict, file_data: FileData) -> dict:
         resp = {
             "_index": self.upload_stager_config.index_name,
             "_id": str(uuid.uuid4()),
@@ -327,6 +334,7 @@ class ElasticsearchUploadStager(UploadStager):
                 "embeddings": data.pop("embeddings", None),
                 "text": data.pop("text", None),
                 "type": data.pop("type", None),
+                RECORD_ID_LABEL: file_data.identifier,
             },
         }
         if "metadata" in data and isinstance(data["metadata"], dict):
@@ -343,10 +351,17 @@ class ElasticsearchUploadStager(UploadStager):
     ) -> Path:
         with open(elements_filepath) as elements_file:
             elements_contents = json.load(elements_file)
-        conformed_elements = [self.conform_dict(data=element) for element in elements_contents]
-        output_path = Path(output_dir) / Path(f"{output_filename}.json")
+        conformed_elements = [
+            self.conform_dict(data=element, file_data=file_data) for element in elements_contents
+        ]
+        if Path(output_filename).suffix != ".json":
+            output_filename = f"{output_filename}.json"
+        else:
+            output_filename = f"{Path(output_filename).stem}.json"
+        output_path = Path(output_dir) / output_filename
+        output_path.parent.mkdir(parents=True, exist_ok=True)
         with open(output_path, "w") as output_file:
-            json.dump(conformed_elements, output_file)
+            json.dump(conformed_elements, output_file, indent=2)
         return output_path
 
 
@@ -363,6 +378,10 @@ class ElasticsearchUploaderConfig(UploaderConfig):
     num_threads: int = Field(
         default=4, description="Number of threads to be used while uploading content"
     )
+    record_id_key: str = Field(
+        default=RECORD_ID_LABEL,
+        description="searchable key to find entries for the same record on previous runs",
+    )
 
 
 @dataclass
@@ -373,7 +392,16 @@ class ElasticsearchUploader(Uploader):
 
     def precheck(self) -> None:
         try:
-            self.connection_config.get_client()
+            with self.connection_config.get_client() as client:
+                if not client.ping():
+                    raise DestinationConnectionError("cluster not detected")
+                indices = client.indices.get_alias(index="*")
+                if self.upload_config.index_name not in indices:
+                    raise SourceConnectionError(
+                        "index {} not found: {}".format(
+                            self.upload_config.index_name, ", ".join(indices.keys())
+                        )
+                    )
         except Exception as e:
             logger.error(f"failed to validate connection: {e}", exc_info=True)
             raise DestinationConnectionError(f"failed to validate connection: {e}")
@@ -384,6 +412,23 @@ class ElasticsearchUploader(Uploader):
 
         return parallel_bulk
 
+    def delete_by_record_id(self, client, file_data: FileData) -> None:
+        logger.debug(
+            f"deleting any content with metadata {RECORD_ID_LABEL}={file_data.identifier} "
+            f"from {self.upload_config.index_name} index"
+        )
+        delete_resp = client.delete_by_query(
+            index=self.upload_config.index_name,
+            body={"query": {"match": {self.upload_config.record_id_key: file_data.identifier}}},
+        )
+        logger.info(
+            "deleted {} records from index {}".format(
+                delete_resp["deleted"], self.upload_config.index_name
+            )
+        )
+        if failures := delete_resp.get("failures"):
+            raise WriteError(f"failed to delete records: {failures}")
+
     def run(self, path: Path, file_data: FileData, **kwargs: Any) -> None:
         parallel_bulk = self.load_parallel_bulk()
         with path.open("r") as file:
@@ -397,28 +442,29 @@ class ElasticsearchUploader(Uploader):
             f"{self.upload_config.num_threads} (number of) threads"
         )
 
-        client = self.connection_config.get_client()
-        if not client.indices.exists(index=self.upload_config.index_name):
-            logger.warning(
-                f"{(self.__class__.__name__).replace('Uploader', '')} index does not exist: "
-                f"{self.upload_config.index_name}. "
-                f"This may cause issues when uploading."
-            )
-        for batch in generator_batching_wbytes(
-            elements_dict, batch_size_limit_bytes=self.upload_config.batch_size_bytes
-        ):
-            for success, info in parallel_bulk(
-                client=client,
-                actions=batch,
-                thread_count=self.upload_config.num_threads,
+        with self.connection_config.get_client() as client:
+            self.delete_by_record_id(client=client, file_data=file_data)
+            if not client.indices.exists(index=self.upload_config.index_name):
+                logger.warning(
+                    f"{(self.__class__.__name__).replace('Uploader', '')} index does not exist: "
+                    f"{self.upload_config.index_name}. "
+                    f"This may cause issues when uploading."
+                )
+            for batch in generator_batching_wbytes(
+                elements_dict, batch_size_limit_bytes=self.upload_config.batch_size_bytes
             ):
-                if not success:
-                    logger.error(
-                        "upload failed for a batch in "
-                        f"{(self.__class__.__name__).replace('Uploader', '')} "
-                        "destination connector:",
-                        info,
-                    )
+                for success, info in parallel_bulk(
+                    client=client,
+                    actions=batch,
+                    thread_count=self.upload_config.num_threads,
+                ):
+                    if not success:
+                        logger.error(
+                            "upload failed for a batch in "
+                            f"{(self.__class__.__name__).replace('Uploader', '')} "
+                            "destination connector:",
+                            info,
+                        )
 
 
 elasticsearch_source_entry = SourceRegistryEntry(

unstructured_ingest/v2/processes/connectors/{opensearch.py → elasticsearch/opensearch.py}

@@ -17,7 +17,7 @@ from unstructured_ingest.v2.processes.connector_registry import (
     DestinationRegistryEntry,
     SourceRegistryEntry,
 )
-from unstructured_ingest.v2.processes.connectors.elasticsearch import (
+from unstructured_ingest.v2.processes.connectors.elasticsearch.elasticsearch import (
     ElasticsearchDownloader,
     ElasticsearchDownloaderConfig,
     ElasticsearchIndexer,

unstructured_ingest/v2/processes/connectors/kafka/kafka.py

@@ -161,6 +161,12 @@ class KafkaIndexer(Indexer, ABC):
                 current_topics = [
                     topic for topic in cluster_meta.topics if topic != "__consumer_offsets"
                 ]
+                if self.index_config.topic not in current_topics:
+                    raise SourceConnectionError(
+                        "expected topic {} not detected in cluster: {}".format(
+                            self.index_config.topic, ", ".join(current_topics)
+                        )
+                    )
                 logger.info(f"successfully checked available topics: {current_topics}")
         except Exception as e:
             logger.error(f"failed to validate connection: {e}", exc_info=True)

unstructured_ingest/v2/processes/connectors/lancedb/__init__.py

@@ -0,0 +1,17 @@
+from __future__ import annotations
+
+from unstructured_ingest.v2.processes.connector_registry import add_destination_entry
+
+from .aws import CONNECTOR_TYPE as LANCEDB_S3_CONNECTOR_TYPE
+from .aws import lancedb_aws_destination_entry
+from .azure import CONNECTOR_TYPE as LANCEDB_AZURE_CONNECTOR_TYPE
+from .azure import lancedb_azure_destination_entry
+from .gcp import CONNECTOR_TYPE as LANCEDB_GCS_CONNECTOR_TYPE
+from .gcp import lancedb_gcp_destination_entry
+from .local import CONNECTOR_TYPE as LANCEDB_LOCAL_CONNECTOR_TYPE
+from .local import lancedb_local_destination_entry
+
+add_destination_entry(LANCEDB_S3_CONNECTOR_TYPE, lancedb_aws_destination_entry)
+add_destination_entry(LANCEDB_AZURE_CONNECTOR_TYPE, lancedb_azure_destination_entry)
+add_destination_entry(LANCEDB_GCS_CONNECTOR_TYPE, lancedb_gcp_destination_entry)
+add_destination_entry(LANCEDB_LOCAL_CONNECTOR_TYPE, lancedb_local_destination_entry)

unstructured_ingest/v2/processes/connectors/lancedb/aws.py

@@ -0,0 +1,43 @@
+from dataclasses import dataclass
+
+from pydantic import Field, Secret
+
+from unstructured_ingest.v2.interfaces.connector import AccessConfig
+from unstructured_ingest.v2.processes.connector_registry import DestinationRegistryEntry
+from unstructured_ingest.v2.processes.connectors.lancedb.lancedb import (
+    LanceDBRemoteConnectionConfig,
+    LanceDBUploader,
+    LanceDBUploaderConfig,
+    LanceDBUploadStager,
+    LanceDBUploadStagerConfig,
+)
+
+CONNECTOR_TYPE = "lancedb_aws"
+
+
+class LanceDBS3AccessConfig(AccessConfig):
+    aws_access_key_id: str = Field(description="The AWS access key ID to use.")
+    aws_secret_access_key: str = Field(description="The AWS secret access key to use.")
+
+
+class LanceDBS3ConnectionConfig(LanceDBRemoteConnectionConfig):
+    access_config: Secret[LanceDBS3AccessConfig]
+
+    def get_storage_options(self) -> dict:
+        return {**self.access_config.get_secret_value().model_dump(), "timeout": self.timeout}
+
+
+@dataclass
+class LanceDBS3Uploader(LanceDBUploader):
+    upload_config: LanceDBUploaderConfig
+    connection_config: LanceDBS3ConnectionConfig
+    connector_type: str = CONNECTOR_TYPE
+
+
+lancedb_aws_destination_entry = DestinationRegistryEntry(
+    connection_config=LanceDBS3ConnectionConfig,
+    uploader=LanceDBS3Uploader,
+    uploader_config=LanceDBUploaderConfig,
+    upload_stager_config=LanceDBUploadStagerConfig,
+    upload_stager=LanceDBUploadStager,
+)

unstructured_ingest/v2/processes/connectors/lancedb/azure.py

@@ -0,0 +1,43 @@
+from dataclasses import dataclass
+
+from pydantic import Field, Secret
+
+from unstructured_ingest.v2.interfaces.connector import AccessConfig
+from unstructured_ingest.v2.processes.connector_registry import DestinationRegistryEntry
+from unstructured_ingest.v2.processes.connectors.lancedb.lancedb import (
+    LanceDBRemoteConnectionConfig,
+    LanceDBUploader,
+    LanceDBUploaderConfig,
+    LanceDBUploadStager,
+    LanceDBUploadStagerConfig,
+)
+
+CONNECTOR_TYPE = "lancedb_azure"
+
+
+class LanceDBAzureAccessConfig(AccessConfig):
+    azure_storage_account_name: str = Field(description="The name of the azure storage account.")
+    azure_storage_account_key: str = Field(description="The serialized azure service account key.")
+
+
+class LanceDBAzureConnectionConfig(LanceDBRemoteConnectionConfig):
+    access_config: Secret[LanceDBAzureAccessConfig]
+
+    def get_storage_options(self) -> dict:
+        return {**self.access_config.get_secret_value().model_dump(), "timeout": self.timeout}
+
+
+@dataclass
+class LanceDBAzureUploader(LanceDBUploader):
+    upload_config: LanceDBUploaderConfig
+    connection_config: LanceDBAzureConnectionConfig
+    connector_type: str = CONNECTOR_TYPE
+
+
+lancedb_azure_destination_entry = DestinationRegistryEntry(
+    connection_config=LanceDBAzureConnectionConfig,
+    uploader=LanceDBAzureUploader,
+    uploader_config=LanceDBUploaderConfig,
+    upload_stager_config=LanceDBUploadStagerConfig,
+    upload_stager=LanceDBUploadStager,
+)

unstructured_ingest/v2/processes/connectors/lancedb/gcp.py

@@ -0,0 +1,44 @@
+from dataclasses import dataclass
+
+from pydantic import Field, Secret
+
+from unstructured_ingest.v2.interfaces.connector import AccessConfig
+from unstructured_ingest.v2.processes.connector_registry import DestinationRegistryEntry
+from unstructured_ingest.v2.processes.connectors.lancedb.lancedb import (
+    LanceDBRemoteConnectionConfig,
+    LanceDBUploader,
+    LanceDBUploaderConfig,
+    LanceDBUploadStager,
+    LanceDBUploadStagerConfig,
+)
+
+CONNECTOR_TYPE = "lancedb_gcs"
+
+
+class LanceDBGCSAccessConfig(AccessConfig):
+    google_service_account_key: str = Field(
+        description="The serialized google service account key."
+    )
+
+
+class LanceDBGCSConnectionConfig(LanceDBRemoteConnectionConfig):
+    access_config: Secret[LanceDBGCSAccessConfig]
+
+    def get_storage_options(self) -> dict:
+        return {**self.access_config.get_secret_value().model_dump(), "timeout": self.timeout}
+
+
+@dataclass
+class LanceDBGSPUploader(LanceDBUploader):
+    upload_config: LanceDBUploaderConfig
+    connection_config: LanceDBGCSConnectionConfig
+    connector_type: str = CONNECTOR_TYPE
+
+
+lancedb_gcp_destination_entry = DestinationRegistryEntry(
+    connection_config=LanceDBGCSConnectionConfig,
+    uploader=LanceDBGSPUploader,
+    uploader_config=LanceDBUploaderConfig,
+    upload_stager_config=LanceDBUploadStagerConfig,
+    upload_stager=LanceDBUploadStager,
+)

unstructured_ingest/v2/processes/connectors/lancedb/lancedb.py

@@ -0,0 +1,161 @@
+from __future__ import annotations
+
+import asyncio
+import json
+from abc import ABC, abstractmethod
+from contextlib import asynccontextmanager
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, AsyncGenerator, Optional
+
+import pandas as pd
+from pydantic import Field
+
+from unstructured_ingest.error import DestinationConnectionError
+from unstructured_ingest.logger import logger
+from unstructured_ingest.utils.data_prep import flatten_dict
+from unstructured_ingest.utils.dep_check import requires_dependencies
+from unstructured_ingest.v2.interfaces.connector import ConnectionConfig
+from unstructured_ingest.v2.interfaces.file_data import FileData
+from unstructured_ingest.v2.interfaces.upload_stager import UploadStager, UploadStagerConfig
+from unstructured_ingest.v2.interfaces.uploader import Uploader, UploaderConfig
+
+CONNECTOR_TYPE = "lancedb"
+
+if TYPE_CHECKING:
+    from lancedb import AsyncConnection
+    from lancedb.table import AsyncTable
+
+
+class LanceDBConnectionConfig(ConnectionConfig, ABC):
+    uri: str = Field(description="The uri of the database.")
+
+    @abstractmethod
+    def get_storage_options(self) -> Optional[dict[str, str]]:
+        raise NotImplementedError
+
+    @asynccontextmanager
+    @requires_dependencies(["lancedb"], extras="lancedb")
+    @DestinationConnectionError.wrap
+    async def get_async_connection(self) -> AsyncGenerator["AsyncConnection", None]:
+        import lancedb
+
+        connection = await lancedb.connect_async(
+            self.uri,
+            storage_options=self.get_storage_options(),
+        )
+        try:
+            yield connection
+        finally:
+            connection.close()
+
+
+class LanceDBRemoteConnectionConfig(LanceDBConnectionConfig):
+    timeout: str = Field(
+        default="30s",
+        description=(
+            "Timeout for the entire request, from connection until the response body has finished"
+            "in a [0-9]+(ns|us|ms|[smhdwy]) format."
+        ),
+        pattern=r"[0-9]+(ns|us|ms|[smhdwy])",
+    )
+
+
+class LanceDBUploadStagerConfig(UploadStagerConfig):
+    pass
+
+
+@dataclass
+class LanceDBUploadStager(UploadStager):
+    upload_stager_config: LanceDBUploadStagerConfig = field(
+        default_factory=LanceDBUploadStagerConfig
+    )
+
+    def run(
+        self,
+        elements_filepath: Path,
+        file_data: FileData,
+        output_dir: Path,
+        output_filename: str,
+        **kwargs: Any,
+    ) -> Path:
+        with open(elements_filepath) as elements_file:
+            elements_contents: list[dict] = json.load(elements_file)
+
+        df = pd.DataFrame(
+            [
+                self._conform_element_contents(element_contents)
+                for element_contents in elements_contents
+            ]
+        )
+
+        output_path = (output_dir / output_filename).with_suffix(".feather")
+        df.to_feather(output_path)
+
+        return output_path
+
+    def _conform_element_contents(self, element: dict) -> dict:
+        return {
+            "vector": element.pop("embeddings", None),
+            **flatten_dict(element, separator="-"),
+        }
+
+
+class LanceDBUploaderConfig(UploaderConfig):
+    table_name: str = Field(description="The name of the table.")
+
+
+@dataclass
+class LanceDBUploader(Uploader):
+    upload_config: LanceDBUploaderConfig
+    connection_config: LanceDBConnectionConfig
+    connector_type: str = CONNECTOR_TYPE
+
+    @DestinationConnectionError.wrap
+    def precheck(self):
+        async def _precheck() -> None:
+            async with self.connection_config.get_async_connection() as conn:
+                table = await conn.open_table(self.upload_config.table_name)
+                table.close()
+
+        asyncio.run(_precheck())
+
+    @asynccontextmanager
+    async def get_table(self) -> AsyncGenerator["AsyncTable", None]:
+        async with self.connection_config.get_async_connection() as conn:
+            table = await conn.open_table(self.upload_config.table_name)
+            try:
+                yield table
+            finally:
+                table.close()
+
+    async def run_async(self, path, file_data, **kwargs):
+        df = pd.read_feather(path)
+        async with self.get_table() as table:
+            schema = await table.schema()
+            df = self._fit_to_schema(df, schema)
+            await table.add(data=df)
+
+    def _fit_to_schema(self, df: pd.DataFrame, schema) -> pd.DataFrame:
+        columns = set(df.columns)
+        schema_fields = set(schema.names)
+        columns_to_drop = columns - schema_fields
+        missing_columns = schema_fields - columns
+
+        if columns_to_drop:
+            logger.info(
+                "Following columns will be dropped to match the table's schema: "
+                f"{', '.join(columns_to_drop)}"
+            )
+        if missing_columns:
+            logger.info(
+                "Following null filled columns will be added to match the table's schema:"
+                f" {', '.join(missing_columns)} "
+            )
+
+        df = df.drop(columns=columns_to_drop)
+
+        for column in missing_columns:
+            df[column] = pd.Series()
+
+        return df

unstructured_ingest/v2/processes/connectors/lancedb/local.py

@@ -0,0 +1,44 @@
+from dataclasses import dataclass
+
+from pydantic import Field, Secret
+
+from unstructured_ingest.v2.interfaces.connector import AccessConfig
+from unstructured_ingest.v2.processes.connector_registry import DestinationRegistryEntry
+from unstructured_ingest.v2.processes.connectors.lancedb.lancedb import (
+    LanceDBConnectionConfig,
+    LanceDBUploader,
+    LanceDBUploaderConfig,
+    LanceDBUploadStager,
+    LanceDBUploadStagerConfig,
+)
+
+CONNECTOR_TYPE = "lancedb_local"
+
+
+class LanceDBLocalAccessConfig(AccessConfig):
+    pass
+
+
+class LanceDBLocalConnectionConfig(LanceDBConnectionConfig):
+    access_config: Secret[LanceDBLocalAccessConfig] = Field(
+        default_factory=LanceDBLocalAccessConfig, validate_default=True
+    )
+
+    def get_storage_options(self) -> None:
+        return None
+
+
+@dataclass
+class LanceDBLocalUploader(LanceDBUploader):
+    upload_config: LanceDBUploaderConfig
+    connection_config: LanceDBLocalConnectionConfig
+    connector_type: str = CONNECTOR_TYPE
+
+
+lancedb_local_destination_entry = DestinationRegistryEntry(
+    connection_config=LanceDBLocalConnectionConfig,
+    uploader=LanceDBLocalUploader,
+    uploader_config=LanceDBUploaderConfig,
+    upload_stager_config=LanceDBUploadStagerConfig,
+    upload_stager=LanceDBUploadStager,
+)