unstructured-ingest 0.1.1__py3-none-any.whl → 0.2.1__py3-none-any.whl

unstructured_ingest/v2/processes/connectors/astradb.py

@@ -1,31 +1,50 @@
+import copy
+import csv
+import hashlib
 import json
+import sys
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Optional
+from time import time
+from typing import TYPE_CHECKING, Any, Generator, Optional
 
 from pydantic import Field, Secret
 
 from unstructured_ingest import __name__ as integration_name
 from unstructured_ingest.__version__ import __version__ as integration_version
-from unstructured_ingest.error import DestinationConnectionError
+from unstructured_ingest.error import (
+    DestinationConnectionError,
+    SourceConnectionError,
+    SourceConnectionNetworkError,
+)
 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,
+    Downloader,
+    DownloaderConfig,
+    DownloadResponse,
     FileData,
+    FileDataSourceMetadata,
+    Indexer,
+    IndexerConfig,
     Uploader,
     UploaderConfig,
     UploadStager,
     UploadStagerConfig,
+    download_responses,
 )
 from unstructured_ingest.v2.logger import logger
 from unstructured_ingest.v2.processes.connector_registry import (
     DestinationRegistryEntry,
+    SourceRegistryEntry,
 )
 
 if TYPE_CHECKING:
+    from astrapy import AsyncCollection as AstraDBAsyncCollection
     from astrapy import Collection as AstraDBCollection
+    from astrapy import DataAPIClient as AstraDBClient
 
 
 CONNECTOR_TYPE = "astradb"
@@ -37,14 +56,253 @@ class AstraDBAccessConfig(AccessConfig):
 
 
 class AstraDBConnectionConfig(ConnectionConfig):
-    connection_type: str = Field(default=CONNECTOR_TYPE, init=False)
+    connector_type: str = Field(default=CONNECTOR_TYPE, init=False)
     access_config: Secret[AstraDBAccessConfig]
 
+    @requires_dependencies(["astrapy"], extras="astradb")
+    def get_client(self) -> "AstraDBClient":
+        from astrapy import DataAPIClient as AstraDBClient
+
+        # Create a client object to interact with the Astra DB
+        # caller_name/version for Astra DB tracking
+        return AstraDBClient(
+            caller_name=integration_name,
+            caller_version=integration_version,
+        )
+
+
+def get_astra_collection(
+    connection_config: AstraDBConnectionConfig,
+    collection_name: str,
+    keyspace: str,
+) -> "AstraDBCollection":
+    # Build the Astra DB object.
+    access_configs = connection_config.access_config.get_secret_value()
+
+    # Create a client object to interact with the Astra DB
+    # caller_name/version for Astra DB tracking
+    client = connection_config.get_client()
+
+    # Get the database object
+    astra_db = client.get_database(
+        api_endpoint=access_configs.api_endpoint,
+        token=access_configs.token,
+        keyspace=keyspace,
+    )
+
+    # Connect to the collection
+    astra_db_collection = astra_db.get_collection(name=collection_name)
+    return astra_db_collection
+
+
+async def get_async_astra_collection(
+    connection_config: AstraDBConnectionConfig,
+    collection_name: str,
+    keyspace: str,
+) -> "AstraDBAsyncCollection":
+    # Build the Astra DB object.
+    access_configs = connection_config.access_config.get_secret_value()
+
+    # Create a client object to interact with the Astra DB
+    client = connection_config.get_client()
+
+    # Get the async database object
+    async_astra_db = client.get_async_database(
+        api_endpoint=access_configs.api_endpoint,
+        token=access_configs.token,
+        keyspace=keyspace,
+    )
+
+    # Get async collection from AsyncDatabase
+    async_astra_db_collection = await async_astra_db.get_collection(name=collection_name)
+    return async_astra_db_collection
+
 
 class AstraDBUploadStagerConfig(UploadStagerConfig):
     pass
 
 
