unstructured-ingest 0.3.13__py3-none-any.whl

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

@@ -0,0 +1,102 @@
+from dataclasses import dataclass, field
+from typing import Optional
+
+from pydantic import Field, Secret
+
+from unstructured_ingest.v2.processes.connector_registry import (
+    DestinationRegistryEntry,
+    SourceRegistryEntry,
+)
+from unstructured_ingest.v2.processes.connectors.databricks.volumes import (
+    DatabricksVolumesAccessConfig,
+    DatabricksVolumesConnectionConfig,
+    DatabricksVolumesDownloader,
+    DatabricksVolumesDownloaderConfig,
+    DatabricksVolumesIndexer,
+    DatabricksVolumesIndexerConfig,
+    DatabricksVolumesUploader,
+    DatabricksVolumesUploaderConfig,
+)
+
+CONNECTOR_TYPE = "databricks_volumes_azure"
+
+
+class DatabricksAzureVolumesAccessConfig(DatabricksVolumesAccessConfig):
+    account_id: Optional[str] = Field(
+        default=None,
+        description="The Databricks account ID for the Databricks " "accounts endpoint.",
+    )
+    profile: 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"],
+    )
+
+
+class DatabricksAzureVolumesConnectionConfig(DatabricksVolumesConnectionConfig):
+    access_config: Secret[DatabricksAzureVolumesAccessConfig]
+
+
+class DatabricksAzureVolumesIndexerConfig(DatabricksVolumesIndexerConfig):
+    pass
+
+
+@dataclass
+class DatabricksAzureVolumesIndexer(DatabricksVolumesIndexer):
+    connection_config: DatabricksAzureVolumesConnectionConfig
+    index_config: DatabricksAzureVolumesIndexerConfig
+    connector_type: str = CONNECTOR_TYPE
+
+
+class DatabricksAzureVolumesDownloaderConfig(DatabricksVolumesDownloaderConfig):
+    pass
+
+
+@dataclass
+class DatabricksAzureVolumesDownloader(DatabricksVolumesDownloader):
+    connection_config: DatabricksAzureVolumesConnectionConfig
+    download_config: DatabricksVolumesDownloaderConfig
+    connector_type: str = CONNECTOR_TYPE
+
+
+class DatabricksAzureVolumesUploaderConfig(DatabricksVolumesUploaderConfig):
+    pass
+
+
+@dataclass
+class DatabricksAzureVolumesUploader(DatabricksVolumesUploader):
+    connection_config: DatabricksAzureVolumesConnectionConfig
+    upload_config: DatabricksAzureVolumesUploaderConfig = field(
+        default_factory=DatabricksAzureVolumesUploaderConfig
+    )
+    connector_type: str = CONNECTOR_TYPE
+
+
+databricks_azure_volumes_destination_entry = DestinationRegistryEntry(
+    connection_config=DatabricksAzureVolumesConnectionConfig,
+    uploader=DatabricksAzureVolumesUploader,
+    uploader_config=DatabricksAzureVolumesUploaderConfig,
+)
+
+databricks_azure_volumes_source_entry = SourceRegistryEntry(
+    connection_config=DatabricksAzureVolumesConnectionConfig,
+    indexer=DatabricksAzureVolumesIndexer,
+    indexer_config=DatabricksAzureVolumesIndexerConfig,
+    downloader=DatabricksAzureVolumesDownloader,
+    downloader_config=DatabricksAzureVolumesDownloaderConfig,
+)

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

