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

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

@@ -0,0 +1,175 @@
+import os
+from abc import ABC
+from dataclasses import dataclass
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Generator, Optional
+from uuid import NAMESPACE_DNS, uuid5
+
+from pydantic import BaseModel, Field
+
+from unstructured_ingest.error import (
+    DestinationConnectionError,
+    SourceConnectionError,
+    SourceConnectionNetworkError,
+)
+from unstructured_ingest.utils.dep_check import requires_dependencies
+from unstructured_ingest.v2.interfaces import (
+    ConnectionConfig,
+    Downloader,
+    DownloaderConfig,
+    DownloadResponse,
+    FileData,
+    FileDataSourceMetadata,
+    Indexer,
+    IndexerConfig,
+    SourceIdentifiers,
+    Uploader,
+    UploaderConfig,
+)
+from unstructured_ingest.v2.logger import logger
+
+if TYPE_CHECKING:
+    from databricks.sdk import WorkspaceClient
+
+
+class DatabricksPathMixin(BaseModel):
+    volume: str = Field(description="Name of volume in the Unity Catalog")
+    catalog: str = Field(description="Name of the catalog in the Databricks Unity Catalog service")
+    volume_path: Optional[str] = Field(
+        default=None, description="Optional path within the volume to write to"
+    )
+    databricks_schema: str = Field(
+        default="default",
+        alias="schema",
+        description="Schema associated with the volume to write to in the Unity Catalog service",
+    )
+
+    @property
+    def path(self) -> str:
+        path = f"/Volumes/{self.catalog}/{self.databricks_schema}/{self.volume}"
+        if self.volume_path:
+            path = f"{path}/{self.volume_path}"
+        return path
+
+
+class DatabricksVolumesConnectionConfig(ConnectionConfig, ABC):
+    host: Optional[str] = Field(
+        default=None,
+        description="The Databricks host URL for either the "
+        "Databricks workspace endpoint or the "
+        "Databricks accounts endpoint.",
+    )
+
+    @requires_dependencies(dependencies=["databricks.sdk"], extras="databricks-volumes")
+    def get_client(self) -> "WorkspaceClient":
+        from databricks.sdk import WorkspaceClient
+
+        return WorkspaceClient(
+            host=self.host,
+            **self.access_config.get_secret_value().model_dump(),
+        )
+
+
+class DatabricksVolumesIndexerConfig(IndexerConfig, DatabricksPathMixin):
+    recursive: bool = False
+
+
+@dataclass
+class DatabricksVolumesIndexer(Indexer, ABC):
+    index_config: DatabricksVolumesIndexerConfig
+    connection_config: DatabricksVolumesConnectionConfig
+
+    def precheck(self) -> None:
+        try:
+            self.connection_config.get_client()
+        except Exception as e:
+            logger.error(f"failed to validate connection: {e}", exc_info=True)
+            raise SourceConnectionError(f"failed to validate connection: {e}")
+
+    def run(self, **kwargs: Any) -> Generator[FileData, None, None]:
+        for file_info in self.connection_config.get_client().dbfs.list(
+            path=self.index_config.path, recursive=self.index_config.recursive
+        ):
+            if file_info.is_dir:
+                continue
+            rel_path = file_info.path.replace(self.index_config.path, "")
+            if rel_path.startswith("/"):
+                rel_path = rel_path[1:]
+            filename = Path(file_info.path).name
+            yield FileData(
+                identifier=str(uuid5(NAMESPACE_DNS, file_info.path)),
+                connector_type=self.connector_type,
+                source_identifiers=SourceIdentifiers(
+                    filename=filename,
+                    rel_path=rel_path,
+                    fullpath=file_info.path,
+                ),
+                additional_metadata={"catalog": self.index_config.catalog, "path": file_info.path},
+                metadata=FileDataSourceMetadata(
+                    url=file_info.path, date_modified=str(file_info.modification_time)
+                ),
+            )
+
+
+class DatabricksVolumesDownloaderConfig(DownloaderConfig):
+    pass
+
+
+@dataclass
+class DatabricksVolumesDownloader(Downloader, ABC):
+    download_config: DatabricksVolumesDownloaderConfig
+    connection_config: DatabricksVolumesConnectionConfig
+
+    def precheck(self) -> None:
+        try:
+            self.connection_config.get_client()
+        except Exception as e:
+            logger.error(f"failed to validate connection: {e}", exc_info=True)
+            raise SourceConnectionError(f"failed to validate connection: {e}")
+
+    def get_download_path(self, file_data: FileData) -> Path:
+        return self.download_config.download_dir / Path(file_data.source_identifiers.relative_path)
+
+    def run(self, file_data: FileData, **kwargs: Any) -> DownloadResponse:
+        download_path = self.get_download_path(file_data=file_data)
+        download_path.parent.mkdir(parents=True, exist_ok=True)
+        volumes_path = file_data.additional_metadata["path"]
+        logger.info(f"Writing {file_data.identifier} to {download_path}")
+        try:
+            with self.connection_config.get_client().dbfs.download(path=volumes_path) as c:
+                read_content = c._read_handle.read()
+            with open(download_path, "wb") as f:
+                f.write(read_content)
+        except Exception as e:
+            logger.error(f"failed to download file {file_data.identifier}: {e}", exc_info=True)
+            raise SourceConnectionNetworkError(f"failed to download file {file_data.identifier}")
+
+        return self.generate_download_response(file_data=file_data, download_path=download_path)
+
+
+class DatabricksVolumesUploaderConfig(UploaderConfig, DatabricksPathMixin):
+    overwrite: bool = Field(
+        default=False, description="If true, an existing file will be overwritten."
+    )
+
+
+@dataclass
+class DatabricksVolumesUploader(Uploader, ABC):
+    upload_config: DatabricksVolumesUploaderConfig
+    connection_config: DatabricksVolumesConnectionConfig
+
+    def precheck(self) -> None:
+        try:
+            assert self.connection_config.get_client().current_user.me().active
+        except Exception as e:
+            logger.error(f"failed to validate connection: {e}", exc_info=True)
+            raise DestinationConnectionError(f"failed to validate connection: {e}")
+
+    def run(self, path: Path, file_data: FileData, **kwargs: Any) -> None:
+        output_path = os.path.join(self.upload_config.path, path.name)
+        with open(path, "rb") as elements_file:
+            self.connection_config.get_client().files.upload(
+                file_path=output_path,
+                contents=elements_file,
+                overwrite=self.upload_config.overwrite,
+            )

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