+class AstraDBIndexerConfig(IndexerConfig):
+    collection_name: str = Field(
+        description="The name of the Astra DB collection. "
+        "Note that the collection name must only include letters, "
+        "numbers, and underscores."
+    )
+    keyspace: Optional[str] = Field(default=None, description="The Astra DB connection keyspace.")
+    namespace: Optional[str] = Field(
+        default=None,
+        description="The Astra DB connection namespace.",
+        deprecated="Please use 'keyspace' instead.",
+    )
+    batch_size: int = Field(default=20, description="Number of records per batch")
+
+
+class AstraDBDownloaderConfig(DownloaderConfig):
+    fields: list[str] = field(default_factory=list)
+
+
+class AstraDBUploaderConfig(UploaderConfig):
+    collection_name: str = Field(
+        description="The name of the Astra DB collection. "
+        "Note that the collection name must only include letters, "
+        "numbers, and underscores."
+    )
+    embedding_dimension: int = Field(
+        default=384, description="The dimensionality of the embeddings"
+    )
+    keyspace: Optional[str] = Field(default=None, description="The Astra DB connection keyspace.")
+    namespace: Optional[str] = Field(
+        default=None,
+        description="The Astra DB connection namespace.",
+        deprecated="Please use 'keyspace' instead.",
+    )
+    requested_indexing_policy: Optional[dict[str, Any]] = Field(
+        default=None,
+        description="The indexing policy to use for the collection.",
+        examples=['{"deny": ["metadata"]}'],
+    )
+    batch_size: int = Field(default=20, description="Number of records per batch")
+
+
+@dataclass
+class AstraDBIndexer(Indexer):
+    connection_config: AstraDBConnectionConfig
+    index_config: AstraDBIndexerConfig
+
+    def get_collection(self) -> "AstraDBCollection":
+        return get_astra_collection(
+            connection_config=self.connection_config,
+            collection_name=self.index_config.collection_name,
+            keyspace=self.index_config.keyspace or self.index_config.namespace,
+        )
+
+    def precheck(self) -> None:
+        try:
+            self.get_collection()
+        except Exception as e:
+            logger.error(f"Failed to validate connection {e}", exc_info=True)
+            raise SourceConnectionError(f"failed to validate connection: {e}")
+
+    def _get_doc_ids(self) -> set[str]:
+        """Fetches all document ids in an index"""
+        # Initialize set of ids
+        ids = set()
+
+        # Get the collection
+        collection = self.get_collection()
+
+        # Perform the find operation to get all items
+        astra_db_docs_cursor = collection.find({}, projection={"_id": True})
+
+        # Iterate over the cursor
+        astra_db_docs = []
+        for result in astra_db_docs_cursor:
+            astra_db_docs.append(result)
+
+        # Create file data for each astra record
+        for astra_record in astra_db_docs:
+            ids.add(astra_record["_id"])
+
+        return ids
+
+    def run(self, **kwargs: Any) -> Generator[FileData, None, None]:
+        all_ids = self._get_doc_ids()
+        ids = list(all_ids)
+        id_batches = batch_generator(ids, self.index_config.batch_size)
+
+        for batch in id_batches:
+            # Make sure the hash is always a positive number to create identified
+            identified = str(hash(batch) + sys.maxsize + 1)
+            fd = FileData(
+                identifier=identified,
+                connector_type=CONNECTOR_TYPE,
+                doc_type="batch",
+                metadata=FileDataSourceMetadata(
+                    date_processed=str(time()),
+                ),
+                additional_metadata={
+                    "ids": list(batch),
+                    "collection_name": self.index_config.collection_name,
+                    "keyspace": self.index_config.keyspace or self.index_config.namespace,
+                },
+            )
+            yield fd
+
+
+@dataclass
+class AstraDBDownloader(Downloader):
+    connection_config: AstraDBConnectionConfig
+    download_config: AstraDBDownloaderConfig
+    connector_type: str = CONNECTOR_TYPE
+
+    def is_async(self) -> bool:
+        return True
+
+    def get_identifier(self, record_id: str) -> str:
+        f = f"{record_id}"
+        if self.download_config.fields:
+            f = "{}-{}".format(
+                f,
+                hashlib.sha256(",".join(self.download_config.fields).encode()).hexdigest()[:8],
+            )
+        return f
+
+    def write_astra_result_to_csv(self, astra_result: dict, download_path: str) -> None:
+        with open(download_path, "w", encoding="utf8") as f:
+            writer = csv.writer(f)
+            writer.writerow(astra_result.keys())
+            writer.writerow(astra_result.values())
+
+    def generate_download_response(self, result: dict, file_data: FileData) -> DownloadResponse:
+        record_id = result["_id"]
+        filename_id = self.get_identifier(record_id=record_id)
+        filename = f"{filename_id}.csv"  # csv to preserve column info
+        download_path = self.download_dir / Path(filename)
+        logger.debug(f"Downloading results from record {record_id} as csv to {download_path}")
+        download_path.parent.mkdir(parents=True, exist_ok=True)
+        try:
+            self.write_astra_result_to_csv(astra_result=result, download_path=download_path)
+        except Exception as e:
+            logger.error(
+                f"failed to download from record {record_id} to {download_path}: {e}",
+                exc_info=True,
+            )
+            raise SourceConnectionNetworkError(f"failed to download file {file_data.identifier}")
+
+        # modify input file_data for download_response
+        copied_file_data = copy.deepcopy(file_data)
+        copied_file_data.identifier = filename
+        copied_file_data.doc_type = "file"
+        copied_file_data.metadata.date_processed = str(time())
+        copied_file_data.metadata.record_locator = {"document_id": record_id}
+        copied_file_data.additional_metadata.pop("ids", None)
+        return super().generate_download_response(
+            file_data=copied_file_data, download_path=download_path
+        )
+
+    def run(self, file_data: FileData, **kwargs: Any) -> download_responses:
+        raise NotImplementedError("Use astradb run_async instead")
+
+    async def run_async(self, file_data: FileData, **kwargs: Any) -> download_responses:
+        # Get metadata from file_data
+        ids: list[str] = file_data.additional_metadata["ids"]
+        collection_name: str = file_data.additional_metadata["collection_name"]
+        keyspace: str = file_data.additional_metadata["keyspace"]
+
+        # Retrieve results from async collection
+        download_responses = []
+        async_astra_collection = await get_async_astra_collection(
+            connection_config=self.connection_config,
+            collection_name=collection_name,
+            keyspace=keyspace,
+        )
+        async for result in async_astra_collection.find({"_id": {"$in": ids}}):
+            download_responses.append(
+                self.generate_download_response(result=result, file_data=file_data)
+            )
+        return download_responses
+
+
 @dataclass
 class AstraDBUploadStager(UploadStager):
     upload_stager_config: AstraDBUploadStagerConfig = field(
@@ -77,29 +335,6 @@ class AstraDBUploadStager(UploadStager):
         return output_path
 
 
-class AstraDBUploaderConfig(UploaderConfig):
-    collection_name: str = Field(
-        description="The name of the Astra DB collection. "
-        "Note that the collection name must only include letters, "
-        "numbers, and underscores."
-    )
-    embedding_dimension: int = Field(
-        default=384, description="The dimensionality of the embeddings"
-    )
-    keyspace: Optional[str] = Field(default=None, description="The Astra DB connection keyspace.")
-    namespace: Optional[str] = Field(
-        default=None,
-        description="The Astra DB connection namespace.",
-        deprecated="Please use 'keyspace' instead.",
-    )
-    requested_indexing_policy: Optional[dict[str, Any]] = Field(
-        default=None,
-        description="The indexing policy to use for the collection.",
-        examples=['{"deny": ["metadata"]}'],
-    )
-    batch_size: int = Field(default=20, description="Number of records per batch")
-
-
 @dataclass
 class AstraDBUploader(Uploader):
     connection_config: AstraDBConnectionConfig
@@ -108,43 +343,23 @@ class AstraDBUploader(Uploader):
 
     def precheck(self) -> None:
         try:
-            self.get_collection()
+            get_astra_collection(
+                connection_config=self.connection_config,
+                collection_name=self.upload_config.collection_name,
+                keyspace=self.upload_config.keyspace or self.upload_config.namespace,
+            )
         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":
-        from astrapy import DataAPIClient as AstraDBClient
-
-        # Choose keyspace or deprecated namespace
-        keyspace_param = self.upload_config.keyspace or self.upload_config.namespace
-
-        # Get the collection_name
-        collection_name = self.upload_config.collection_name
-
-        # Build the Astra DB object.
-        access_configs = self.connection_config.access_config.get_secret_value()
-
-        # Create a client object to interact with the Astra DB
-        # caller_name/version for Astra DB tracking
-        my_client = AstraDBClient(
-            caller_name=integration_name,
-            caller_version=integration_version,
-        )
-
-        # Get the database object
-        astra_db = my_client.get_database(
-            api_endpoint=access_configs.api_endpoint,
-            token=access_configs.token,
-            keyspace=keyspace_param,
+        return get_astra_collection(
+            connection_config=self.connection_config,
+            collection_name=self.upload_config.collection_name,
+            keyspace=self.upload_config.keyspace or self.upload_config.namespace,
         )
 
-        # Connect to the newly created collection
-        astra_db_collection = astra_db.get_collection(name=collection_name)
-
-        return astra_db_collection
-
     def run(self, path: Path, file_data: FileData, **kwargs: Any) -> None:
         with path.open("r") as file:
             elements_dict = json.load(file)
@@ -160,6 +375,14 @@ class AstraDBUploader(Uploader):
             collection.insert_many(chunk)
 
 
+astra_db_source_entry = SourceRegistryEntry(
+    indexer=AstraDBIndexer,
+    indexer_config=AstraDBIndexerConfig,
+    downloader=AstraDBDownloader,
+    downloader_config=AstraDBDownloaderConfig,
+    connection_config=AstraDBConnectionConfig,
+)
+
 astra_db_destination_entry = DestinationRegistryEntry(
     connection_config=AstraDBConnectionConfig,
     upload_stager_config=AstraDBUploadStagerConfig,

unstructured_ingest/v2/processes/connectors/databricks/volumes.py

@@ -166,7 +166,9 @@ class DatabricksVolumesUploader(Uploader, ABC):
             raise DestinationConnectionError(f"failed to validate connection: {e}")
 
     def run(self, path: Path, file_data: FileData, **kwargs: Any) -> None:
-        output_path = os.path.join(self.upload_config.path, path.name)
+        output_path = os.path.join(
+            self.upload_config.path, f"{file_data.source_identifiers.filename}.json"
+        )
         with open(path, "rb") as elements_file:
             self.connection_config.get_client().files.upload(
                 file_path=output_path,

unstructured_ingest/v2/processes/connectors/delta_table.py

@@ -0,0 +1,185 @@
+import json
+import os
+from dataclasses import dataclass, field
+from multiprocessing import Process
+from pathlib import Path
+from typing import Any, Optional
+
+import pandas as pd
+from pydantic import Field, Secret
+
+from unstructured_ingest.error import DestinationConnectionError
+from unstructured_ingest.utils.dep_check import requires_dependencies
+from unstructured_ingest.utils.table import convert_to_pandas_dataframe
+from unstructured_ingest.v2.interfaces import (
+    AccessConfig,
+    ConnectionConfig,
+    FileData,
+    Uploader,
+    UploaderConfig,
+    UploadStager,
+    UploadStagerConfig,
+)
+from unstructured_ingest.v2.logger import logger
+from unstructured_ingest.v2.processes.connector_registry import DestinationRegistryEntry
+
+CONNECTOR_TYPE = "delta_table"
+
+
+class DeltaTableAccessConfig(AccessConfig):
+    aws_access_key_id: Optional[str] = Field(default=None, description="AWS Access Key Id")
+    aws_secret_access_key: Optional[str] = Field(default=None, description="AWS Secret Access Key")
+
+
+class DeltaTableConnectionConfig(ConnectionConfig):
+    access_config: Secret[DeltaTableAccessConfig] = Field(
+        default=DeltaTableAccessConfig(), validate_default=True
+    )
+    aws_region: Optional[str] = Field(default=None, description="AWS Region")
+    table_uri: str = Field(
+        default=None,
+        description=(
+            "Local path or path to the target folder in the S3 bucket, "
+            "formatted as s3://my-bucket/my-folder/"
+        ),
+    )
+
+    def update_storage_options(self, storage_options: dict) -> None:
+        secrets = self.access_config.get_secret_value()
+        if self.aws_region and secrets.aws_access_key_id and secrets.aws_secret_access_key:
+            storage_options["AWS_REGION"] = self.aws_region
+            storage_options["AWS_ACCESS_KEY_ID"] = secrets.aws_access_key_id
+            storage_options["AWS_SECRET_ACCESS_KEY"] = secrets.aws_secret_access_key
+            # Delta-rs doesn't support concurrent S3 writes without external locks (DynamoDB).
+            # This flag allows single-writer uploads to S3 without using locks, according to:
+            # https://delta-io.github.io/delta-rs/usage/writing/writing-to-s3-with-locking-provider/
+            storage_options["AWS_S3_ALLOW_UNSAFE_RENAME"] = "true"
+
+
+class DeltaTableUploadStagerConfig(UploadStagerConfig):
+    pass
+
+
+@dataclass
+class DeltaTableUploadStager(UploadStager):
+    upload_stager_config: DeltaTableUploadStagerConfig = field(
+        default_factory=lambda: DeltaTableUploadStagerConfig()
+    )
+
+    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_path = Path(output_dir) / Path(f"{output_filename}.parquet")
+
+        df = convert_to_pandas_dataframe(elements_dict=elements_contents)
+        df.to_parquet(output_path)
+
+        return output_path
+
+
+class DeltaTableUploaderConfig(UploaderConfig):
+    pass
+
+
+@dataclass
+class DeltaTableUploader(Uploader):
+    upload_config: DeltaTableUploaderConfig
+    connection_config: DeltaTableConnectionConfig
+    connector_type: str = CONNECTOR_TYPE
+
+    @requires_dependencies(["s3fs", "fsspec"], extras="s3")
+    def precheck(self):
+        secrets = self.connection_config.access_config.get_secret_value()
+        if (
+            self.connection_config.aws_region
+            and secrets.aws_access_key_id
+            and secrets.aws_secret_access_key
+        ):
+            from fsspec import get_filesystem_class
+
+            try:
+                fs = get_filesystem_class("s3")(
+                    key=secrets.aws_access_key_id, secret=secrets.aws_secret_access_key
+                )
+                fs.write_bytes(path=self.connection_config.table_uri, value=b"")
+
+            except Exception as e:
+                logger.error(f"failed to validate connection: {e}", exc_info=True)
+                raise DestinationConnectionError(f"failed to validate connection: {e}")
+
+    def process_csv(self, csv_paths: list[Path]) -> pd.DataFrame:
+        logger.debug(f"uploading content from {len(csv_paths)} csv files")
+        df = pd.concat((pd.read_csv(path) for path in csv_paths), ignore_index=True)
+        return df
+
+    def process_json(self, json_paths: list[Path]) -> pd.DataFrame:
+        logger.debug(f"uploading content from {len(json_paths)} json files")
+        all_records = []
+        for p in json_paths:
+            with open(p) as json_file:
+                all_records.extend(json.load(json_file))
+
+        return pd.DataFrame(data=all_records)
+
+    def process_parquet(self, parquet_paths: list[Path]) -> pd.DataFrame:
+        logger.debug(f"uploading content from {len(parquet_paths)} parquet files")
+        df = pd.concat((pd.read_parquet(path) for path in parquet_paths), ignore_index=True)
+        return df
+
+    def read_dataframe(self, path: Path) -> pd.DataFrame:
+        if path.suffix == ".csv":
+            return self.process_csv(csv_paths=[path])
+        elif path.suffix == ".json":
+            return self.process_json(json_paths=[path])
+        elif path.suffix == ".parquet":
+            return self.process_parquet(parquet_paths=[path])
+        else:
+            raise ValueError(f"Unsupported file type, must be parquet, json or csv file: {path}")
+
+    @requires_dependencies(["deltalake"], extras="delta-table")
+    def run(self, path: Path, file_data: FileData, **kwargs: Any) -> None:
+        from deltalake.writer import write_deltalake
+
+        df = self.read_dataframe(path)
+        updated_upload_path = os.path.join(
+            self.connection_config.table_uri, file_data.source_identifiers.relative_path
+        )
+        logger.info(
+            f"writing {len(df)} rows to destination table "
+            f"at {updated_upload_path}\ndtypes: {df.dtypes}",
+        )
+        storage_options = {}
+        self.connection_config.update_storage_options(storage_options=storage_options)
+
+        writer_kwargs = {
+            "table_or_uri": updated_upload_path,
+            "data": df,
+            "mode": "overwrite",
+            "storage_options": storage_options,
+        }
+        # NOTE: deltalake writer on Linux sometimes can finish but still trigger a SIGABRT and cause
+        # ingest to fail, even though all tasks are completed normally. Putting the writer into a
+        # process mitigates this issue by ensuring python interpreter waits properly for deltalake's
+        # rust backend to finish
+        writer = Process(
+            target=write_deltalake,
+            kwargs=writer_kwargs,
+        )
+        writer.start()
+        writer.join()
+
+
+delta_table_destination_entry = DestinationRegistryEntry(
+    connection_config=DeltaTableConnectionConfig,
+    uploader=DeltaTableUploader,
+    uploader_config=DeltaTableUploaderConfig,
+    upload_stager=DeltaTableUploadStager,
+    upload_stager_config=DeltaTableUploadStagerConfig,
+)

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

@@ -176,6 +176,7 @@ class FsspecIndexer(Indexer):
                 ),
                 metadata=self.get_metadata(file_data=file_data),
                 additional_metadata=additional_metadata,
+                display_name=file_path,
             )