@@ -0,0 +1,85 @@
+from dataclasses import dataclass, field
+from typing import Optional
+
+from pydantic import Field, Secret
+
+from unstructured_ingest.v2.processes.connector_registry import (
+    DestinationRegistryEntry,
+    SourceRegistryEntry,
+)
+from unstructured_ingest.v2.processes.connectors.databricks.volumes import (
+    DatabricksVolumesAccessConfig,
+    DatabricksVolumesConnectionConfig,
+    DatabricksVolumesDownloader,
+    DatabricksVolumesDownloaderConfig,
+    DatabricksVolumesIndexer,
+    DatabricksVolumesIndexerConfig,
+    DatabricksVolumesUploader,
+    DatabricksVolumesUploaderConfig,
+)
+
+CONNECTOR_TYPE = "databricks_volumes_gcp"
+
+
+class DatabricksGoogleVolumesAccessConfig(DatabricksVolumesAccessConfig):
+    account_id: Optional[str] = Field(
+        default=None,
+        description="The Databricks account ID for the Databricks " "accounts endpoint.",
+    )
+    profile: Optional[str] = None
+    google_credentials: Optional[str] = None
+    google_service_account: Optional[str] = None
+
+
+class DatabricksGoogleVolumesConnectionConfig(DatabricksVolumesConnectionConfig):
+    access_config: Secret[DatabricksGoogleVolumesAccessConfig]
+
+
+class DatabricksGoogleVolumesIndexerConfig(DatabricksVolumesIndexerConfig):
+    pass
+
+
+@dataclass
+class DatabricksGoogleVolumesIndexer(DatabricksVolumesIndexer):
+    connection_config: DatabricksGoogleVolumesConnectionConfig
+    index_config: DatabricksGoogleVolumesIndexerConfig
+    connector_type: str = CONNECTOR_TYPE
+
+
+class DatabricksGoogleVolumesDownloaderConfig(DatabricksVolumesDownloaderConfig):
+    pass
+
+
+@dataclass
+class DatabricksGoogleVolumesDownloader(DatabricksVolumesDownloader):
+    connection_config: DatabricksGoogleVolumesConnectionConfig
+    download_config: DatabricksVolumesDownloaderConfig
+    connector_type: str = CONNECTOR_TYPE
+
+
+class DatabricksGoogleVolumesUploaderConfig(DatabricksVolumesUploaderConfig):
+    pass
+
+
+@dataclass
+class DatabricksGoogleVolumesUploader(DatabricksVolumesUploader):
+    connection_config: DatabricksGoogleVolumesConnectionConfig
+    upload_config: DatabricksGoogleVolumesUploaderConfig = field(
+        default_factory=DatabricksGoogleVolumesUploaderConfig
+    )
+    connector_type: str = CONNECTOR_TYPE
+
+
+databricks_gcp_volumes_destination_entry = DestinationRegistryEntry(
+    connection_config=DatabricksGoogleVolumesConnectionConfig,
+    uploader=DatabricksGoogleVolumesUploader,
+    uploader_config=DatabricksGoogleVolumesUploaderConfig,
+)
+
+databricks_gcp_volumes_source_entry = SourceRegistryEntry(
+    connection_config=DatabricksGoogleVolumesConnectionConfig,
+    indexer=DatabricksGoogleVolumesIndexer,
+    indexer_config=DatabricksGoogleVolumesIndexerConfig,
+    downloader=DatabricksGoogleVolumesDownloader,
+    downloader_config=DatabricksGoogleVolumesDownloaderConfig,
+)

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

@@ -0,0 +1,86 @@
+from dataclasses import dataclass
+from typing import Optional
+
+from pydantic import Field, Secret
+
+from unstructured_ingest.v2.processes.connector_registry import (
+    DestinationRegistryEntry,
+    SourceRegistryEntry,
+)
+from unstructured_ingest.v2.processes.connectors.databricks.volumes import (
+    DatabricksVolumesAccessConfig,
+    DatabricksVolumesConnectionConfig,
+    DatabricksVolumesDownloader,
+    DatabricksVolumesDownloaderConfig,
+    DatabricksVolumesIndexer,
+    DatabricksVolumesIndexerConfig,
+    DatabricksVolumesUploader,
+    DatabricksVolumesUploaderConfig,
+)
+
+CONNECTOR_TYPE = "databricks_volumes"
+
+
+class DatabricksNativeVolumesAccessConfig(DatabricksVolumesAccessConfig):
+    client_id: Optional[str] = Field(default=None, description="Client ID of the OAuth app.")
+    client_secret: Optional[str] = Field(
+        default=None, description="Client Secret of the OAuth app."
+    )
+    profile: 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.",
+    )
+
+
+class DatabricksNativeVolumesConnectionConfig(DatabricksVolumesConnectionConfig):
+    access_config: Secret[DatabricksNativeVolumesAccessConfig]
+
+
+class DatabricksNativeVolumesIndexerConfig(DatabricksVolumesIndexerConfig):
+    pass
+
+
+@dataclass
+class DatabricksNativeVolumesIndexer(DatabricksVolumesIndexer):
+    connection_config: DatabricksNativeVolumesConnectionConfig
+    index_config: DatabricksNativeVolumesIndexerConfig
+    connector_type: str = CONNECTOR_TYPE
+
+
+class DatabricksNativeVolumesDownloaderConfig(DatabricksVolumesDownloaderConfig):
+    pass
+
+
+@dataclass
+class DatabricksNativeVolumesDownloader(DatabricksVolumesDownloader):
+    connection_config: DatabricksNativeVolumesConnectionConfig
+    download_config: DatabricksVolumesDownloaderConfig
+    connector_type: str = CONNECTOR_TYPE
+
+
+class DatabricksNativeVolumesUploaderConfig(DatabricksVolumesUploaderConfig):
+    pass
+
+
+@dataclass
+class DatabricksNativeVolumesUploader(DatabricksVolumesUploader):
+    connection_config: DatabricksNativeVolumesConnectionConfig
+    upload_config: DatabricksNativeVolumesUploaderConfig
+    connector_type: str = CONNECTOR_TYPE
+
+
+databricks_native_volumes_destination_entry = DestinationRegistryEntry(
+    connection_config=DatabricksNativeVolumesConnectionConfig,
+    uploader=DatabricksNativeVolumesUploader,
+    uploader_config=DatabricksNativeVolumesUploaderConfig,
+)
+
+databricks_native_volumes_source_entry = SourceRegistryEntry(
+    connection_config=DatabricksNativeVolumesConnectionConfig,
+    indexer=DatabricksNativeVolumesIndexer,
+    indexer_config=DatabricksNativeVolumesIndexerConfig,
+    downloader=DatabricksNativeVolumesDownloader,
+    downloader_config=DatabricksNativeVolumesDownloaderConfig,
+)