@@ -0,0 +1,87 @@
+from dataclasses import dataclass, field
+from typing import Optional
+
+from pydantic import Field, Secret
+
+from unstructured_ingest.v2.interfaces import AccessConfig
+from unstructured_ingest.v2.processes.connector_registry import (
+    DestinationRegistryEntry,
+    SourceRegistryEntry,
+)
+from unstructured_ingest.v2.processes.connectors.databricks.volumes import (
+    DatabricksVolumesConnectionConfig,
+    DatabricksVolumesDownloader,
+    DatabricksVolumesDownloaderConfig,
+    DatabricksVolumesIndexer,
+    DatabricksVolumesIndexerConfig,
+    DatabricksVolumesUploader,
+    DatabricksVolumesUploaderConfig,
+)
+
+CONNECTOR_TYPE = "databricks_volumes_aws"
+
+
+class DatabricksAWSVolumesAccessConfig(AccessConfig):
+    account_id: Optional[str] = Field(
+        default=None,
+        description="The Databricks account ID for the Databricks " "accounts endpoint",
+    )
+    profile: Optional[str] = None
+    token: Optional[str] = Field(
+        default=None,
+        description="The Databricks personal access token (PAT)",
+    )
+
+
+class DatabricksAWSVolumesConnectionConfig(DatabricksVolumesConnectionConfig):
+    access_config: Secret[DatabricksAWSVolumesAccessConfig]
+
+
+class DatabricksAWSVolumesIndexerConfig(DatabricksVolumesIndexerConfig):
+    pass
+
+
+@dataclass
+class DatabricksAWSVolumesIndexer(DatabricksVolumesIndexer):
+    connection_config: DatabricksAWSVolumesConnectionConfig
+    index_config: DatabricksAWSVolumesIndexerConfig
+    connector_type: str = CONNECTOR_TYPE
+
+
+class DatabricksAWSVolumesDownloaderConfig(DatabricksVolumesDownloaderConfig):
+    pass
+
+
+@dataclass
+class DatabricksAWSVolumesDownloader(DatabricksVolumesDownloader):
+    connection_config: DatabricksAWSVolumesConnectionConfig
+    download_config: DatabricksVolumesDownloaderConfig
+    connector_type: str = CONNECTOR_TYPE
+
+
+class DatabricksAWSVolumesUploaderConfig(DatabricksVolumesUploaderConfig):
+    pass
+
+
+@dataclass
+class DatabricksAWSVolumesUploader(DatabricksVolumesUploader):
+    connection_config: DatabricksAWSVolumesConnectionConfig
+    upload_config: DatabricksAWSVolumesUploaderConfig = field(
+        default_factory=DatabricksAWSVolumesUploaderConfig
+    )
+    connector_type: str = CONNECTOR_TYPE
+
+
+databricks_aws_volumes_destination_entry = DestinationRegistryEntry(
+    connection_config=DatabricksAWSVolumesConnectionConfig,
+    uploader=DatabricksAWSVolumesUploader,
+    uploader_config=DatabricksAWSVolumesUploaderConfig,
+)
+
+databricks_aws_volumes_source_entry = SourceRegistryEntry(
+    connection_config=DatabricksAWSVolumesConnectionConfig,
+    indexer=DatabricksAWSVolumesIndexer,
+    indexer_config=DatabricksAWSVolumesIndexerConfig,
+    downloader=DatabricksAWSVolumesDownloader,
+    downloader_config=DatabricksAWSVolumesDownloaderConfig,
+)

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.interfaces import AccessConfig
+from unstructured_ingest.v2.processes.connector_registry import (
+    DestinationRegistryEntry,
+    SourceRegistryEntry,
+)
+from unstructured_ingest.v2.processes.connectors.databricks.volumes import (
+    DatabricksVolumesConnectionConfig,
+    DatabricksVolumesDownloader,
+    DatabricksVolumesDownloaderConfig,
+    DatabricksVolumesIndexer,
+    DatabricksVolumesIndexerConfig,
+    DatabricksVolumesUploader,
+    DatabricksVolumesUploaderConfig,
+)
+
+CONNECTOR_TYPE = "databricks_volumes_azure"
+
+
+class DatabricksAzureVolumesAccessConfig(AccessConfig):
+    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.interfaces import AccessConfig
+from unstructured_ingest.v2.processes.connector_registry import (
+    DestinationRegistryEntry,
+    SourceRegistryEntry,
+)
+from unstructured_ingest.v2.processes.connectors.databricks.volumes import (
+    DatabricksVolumesConnectionConfig,
+    DatabricksVolumesDownloader,
+    DatabricksVolumesDownloaderConfig,
+    DatabricksVolumesIndexer,
+    DatabricksVolumesIndexerConfig,
+    DatabricksVolumesUploader,
+    DatabricksVolumesUploaderConfig,
+)
+
+CONNECTOR_TYPE = "databricks_volumes_gcp"
+
+
+class DatabricksGoogleVolumesAccessConfig(AccessConfig):
+    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.interfaces import AccessConfig
+from unstructured_ingest.v2.processes.connector_registry import (
+    DestinationRegistryEntry,
+    SourceRegistryEntry,
+)
+from unstructured_ingest.v2.processes.connectors.databricks.volumes import (
+    DatabricksVolumesConnectionConfig,
+    DatabricksVolumesDownloader,
+    DatabricksVolumesDownloaderConfig,
+    DatabricksVolumesIndexer,
+    DatabricksVolumesIndexerConfig,
+    DatabricksVolumesUploader,
+    DatabricksVolumesUploaderConfig,
+)
+
+CONNECTOR_TYPE = "databricks_volumes"
+
+
+class DatabricksNativeVolumesAccessConfig(AccessConfig):
+    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/fsspec/fsspec.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import random
 from dataclasses import dataclass, field
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Generator, Optional, TypeVar
@@ -63,6 +64,7 @@ class FileConfig(BaseModel):
 
 class FsspecIndexerConfig(FileConfig, IndexerConfig):
     recursive: bool = False
