unstructured-ingest 0.3.8__py3-none-any.whl → 0.3.10__py3-none-any.whl

unstructured_ingest/v2/processes/connectors/duckdb/base.py

@@ -1,5 +1,3 @@
-import json
-import uuid
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Any
@@ -7,6 +5,7 @@ from typing import Any
 import pandas as pd
 
 from unstructured_ingest.v2.interfaces import FileData, UploadStager
+from unstructured_ingest.v2.utils import get_enhanced_element_id
 
 _COLUMNS = (
     "id",
@@ -56,6 +55,22 @@ _COLUMNS = (
 @dataclass
 class BaseDuckDBUploadStager(UploadStager):
 
+    def conform_dict(self, element_dict: dict, file_data: FileData) -> dict:
+        data = element_dict.copy()
+        metadata: dict[str, Any] = data.pop("metadata", {})
+        data_source = metadata.pop("data_source", {})
+        coordinates = metadata.pop("coordinates", {})
+
+        data.update(metadata)
+        data.update(data_source)
+        data.update(coordinates)
+
+        data["id"] = get_enhanced_element_id(element_dict=data, file_data=file_data)
+
+        # remove extraneous, not supported columns
+        data = {k: v for k, v in data.items() if k in _COLUMNS}
+        return data
+
     def run(
         self,
         elements_filepath: Path,
@@ -64,29 +79,14 @@ class BaseDuckDBUploadStager(UploadStager):
         output_filename: str,
         **kwargs: Any,
     ) -> Path:
-        with open(elements_filepath) as elements_file:
-            elements_contents: list[dict] = json.load(elements_file)
-        output_path = Path(output_dir) / Path(f"{output_filename}.json")
-        output_path.parent.mkdir(parents=True, exist_ok=True)
-
-        output = []
-        for data in elements_contents:
-            metadata: dict[str, Any] = data.pop("metadata", {})
-            data_source = metadata.pop("data_source", {})
-            coordinates = metadata.pop("coordinates", {})
-
-            data.update(metadata)
-            data.update(data_source)
-            data.update(coordinates)
-
-            data["id"] = str(uuid.uuid4())
-
-            # remove extraneous, not supported columns
-            data = {k: v for k, v in data.items() if k in _COLUMNS}
-
-            output.append(data)
+        elements_contents = self.get_data(elements_filepath=elements_filepath)
+        output_path = self.get_output_path(output_filename=output_filename, output_dir=output_dir)
 
-        df = pd.DataFrame.from_dict(output)
+        output = [
+            self.conform_dict(element_dict=element_dict, file_data=file_data)
+            for element_dict in elements_contents
+        ]
+        df = pd.DataFrame(data=output)
 
         for column in filter(
             lambda x: x in df.columns,
@@ -94,6 +94,6 @@ class BaseDuckDBUploadStager(UploadStager):
         ):
             df[column] = df[column].apply(str)
 
-        with output_path.open("w") as output_file:
-            df.to_json(output_file, orient="records", lines=True)
+        data = df.to_dict(orient="records")
+        self.write_output(output_path=output_path, data=data)
         return output_path

unstructured_ingest/v2/processes/connectors/duckdb/duckdb.py

@@ -1,11 +1,13 @@
+from contextlib import contextmanager
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Callable, Optional
+from typing import TYPE_CHECKING, Any, Generator, Optional
 
 import pandas as pd
 from pydantic import Field, Secret
 
 from unstructured_ingest.error import DestinationConnectionError
+from unstructured_ingest.utils.data_prep import get_data_df
 from unstructured_ingest.utils.dep_check import requires_dependencies
 from unstructured_ingest.v2.interfaces import (
     AccessConfig,
@@ -55,6 +57,20 @@ class DuckDBConnectionConfig(ConnectionConfig):
                 "through the `database` argument"
             )
 
+    @requires_dependencies(["duckdb"], extras="duckdb")
+    @contextmanager
+    def get_client(self) -> Generator["DuckDBConnection", None, None]:
+        import duckdb
+
+        with duckdb.connect(self.database) as client:
+            yield client
+
+    @contextmanager
+    def get_cursor(self) -> Generator["DuckDBConnection", None, None]:
+        with self.get_client() as client:
+            with client.cursor() as cursor:
+                yield cursor
+
 
 class DuckDBUploadStagerConfig(UploadStagerConfig):
     pass
@@ -79,34 +95,27 @@ class DuckDBUploader(Uploader):
 
     def precheck(self) -> None:
         try:
-            cursor = self.connection().cursor()
-            cursor.execute("SELECT 1;")
-            cursor.close()
+            with self.connection_config.get_cursor() as cursor:
+                cursor.execute("SELECT 1;")
         except Exception as e:
             logger.error(f"failed to validate connection: {e}", exc_info=True)
             raise DestinationConnectionError(f"failed to validate connection: {e}")
 
-    @property
-    def connection(self) -> Callable[[], "DuckDBConnection"]:
-        return self._make_duckdb_connection
+    def upload_dataframe(self, df: pd.DataFrame) -> None:
+        logger.debug(f"uploading {len(df)} entries to {self.connection_config.database} ")
 
-    @requires_dependencies(["duckdb"], extras="duckdb")
-    def _make_duckdb_connection(self) -> "DuckDBConnection":
-        import duckdb
-
-        return duckdb.connect(self.connection_config.database)
-
-    def upload_contents(self, path: Path) -> None:
-        df_elements = pd.read_json(path, orient="records", lines=True)
-        logger.debug(f"uploading {len(df_elements)} entries to {self.connection_config.database} ")
-
-        with self.connection() as conn:
+        with self.connection_config.get_client() as conn:
             conn.query(
-                f"INSERT INTO {self.connection_config.db_schema}.{self.connection_config.table} BY NAME SELECT * FROM df_elements"  # noqa: E501
+                f"INSERT INTO {self.connection_config.db_schema}.{self.connection_config.table} BY NAME SELECT * FROM df"  # noqa: E501
             )
 
+    def run_data(self, data: list[dict], file_data: FileData, **kwargs: Any) -> None:
+        df = pd.DataFrame(data=data)
+        self.upload_dataframe(df=df)
+
     def run(self, path: Path, file_data: FileData, **kwargs: Any) -> None:
-        self.upload_contents(path=path)
+        df = get_data_df(path)
+        self.upload_dataframe(df=df)
 
 
 duckdb_destination_entry = DestinationRegistryEntry(

unstructured_ingest/v2/processes/connectors/duckdb/motherduck.py

@@ -1,12 +1,14 @@
+from contextlib import contextmanager
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Callable, Optional
+from typing import TYPE_CHECKING, Any, Generator, Optional
 
 import pandas as pd
 from pydantic import Field, Secret
 
 from unstructured_ingest.__version__ import __version__ as unstructured_io_ingest_version
 from unstructured_ingest.error import DestinationConnectionError
+from unstructured_ingest.utils.data_prep import get_data_df
 from unstructured_ingest.utils.dep_check import requires_dependencies
 from unstructured_ingest.v2.interfaces import (
     AccessConfig,
@@ -27,13 +29,12 @@ CONNECTOR_TYPE = "motherduck"
 
 
 class MotherDuckAccessConfig(AccessConfig):
-    md_token: Optional[str] = Field(default=None, description="MotherDuck token")
+    md_token: str = Field(default=None, description="MotherDuck token")
 
 
 class MotherDuckConnectionConfig(ConnectionConfig):
     connector_type: str = Field(default=CONNECTOR_TYPE, init=False)
-    database: Optional[str] = Field(
-        default=None,
+    database: str = Field(
         description="Database name. Name of the MotherDuck database.",
     )
     db_schema: Optional[str] = Field(
@@ -48,17 +49,26 @@ class MotherDuckConnectionConfig(ConnectionConfig):
         default=MotherDuckAccessConfig(), validate_default=True
     )
 
-    def __post_init__(self):
-        if self.database is None:
-            raise ValueError(
-                "A MotherDuck connection requires a database (string) to be passed "
-                "through the `database` argument"
-            )
-        if self.access_config.md_token is None:
-            raise ValueError(
-                "A MotherDuck connection requires a md_token (MotherDuck token) to be passed "
-                "using MotherDuckAccessConfig through the `access_config` argument"
-            )
+    @requires_dependencies(["duckdb"], extras="duckdb")
+    @contextmanager
+    def get_client(self) -> Generator["MotherDuckConnection", None, None]:
+        import duckdb
+
+        access_config = self.access_config.get_secret_value()
+        with duckdb.connect(
+            f"md:?motherduck_token={access_config.md_token}",
+            config={
+                "custom_user_agent": f"unstructured-io-ingest/{unstructured_io_ingest_version}"
+            },
+        ) as conn:
+            conn.sql(f"USE {self.database}")
+            yield conn
+
+    @contextmanager
+    def get_cursor(self) -> Generator["MotherDuckConnection", None, None]:
+        with self.get_client() as client:
+            with client.cursor() as cursor:
+                yield cursor
 
 
 class MotherDuckUploadStagerConfig(UploadStagerConfig):
@@ -84,44 +94,27 @@ class MotherDuckUploader(Uploader):
 
     def precheck(self) -> None:
         try:
-            cursor = self.connection().cursor()
-            cursor.execute("SELECT 1;")
-            cursor.close()
+            with self.connection_config.get_cursor() as cursor:
+                cursor.execute("SELECT 1;")
         except Exception as e:
             logger.error(f"failed to validate connection: {e}", exc_info=True)
             raise DestinationConnectionError(f"failed to validate connection: {e}")
 
-    @property
-    def connection(self) -> Callable[[], "MotherDuckConnection"]:
-        return self._make_motherduck_connection
-
-    @requires_dependencies(["duckdb"], extras="duckdb")
-    def _make_motherduck_connection(self) -> "MotherDuckConnection":
-        import duckdb
+    def upload_dataframe(self, df: pd.DataFrame) -> None:
+        logger.debug(f"uploading {len(df)} entries to {self.connection_config.database} ")
 
-        access_config = self.connection_config.access_config.get_secret_value()
-        conn = duckdb.connect(
-            f"md:?motherduck_token={access_config.md_token}",
-            config={
-                "custom_user_agent": f"unstructured-io-ingest/{unstructured_io_ingest_version}"
-            },
-        )
-
-        conn.sql(f"USE {self.connection_config.database}")
-
-        return conn
-
-    def upload_contents(self, path: Path) -> None:
-        df_elements = pd.read_json(path, orient="records", lines=True)
-        logger.debug(f"uploading {len(df_elements)} entries to {self.connection_config.database} ")
-
-        with self.connection() as conn:
+        with self.connection_config.get_client() as conn:
             conn.query(
-                f"INSERT INTO {self.connection_config.db_schema}.{self.connection_config.table} BY NAME SELECT * FROM df_elements"  # noqa: E501
+                f"INSERT INTO {self.connection_config.db_schema}.{self.connection_config.table} BY NAME SELECT * FROM df"  # noqa: E501
             )
 
+    def run_data(self, data: list[dict], file_data: FileData, **kwargs: Any) -> None:
+        df = pd.DataFrame(data=data)
+        self.upload_dataframe(df=df)
+
     def run(self, path: Path, file_data: FileData, **kwargs: Any) -> None:
-        self.upload_contents(path=path)
+        df = get_data_df(path)
+        self.upload_dataframe(df=df)
 
 
 motherduck_destination_entry = DestinationRegistryEntry(

unstructured_ingest/v2/processes/connectors/elasticsearch/elasticsearch.py

@@ -1,7 +1,5 @@
 import collections
 import hashlib
-import json
-import sys
 from contextlib import contextmanager
 from dataclasses import dataclass, field
 from pathlib import Path
@@ -16,11 +14,17 @@ from unstructured_ingest.error import (
     SourceConnectionNetworkError,
     WriteError,
 )
-from unstructured_ingest.utils.data_prep import flatten_dict, generator_batching_wbytes
+from unstructured_ingest.utils.data_prep import (
+    batch_generator,
+    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,
+    BatchFileData,
+    BatchItem,
     ConnectionConfig,
     Downloader,
     DownloaderConfig,
@@ -49,6 +53,14 @@ if TYPE_CHECKING:
 CONNECTOR_TYPE = "elasticsearch"
 
 
+class ElastisearchAdditionalMetadata(BaseModel):
+    index_name: str
+
+
+class ElasticsearchBatchFileData(BatchFileData):
+    additional_metadata: ElastisearchAdditionalMetadata
+
+
 class ElasticsearchAccessConfig(AccessConfig):
     password: Optional[str] = Field(
         default=None, description="password when using basic auth or connecting to a cloud instance"
@@ -175,36 +187,21 @@ class ElasticsearchIndexer(Indexer):
 
             return {hit["_id"] for hit in hits}
 
-    def run(self, **kwargs: Any) -> Generator[FileData, None, None]:
+    def run(self, **kwargs: Any) -> Generator[ElasticsearchBatchFileData, None, None]:
         all_ids = self._get_doc_ids()
         ids = list(all_ids)
-        id_batches: list[frozenset[str]] = [
-            frozenset(
-                ids[
-                    i
-                    * self.index_config.batch_size : (i + 1)  # noqa
-                    * self.index_config.batch_size
-                ]
-            )
-            for i in range(
-                (len(ids) + self.index_config.batch_size - 1) // self.index_config.batch_size
-            )
-        ]
-        for batch in id_batches:
+        for batch in batch_generator(ids, self.index_config.batch_size):
             # Make sure the hash is always a positive number to create identified
-            identified = str(hash(batch) + sys.maxsize + 1)
-            yield FileData(
-                identifier=identified,
+            yield ElasticsearchBatchFileData(
                 connector_type=CONNECTOR_TYPE,
-                doc_type="batch",
                 metadata=FileDataSourceMetadata(
                     url=f"{self.connection_config.hosts[0]}/{self.index_config.index_name}",
                     date_processed=str(time()),
                 ),
-                additional_metadata={
-                    "ids": list(batch),
-                    "index_name": self.index_config.index_name,
-                },
+                additional_metadata=ElastisearchAdditionalMetadata(
+                    index_name=self.index_config.index_name,
+                ),
+                batch_items=[BatchItem(identifier=b) for b in batch],
             )
 
 
@@ -238,7 +235,7 @@ class ElasticsearchDownloader(Downloader):
         return concatenated_values
 
     def generate_download_response(
-        self, result: dict, index_name: str, file_data: FileData
+        self, result: dict, index_name: str, file_data: ElasticsearchBatchFileData
     ) -> DownloadResponse:
         record_id = result["_id"]
         filename_id = self.get_identifier(index_name=index_name, record_id=record_id)
@@ -258,22 +255,19 @@ class ElasticsearchDownloader(Downloader):
                 exc_info=True,
             )
             raise SourceConnectionNetworkError(f"failed to download file {file_data.identifier}")
-        return DownloadResponse(
-            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()),
-                    record_locator={
-                        "hosts": self.connection_config.hosts,
-                        "index_name": index_name,
-                        "document_id": record_id,
-                    },
-                ),
-            ),
-            path=download_path,
+        cast_file_data = FileData.cast(file_data=file_data)
+        cast_file_data.identifier = filename_id
+        cast_file_data.metadata.date_processed = str(time())
+        cast_file_data.metadata.version = str(result["_version"]) if "_version" in result else None
+        cast_file_data.metadata.record_locator = {
+            "hosts": self.connection_config.hosts,
+            "index_name": index_name,
+            "document_id": record_id,
+        }
+        cast_file_data.source_identifiers = SourceIdentifiers(filename=filename, fullpath=filename)
+        return super().generate_download_response(
+            file_data=cast_file_data,
+            download_path=download_path,
         )
 
     def run(self, file_data: FileData, **kwargs: Any) -> download_responses:
@@ -286,11 +280,12 @@ class ElasticsearchDownloader(Downloader):
 
         return AsyncElasticsearch, async_scan
 
-    async def run_async(self, file_data: FileData, **kwargs: Any) -> download_responses:
+    async def run_async(self, file_data: BatchFileData, **kwargs: Any) -> download_responses:
+        elasticsearch_filedata = ElasticsearchBatchFileData.cast(file_data=file_data)
         AsyncClient, async_scan = self.load_async()
 
-        index_name: str = file_data.additional_metadata["index_name"]
-        ids: list[str] = file_data.additional_metadata["ids"]
+        index_name: str = elasticsearch_filedata.additional_metadata.index_name
+        ids: list[str] = [item.identifier for item in elasticsearch_filedata.batch_items]
 
         scan_query = {
             "_source": self.download_config.fields,
@@ -308,7 +303,7 @@ class ElasticsearchDownloader(Downloader):
             ):
                 download_responses.append(
                     self.generate_download_response(
-                        result=result, index_name=index_name, file_data=file_data
+                        result=result, index_name=index_name, file_data=elasticsearch_filedata
                     )
                 )
         return download_responses
@@ -324,7 +319,8 @@ class ElasticsearchUploadStagerConfig(UploadStagerConfig):
 class ElasticsearchUploadStager(UploadStager):
     upload_stager_config: ElasticsearchUploadStagerConfig
 
-    def conform_dict(self, data: dict, file_data: FileData) -> dict:
+    def conform_dict(self, element_dict: dict, file_data: FileData) -> dict:
+        data = element_dict.copy()
         resp = {
             "_index": self.upload_stager_config.index_name,
             "_id": get_enhanced_element_id(element_dict=data, file_data=file_data),
@@ -340,29 +336,6 @@ class ElasticsearchUploadStager(UploadStager):
             resp["_source"]["metadata"] = flatten_dict(data["metadata"], separator="-")
         return resp
 
-    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 = json.load(elements_file)
-        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, indent=2)
-        return output_path
-
 
 class ElasticsearchUploaderConfig(UploaderConfig):
     index_name: str = Field(
@@ -427,16 +400,14 @@ class ElasticsearchUploader(Uploader):
             raise WriteError(f"failed to delete records: {failures}")
 
     @requires_dependencies(["elasticsearch"], extras="elasticsearch")
-    def run(self, path: Path, file_data: FileData, **kwargs: Any) -> None:  # type: ignore
+    def run_data(self, data: list[dict], file_data: FileData, **kwargs: Any) -> None:  # noqa: E501
         from elasticsearch.helpers.errors import BulkIndexError
 
         parallel_bulk = self.load_parallel_bulk()
-        with path.open("r") as file:
-            elements_dict = json.load(file)
         upload_destination = self.connection_config.hosts or self.connection_config.cloud_id
 
         logger.info(
-            f"writing {len(elements_dict)} elements via document batches to destination "
+            f"writing {len(data)} elements via document batches to destination "
             f"index named {self.upload_config.index_name} at {upload_destination} with "
             f"batch size (in bytes) {self.upload_config.batch_size_bytes} with "
             f"{self.upload_config.num_threads} (number of) threads"
@@ -451,7 +422,7 @@ class ElasticsearchUploader(Uploader):
                     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
+                data, batch_size_limit_bytes=self.upload_config.batch_size_bytes
             ):
                 try:
                     iterator = parallel_bulk(

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

@@ -1,14 +1,14 @@
 from __future__ import annotations
 
+from contextlib import contextmanager
 from dataclasses import dataclass, field
-from pathlib import Path
 from time import time
-from typing import Any, Generator, Optional
+from typing import TYPE_CHECKING, Any, Generator, Optional
 
 from pydantic import Field, Secret
 
 from unstructured_ingest.utils.dep_check import requires_dependencies
-from unstructured_ingest.v2.interfaces import DownloadResponse, FileData, FileDataSourceMetadata
+from unstructured_ingest.v2.interfaces import FileDataSourceMetadata
 from unstructured_ingest.v2.processes.connector_registry import (
     DestinationRegistryEntry,
     SourceRegistryEntry,
@@ -25,6 +25,9 @@ from unstructured_ingest.v2.processes.connectors.fsspec.fsspec import (
 )
 from unstructured_ingest.v2.processes.connectors.fsspec.utils import json_serial, sterilize_dict
 
+if TYPE_CHECKING:
+    from adlfs import AzureBlobFileSystem
+
 CONNECTOR_TYPE = "azure"
 
 
@@ -89,6 +92,12 @@ class AzureConnectionConfig(FsspecConnectionConfig):
         }
         return access_configs
 
+    @requires_dependencies(["adlfs", "fsspec"], extras="azure")
+    @contextmanager
+    def get_client(self, protocol: str) -> Generator["AzureBlobFileSystem", None, None]:
+        with super().get_client(protocol=protocol) as client:
+            yield client
+
 
 @dataclass
 class AzureIndexer(FsspecIndexer):
@@ -96,17 +105,9 @@ class AzureIndexer(FsspecIndexer):
     index_config: AzureIndexerConfig
     connector_type: str = CONNECTOR_TYPE
 
-    @requires_dependencies(["adlfs", "fsspec"], extras="azure")
-    def precheck(self) -> None:
-        super().precheck()
-
     def sterilize_info(self, file_data: dict) -> dict:
         return sterilize_dict(data=file_data, default=azure_json_serial)
 
-    @requires_dependencies(["adlfs", "fsspec"], extras="azure")
-    def run(self, **kwargs: Any) -> Generator[FileData, None, None]:
-        return super().run(**kwargs)
-
     def get_metadata(self, file_data: dict) -> FileDataSourceMetadata:
         path = file_data["name"]
         date_created = (
@@ -149,14 +150,6 @@ class AzureDownloader(FsspecDownloader):
     connector_type: str = CONNECTOR_TYPE
     download_config: Optional[AzureDownloaderConfig] = field(default_factory=AzureDownloaderConfig)
 
-    @requires_dependencies(["adlfs", "fsspec"], extras="azure")
-    def run(self, file_data: FileData, **kwargs: Any) -> DownloadResponse:
-        return super().run(file_data=file_data, **kwargs)
-
-    @requires_dependencies(["adlfs", "fsspec"], extras="azure")
-    async def run_async(self, file_data: FileData, **kwargs: Any) -> DownloadResponse:
-        return await super().run_async(file_data=file_data, **kwargs)
-
 
 class AzureUploaderConfig(FsspecUploaderConfig):
     pass
@@ -168,22 +161,6 @@ class AzureUploader(FsspecUploader):
     connection_config: AzureConnectionConfig
     upload_config: AzureUploaderConfig = field(default=None)
 
-    @requires_dependencies(["adlfs", "fsspec"], extras="azure")
-    def __post_init__(self):
-        super().__post_init__()
-
-    @requires_dependencies(["adlfs", "fsspec"], extras="azure")
-    def precheck(self) -> None:
-        super().precheck()
-
-    @requires_dependencies(["adlfs", "fsspec"], extras="azure")
-    def run(self, path: Path, file_data: FileData, **kwargs: Any) -> None:
-        return super().run(path=path, file_data=file_data, **kwargs)
-
-    @requires_dependencies(["adlfs", "fsspec"], extras="azure")
-    async def run_async(self, path: Path, file_data: FileData, **kwargs: Any) -> None:
-        return await super().run_async(path=path, file_data=file_data, **kwargs)
-
 
 azure_source_entry = SourceRegistryEntry(
     indexer=AzureIndexer,