unstructured_ingest/v2/processes/connectors/delta_table.py

@@ -0,0 +1,191 @@
+import json
+import os
+import traceback
+from dataclasses import dataclass, field
+from multiprocessing import Process, Queue
+from pathlib import Path
+from typing import Any, Optional
+from urllib.parse import urlparse
+
+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.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"
+
+
+@requires_dependencies(["deltalake"], extras="delta-table")
+def write_deltalake_with_error_handling(queue, **kwargs):
+    from deltalake.writer import write_deltalake
+
+    try:
+        write_deltalake(**kwargs)
+    except Exception:
+        queue.put(traceback.format_exc())
+
+
+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(["boto3"], extras="delta-table")
+    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 boto3 import client
+
+            url = urlparse(self.connection_config.table_uri)
+            bucket_name = url.netloc
+            dir_path = url.path.lstrip("/")
+
+            try:
+                s3_client = client(
+                    "s3",
+                    aws_access_key_id=secrets.aws_access_key_id,
+                    aws_secret_access_key=secrets.aws_secret_access_key,
+                )
+                s3_client.put_object(Bucket=bucket_name, Key=dir_path, Body=b"")
+
+                response = s3_client.get_bucket_location(Bucket=bucket_name)
+
+                if self.connection_config.aws_region != response.get("LocationConstraint"):
+                    raise ValueError("Wrong AWS Region was provided.")
+
+            except Exception as e:
+                logger.error(f"failed to validate connection: {e}", exc_info=True)
+                raise DestinationConnectionError(f"failed to validate connection: {e}")
+
+    def upload_dataframe(self, df: pd.DataFrame, file_data: FileData) -> None:
+        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,
+        }
+        queue = Queue()
+        # 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_with_error_handling,
+            kwargs={"queue": queue, **writer_kwargs},
+        )
+        writer.start()
+        writer.join()
+
+        # Check if the queue has any exception message
+        if not queue.empty():
+            error_message = queue.get()
+            logger.error(f"Exception occurred in write_deltalake: {error_message}")
+            raise RuntimeError(f"Error in write_deltalake: {error_message}")
+
+    def run_data(self, data: list[dict], file_data: FileData, **kwargs: Any) -> None:
+        df = pd.DataFrame(data=data)
+        self.upload_dataframe(df=df, file_data=file_data)
+
+    def run(self, path: Path, file_data: FileData, **kwargs: Any) -> None:
+        df = get_data_df(path)
+        self.upload_dataframe(df=df, file_data=file_data)
+
+
+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/discord.py