+    sample_n_files: Optional[int] = None
 
 
 class FsspecAccessConfig(AccessConfig):
@@ -128,8 +130,23 @@ class FsspecIndexer(Indexer):
         filtered_files = [
             file for file in files if file.get("size") > 0 and file.get("type") == "file"
         ]
+
+        if self.index_config.sample_n_files:
+            filtered_files = self.sample_n_files(filtered_files, self.index_config.sample_n_files)
+
         return filtered_files
 
+    def sample_n_files(self, files: list[dict[str, Any]], n) -> list[dict[str, Any]]:
+        if len(files) <= n:
+            logger.warning(
+                f"number of files to be sampled={n} is not smaller than the number"
+                f" of files found ({len(files)}). Returning all of the files as the"
+                " sample."
+            )
+            return files
+
+        return random.sample(files, n)
+
     def get_metadata(self, file_data: dict) -> FileDataSourceMetadata:
         raise NotImplementedError()
 

unstructured_ingest/v2/processes/connectors/kdbai.py

@@ -26,7 +26,7 @@ from unstructured_ingest.v2.processes.connector_registry import (
 )
 
 if TYPE_CHECKING:
-    from kdbai_client import Session, Table
+    from kdbai_client import Database, Session, Table
 
 CONNECTOR_TYPE = "kdbai"
 
