unstructured-ingest 0.0.24__py3-none-any.whl → 0.1.0__py3-none-any.whl

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/databricks_volumes.py

@@ -1,21 +1,35 @@
 import os
 from dataclasses import dataclass
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING, Any, Generator, Optional
 
 from pydantic import Field, Secret
 
-from unstructured_ingest.error import DestinationConnectionError
+from unstructured_ingest.error import (
+    DestinationConnectionError,
+    SourceConnectionError,
+    SourceConnectionNetworkError,
+)
 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,
     Uploader,
     UploaderConfig,
 )
 from unstructured_ingest.v2.logger import logger
-from unstructured_ingest.v2.processes.connector_registry import DestinationRegistryEntry
+from unstructured_ingest.v2.processes.connector_registry import (
+    DestinationRegistryEntry,
+    SourceRegistryEntry,
+)
 
 if TYPE_CHECKING:
     from databricks.sdk import WorkspaceClient
@@ -32,16 +46,6 @@ class DatabricksVolumesAccessConfig(AccessConfig):
         "https://accounts.azuredatabricks.net/ (Azure), "
         "or https://accounts.gcp.databricks.com/ (GCP).",
     )
-    username: Optional[str] = Field(
-        default=None,
-        description="The Databricks username part of basic authentication. "
-        "Only possible when Host is *.cloud.databricks.com (AWS).",
-    )
-    password: Optional[str] = Field(
-        default=None,
-        description="The Databricks password part of basic authentication. "
-        "Only possible when Host is *.cloud.databricks.com (AWS).",
-    )
     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."
@@ -78,7 +82,6 @@ class DatabricksVolumesAccessConfig(AccessConfig):
         "argument. This argument also holds the currently "
         "selected auth.",
     )
-    cluster_id: Optional[str] = None
     google_credentials: Optional[str] = None
     google_service_account: Optional[str] = None
 
@@ -93,17 +96,11 @@ class DatabricksVolumesConnectionConfig(ConnectionConfig):
         "Databricks workspace endpoint or the "
         "Databricks accounts endpoint.",
     )
-
-
-class DatabricksVolumesUploaderConfig(UploaderConfig):
     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"
     )
-    overwrite: bool = Field(
-        default=False, description="If true, an existing file will be overwritten."
-    )
     databricks_schema: str = Field(
         default="default",
         alias="schema",
@@ -117,33 +114,121 @@ class DatabricksVolumesUploaderConfig(UploaderConfig):
             path = f"{path}/{self.volume_path}"
         return path
 
-
-@dataclass
-class DatabricksVolumesUploader(Uploader):
-    connector_type: str = CONNECTOR_TYPE
-    upload_config: DatabricksVolumesUploaderConfig
-    connection_config: DatabricksVolumesConnectionConfig
-
     @requires_dependencies(dependencies=["databricks.sdk"], extras="databricks-volumes")
     def get_client(self) -> "WorkspaceClient":
         from databricks.sdk import WorkspaceClient
 
         return WorkspaceClient(
-            host=self.connection_config.host,
-            **self.connection_config.access_config.get_secret_value().model_dump(),
+            host=self.host,
+            **self.access_config.get_secret_value().model_dump(),
         )
 
+
+@dataclass
+class DatabricksVolumesIndexerConfig(IndexerConfig):
+    recursive: bool = False
+
+
+@dataclass
+class DatabricksVolumesIndexer(Indexer):
+    index_config: DatabricksVolumesIndexerConfig
+    connection_config: DatabricksVolumesConnectionConfig
+    connector_type: str = CONNECTOR_TYPE
+
     def precheck(self) -> None:
         try:
-            assert self.get_client().current_user.me().active
+            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.connection_config.path, recursive=self.index_config.recursive
+        ):
+            if file_info.is_dir:
+                continue
+            rel_path = file_info.path.replace(self.connection_config.path, "")
+            if rel_path.startswith("/"):
+                rel_path = rel_path[1:]
+            filename = Path(file_info.path).name
+            yield FileData(
+                identifier=file_info.path,
+                connector_type=CONNECTOR_TYPE,
+                source_identifiers=SourceIdentifiers(
+                    filename=filename,
+                    rel_path=rel_path,
+                    fullpath=file_info.path,
+                ),
+                additional_metadata={
+                    "catalog": self.connection_config.catalog,
+                },
+                metadata=FileDataSourceMetadata(
+                    url=file_info.path, date_modified=str(file_info.modification_time)
+                ),
+            )
+
+
+@dataclass
+class DatabricksVolumesDownloaderConfig(DownloaderConfig):
+    pass
+
+
+@dataclass
+class DatabricksVolumesDownloader(Downloader):
+    download_config: DatabricksVolumesDownloaderConfig
+    connection_config: DatabricksVolumesConnectionConfig
+    connector_type: str = CONNECTOR_TYPE
+
+    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)
+        logger.info(f"Writing {file_data.identifier} to {download_path}")
+        try:
+            with self.connection_config.get_client().dbfs.download(path=file_data.identifier) 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):
+    overwrite: bool = Field(
+        default=False, description="If true, an existing file will be overwritten."
+    )
+
+
+@dataclass
+class DatabricksVolumesUploader(Uploader):
+    upload_config: DatabricksVolumesUploaderConfig
+    connection_config: DatabricksVolumesConnectionConfig
+    connector_type: str = CONNECTOR_TYPE
+
+    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)
+        output_path = os.path.join(self.connection_config.path, path.name)
         with open(path, "rb") as elements_file:
-            self.get_client().files.upload(
+            self.connection_config.get_client().files.upload(
                 file_path=output_path,
                 contents=elements_file,
                 overwrite=self.upload_config.overwrite,
@@ -155,3 +240,11 @@ databricks_volumes_destination_entry = DestinationRegistryEntry(
     uploader=DatabricksVolumesUploader,
     uploader_config=DatabricksVolumesUploaderConfig,
 )
+
+databricks_volumes_source_entry = SourceRegistryEntry(
+    connection_config=DatabricksVolumesConnectionConfig,
+    indexer=DatabricksVolumesIndexer,
+    indexer_config=DatabricksVolumesIndexerConfig,
+    downloader=DatabricksVolumesDownloader,
+    downloader_config=DatabricksVolumesDownloaderConfig,
+)

unstructured_ingest/v2/processes/connectors/mongodb.py

@@ -1,26 +1,37 @@
 import json
+import sys
 from dataclasses import dataclass, field
+from datetime import datetime
 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.__version__ import __version__ as unstructured_version
-from unstructured_ingest.error import DestinationConnectionError
-from unstructured_ingest.utils.data_prep import batch_generator
+from unstructured_ingest.error import DestinationConnectionError, SourceConnectionError
+from unstructured_ingest.utils.data_prep import batch_generator, flatten_dict
 from unstructured_ingest.utils.dep_check import requires_dependencies
 from unstructured_ingest.v2.interfaces import (
     AccessConfig,
     ConnectionConfig,
+    Downloader,
+    DownloaderConfig,
     FileData,
+    FileDataSourceMetadata,
+    Indexer,
+    IndexerConfig,
+    SourceIdentifiers,
     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:
@@ -53,6 +64,207 @@ class MongoDBUploadStagerConfig(UploadStagerConfig):
     pass
 
 
+class MongoDBIndexerConfig(IndexerConfig):
+    batch_size: int = Field(default=100, description="Number of records per batch")
+
+
+class MongoDBDownloaderConfig(DownloaderConfig):
+    pass
+
+
+@dataclass
+class MongoDBIndexer(Indexer):
+    connection_config: MongoDBConnectionConfig
+    index_config: MongoDBIndexerConfig
+    connector_type: str = CONNECTOR_TYPE
+
+    def precheck(self) -> None:
+        """Validates the connection to the MongoDB server."""
+        try:
+            client = self.create_client()
+            client.admin.command("ping")
+        except Exception as e:
+            logger.error(f"Failed to validate connection: {e}", exc_info=True)
+            raise SourceConnectionError(f"Failed to validate connection: {e}")
+
+    @requires_dependencies(["pymongo"], extras="mongodb")
+    def create_client(self) -> "MongoClient":
+        from pymongo import MongoClient
+        from pymongo.driver_info import DriverInfo
+        from pymongo.server_api import ServerApi
+
+        access_config = self.connection_config.access_config.get_secret_value()
+
+        if access_config.uri:
+            return MongoClient(
+                access_config.uri,
+                server_api=ServerApi(version=SERVER_API_VERSION),
+                driver=DriverInfo(name="unstructured", version=unstructured_version),
+            )
+        else:
+            return MongoClient(
+                host=self.connection_config.host,
+                port=self.connection_config.port,
+                server_api=ServerApi(version=SERVER_API_VERSION),
+            )
+
+    def run(self, **kwargs: Any) -> Generator[FileData, None, None]:
+        """Generates FileData objects for each document in the MongoDB collection."""
+        client = self.create_client()
+        database = client[self.connection_config.database]
+        collection = database[self.connection_config.collection]
+
+        # Get list of document IDs
+        ids = collection.distinct("_id")
+        batch_size = self.index_config.batch_size if self.index_config else 100
+
+        for id_batch in batch_generator(ids, batch_size=batch_size):
+            # Make sure the hash is always a positive number to create identifier
+            batch_id = str(hash(frozenset(id_batch)) + sys.maxsize + 1)
+
+            metadata = FileDataSourceMetadata(
+                date_processed=str(time()),
+                record_locator={
+                    "database": self.connection_config.database,
+                    "collection": self.connection_config.collection,
+                },
+            )
+
+            file_data = FileData(
+                identifier=batch_id,
+                doc_type="batch",
+                connector_type=self.connector_type,
+                metadata=metadata,
+                additional_metadata={
+                    "ids": [str(doc_id) for doc_id in id_batch],
+                },
+            )
+            yield file_data
+
+
+@dataclass
+class MongoDBDownloader(Downloader):
+    download_config: MongoDBDownloaderConfig
+    connection_config: MongoDBConnectionConfig
+    connector_type: str = CONNECTOR_TYPE
+
+    @requires_dependencies(["pymongo"], extras="mongodb")
+    def create_client(self) -> "MongoClient":
+        from pymongo import MongoClient
+        from pymongo.driver_info import DriverInfo
+        from pymongo.server_api import ServerApi
+
+        access_config = self.connection_config.access_config.get_secret_value()
+
+        if access_config.uri:
+            return MongoClient(
+                access_config.uri,
+                server_api=ServerApi(version=SERVER_API_VERSION),
+                driver=DriverInfo(name="unstructured", version=unstructured_version),
+            )
+        else:
+            return MongoClient(
+                host=self.connection_config.host,
+                port=self.connection_config.port,
+                server_api=ServerApi(version=SERVER_API_VERSION),
+            )
+
+    @SourceConnectionError.wrap
+    @requires_dependencies(["bson"], extras="mongodb")
+    def run(self, file_data: FileData, **kwargs: Any) -> download_responses:
+        """Fetches the document from MongoDB and writes it to a file."""
+        from bson.errors import InvalidId
+        from bson.objectid import ObjectId
+
+        client = self.create_client()
+        database = client[self.connection_config.database]
+        collection = database[self.connection_config.collection]
+
+        ids = file_data.additional_metadata.get("ids", [])
+        if not ids:
+            raise ValueError("No document IDs provided in additional_metadata")
+
+        object_ids = []
+        for doc_id in ids:
+            try:
+                object_ids.append(ObjectId(doc_id))
+            except InvalidId as e:
+                error_message = f"Invalid ObjectId for doc_id '{doc_id}': {str(e)}"
+                logger.error(error_message)
+                raise ValueError(error_message) from e
+
+        try:
+            docs = list(collection.find({"_id": {"$in": object_ids}}))
+        except Exception as e:
+            logger.error(f"Failed to fetch documents: {e}", exc_info=True)
+            raise e
+
+        download_responses = []
+        for doc in docs:
+            doc_id = doc["_id"]
+            doc.pop("_id", None)
+
+            # Extract date_created from the document or ObjectId
+            date_created = None
+            if "date_created" in doc:
+                # If the document has a 'date_created' field, use it
+                date_created = doc["date_created"]
+                if isinstance(date_created, datetime):
+                    date_created = date_created.isoformat()
+                else:
+                    # Convert to ISO format if it's a string
+                    date_created = str(date_created)
+            elif isinstance(doc_id, ObjectId):
+                # Use the ObjectId's generation time
+                date_created = doc_id.generation_time.isoformat()
+
+            flattened_dict = flatten_dict(dictionary=doc)
+            concatenated_values = "\n".join(str(value) for value in flattened_dict.values())
+
+            # Create a FileData object for each document with source_identifiers
+            individual_file_data = FileData(
+                identifier=str(doc_id),
+                connector_type=self.connector_type,
+                source_identifiers=SourceIdentifiers(
+                    filename=str(doc_id),
+                    fullpath=str(doc_id),
+                    rel_path=str(doc_id),
+                ),
+            )
+
+            # Determine the download path
+            download_path = self.get_download_path(individual_file_data)
+            if download_path is None:
+                raise ValueError("Download path could not be determined")
+
+            download_path.parent.mkdir(parents=True, exist_ok=True)
+            download_path = download_path.with_suffix(".txt")
+
+            # Write the concatenated values to the file
+            with open(download_path, "w", encoding="utf8") as f:
+                f.write(concatenated_values)
+
+            individual_file_data.local_download_path = str(download_path)
+
+            # Update metadata
+            individual_file_data.metadata = FileDataSourceMetadata(
+                date_created=date_created,  # Include date_created here
+                date_processed=str(time()),
+                record_locator={
+                    "database": self.connection_config.database,
+                    "collection": self.connection_config.collection,
+                    "document_id": str(doc_id),
+                },
+            )
+
+            download_response = self.generate_download_response(
+                file_data=individual_file_data, download_path=download_path
+            )
+            download_responses.append(download_response)
+
+        return download_responses
+
+
 @dataclass
 class MongoDBUploadStager(UploadStager):
     upload_stager_config: MongoDBUploadStagerConfig = field(
@@ -138,3 +350,11 @@ mongodb_destination_entry = DestinationRegistryEntry(
     upload_stager=MongoDBUploadStager,
     upload_stager_config=MongoDBUploadStagerConfig,
 )
+
+mongodb_source_entry = SourceRegistryEntry(
+    connection_config=MongoDBConnectionConfig,
+    indexer_config=MongoDBIndexerConfig,
+    indexer=MongoDBIndexer,
+    downloader_config=MongoDBDownloaderConfig,
+    downloader=MongoDBDownloader,
+)