@@ -0,0 +1,158 @@
+import datetime as dt
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Generator, Optional
+
+from pydantic import Field, Secret
+
+from unstructured_ingest.error import SourceConnectionError
+from unstructured_ingest.utils.dep_check import requires_dependencies
+from unstructured_ingest.v2.interfaces import (
+    AccessConfig,
+    ConnectionConfig,
+    Downloader,
+    DownloaderConfig,
+    DownloadResponse,
+    FileData,
+    FileDataSourceMetadata,
+    Indexer,
+    IndexerConfig,
+    SourceIdentifiers,
+)
+from unstructured_ingest.v2.logger import logger
+from unstructured_ingest.v2.processes.connector_registry import SourceRegistryEntry
+
+if TYPE_CHECKING:
+    from discord import Client as DiscordClient
+
+CONNECTOR_TYPE = "discord"
+
+
+class DiscordAccessConfig(AccessConfig):
+    token: str = Field(description="Discord API token")
+
+
+class DiscordConnectionConfig(ConnectionConfig):
+    access_config: Secret[DiscordAccessConfig] = Field(
+        default=DiscordAccessConfig, validate_default=True
+    )
+
+    @requires_dependencies(["discord"], extras="discord")
+    def get_client(self) -> "DiscordClient":
+        import discord
+
+        intents = discord.Intents.default()
+        intents.message_content = True
+        return discord.Client(intents=intents)
+
+
+class DiscordIndexerConfig(IndexerConfig):
+    channels: list[str] = Field(
+        default=None,
+        description="List of Discord channel IDs to process",
+    )
+
+
+@dataclass
+class DiscordIndexer(Indexer):
+    connection_config: DiscordConnectionConfig
+    index_config: DiscordIndexerConfig
+
+    def run(self, **kwargs: Any) -> Generator[FileData, None, None]:
+        self.connection_config.get_client()
+        channels_to_process: set[str] = set(self.index_config.channels or [])
+
+        for channel_id in list(channels_to_process):
+            file_data = self.get_channel_file_data(channel_id=channel_id)
+            if file_data:
+                yield file_data
+
+    def precheck(self) -> None:
+        if not self.connection_config.access_config.get_secret_value().token:
+            raise SourceConnectionError("Discord token is missing")
+        if not self.index_config.channels:
+            raise SourceConnectionError("No channels provided")
+
+    def get_channel_file_data(self, channel_id: str) -> Optional[FileData]:
+        # Fetch channel metadata
+        identifier = channel_id
+        channel_id = f"{channel_id}.txt"
+        source_identifiers = SourceIdentifiers(
+            filename=channel_id,
+            fullpath=channel_id,
+        )
+        metadata = FileDataSourceMetadata(
+            record_locator={"channel_id": identifier},
+            date_processed=str(dt.datetime.utcnow().isoformat()),
+        )
+        return FileData(
+            identifier=identifier,
+            connector_type=CONNECTOR_TYPE,
+            source_identifiers=source_identifiers,
+            metadata=metadata,
+        )
+
+
+class DiscordDownloaderConfig(DownloaderConfig):
+    limit: Optional[int] = Field(
+        default=100, description="Limit on how many messages per channel to pull in"
+    )
+
+
+@dataclass
+class DiscordDownloader(Downloader):
+    connection_config: DiscordConnectionConfig
+    download_config: DiscordDownloaderConfig
+    connector_type: str = CONNECTOR_TYPE
+
+    def is_async(self) -> bool:
+        return True
+
+    def run(self, file_data: FileData, **kwargs: Any) -> DownloadResponse:
+        # Synchronous run is not implemented
+        raise NotImplementedError()
+
+    async def run_async(self, file_data: FileData, **kwargs: Any) -> DownloadResponse:
+        record_locator = file_data.metadata.record_locator
+
+        if "channel_id" not in record_locator:
+            raise ValueError(f"No channel id in file data record locator: {record_locator}")
+
+        client = self.connection_config.get_client()
+        download_path = self.get_download_path(file_data=file_data)
+        download_path.parent.mkdir(parents=True, exist_ok=True)
+
+        messages = []
+        channel_id = record_locator["channel_id"]
+
+        @client.event
+        async def on_ready():
+            logger.debug("Discord Bot is ready")
+            channel = client.get_channel(int(channel_id))
+            if not channel:
+                raise ValueError(f"channel not found for id: {channel_id}")
+            logger.debug(f"Processing messages for channel: {channel.name}")
+            async for msg in channel.history(limit=self.download_config.limit):
+                messages.append(msg)
+            logger.debug(f"Fetched {len(messages)} messages")
+            await client.close()
+
+        try:
+            await client.start(self.connection_config.access_config.get_secret_value().token)
+        finally:
+            await client.close()
+
+        content = "\n".join([message.content for message in messages])
+
+        with open(download_path, "w") as file:
+            file.write(content)
+
+        return self.generate_download_response(file_data=file_data, download_path=download_path)
+
+
+discord_source_entry = SourceRegistryEntry(
+    indexer=DiscordIndexer,
+    indexer_config=DiscordIndexerConfig,
+    downloader=DiscordDownloader,
+    downloader_config=DiscordDownloaderConfig,
+    connection_config=DiscordConnectionConfig,
+)

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

@@ -0,0 +1,15 @@
+from __future__ import annotations
+
+from unstructured_ingest.v2.processes.connector_registry import (
+    add_destination_entry,
+)
+
+from .duckdb import CONNECTOR_TYPE as DUCKDB_CONNECTOR_TYPE
+from .duckdb import duckdb_destination_entry
+from .motherduck import CONNECTOR_TYPE as MOTHERDUCK_CONNECTOR_TYPE
+from .motherduck import motherduck_destination_entry
+
+add_destination_entry(destination_type=DUCKDB_CONNECTOR_TYPE, entry=duckdb_destination_entry)
+add_destination_entry(
+    destination_type=MOTHERDUCK_CONNECTOR_TYPE, entry=motherduck_destination_entry
+)