@@ -99,6 +99,9 @@ class KdbaiUploadStager(UploadStager):
 
 
 class KdbaiUploaderConfig(UploaderConfig):
+    database_name: str = Field(
+        default="default", description="The name of the KDBAI database to write into."
+    )
     table_name: str = Field(description="The name of the KDBAI table to write into.")
     batch_size: int = Field(default=100, description="Number of records per batch")
 
@@ -111,24 +114,29 @@ class KdbaiUploader(Uploader):
 
     def precheck(self) -> None:
         try:
-            self.get_table()
+            self.get_database()
         except Exception as e:
             logger.error(f"Failed to validate connection {e}", exc_info=True)
             raise DestinationConnectionError(f"failed to validate connection: {e}")
 
-    def get_table(self) -> "Table":
+    def get_database(self) -> "Database":
         session: Session = self.connection_config.get_session()
-        table = session.table(self.upload_config.table_name)
+        db = session.database(self.upload_config.database_name)
+        return db
+
+    def get_table(self) -> "Table":
+        db = self.get_database()
+        table = db.table(self.upload_config.table_name)
         return table
 
     def upsert_batch(self, batch: pd.DataFrame):
         table = self.get_table()
-        table.insert(data=batch)
+        table.insert(batch)
 
     def process_dataframe(self, df: pd.DataFrame):
         logger.debug(
             f"uploading {len(df)} entries to {self.connection_config.endpoint} "
-            f"db in table {self.upload_config.table_name}"
+            f"db {self.upload_config.database_name} in table {self.upload_config.table_name}"
         )
         for _, batch_df in df.groupby(np.arange(len(df)) // self.upload_config.batch_size):
             self.upsert_batch(batch=batch_df)