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

unstructured_ingest/v2/processes/connectors/kafka/kafka.py

@@ -0,0 +1,251 @@
+import json
+from abc import ABC, abstractmethod
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from pathlib import Path
+from time import time
+from typing import TYPE_CHECKING, Any, ContextManager, Generator, Optional
+
+from pydantic import Field, Secret
+
+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,
+    SourceIdentifiers,
+    Uploader,
+    UploaderConfig,
+)
+from unstructured_ingest.v2.logger import logger
+
+if TYPE_CHECKING:
+    from confluent_kafka import Consumer, Producer
+
+
+class KafkaAccessConfig(AccessConfig, ABC):
+    pass
+
+
+class KafkaConnectionConfig(ConnectionConfig, ABC):
+    access_config: Secret[KafkaAccessConfig]
+    bootstrap_server: str
+    port: int
+
+    @abstractmethod
+    def get_consumer_configuration(self) -> dict:
+        pass
+
+    @abstractmethod
+    def get_producer_configuration(self) -> dict:
+        pass
+
+    @contextmanager
+    @requires_dependencies(["confluent_kafka"], extras="kafka")
+    def get_consumer(self) -> ContextManager["Consumer"]:
+        from confluent_kafka import Consumer
+
+        consumer = Consumer(self.get_consumer_configuration())
+        try:
+            logger.debug("kafka consumer connected")
+            yield consumer
+        finally:
+            consumer.close()
+
+    @requires_dependencies(["confluent_kafka"], extras="kafka")
+    def get_producer(self) -> "Producer":
+        from confluent_kafka import Producer
+
+        producer = Producer(self.get_producer_configuration())
+        return producer
+
+
+class KafkaIndexerConfig(IndexerConfig):
+    topic: str = Field(description="which topic to consume from")
+    num_messages_to_consume: Optional[int] = 100
+    timeout: Optional[float] = Field(default=1.0, description="polling timeout")
+
+    def update_consumer(self, consumer: "Consumer") -> None:
+        consumer.subscribe([self.topic])
+
+
+@dataclass
+class KafkaIndexer(Indexer, ABC):
+    connection_config: KafkaConnectionConfig
+    index_config: KafkaIndexerConfig
+
+    @contextmanager
+    def get_consumer(self) -> ContextManager["Consumer"]:
+        with self.connection_config.get_consumer() as consumer:
+            self.index_config.update_consumer(consumer=consumer)
+            yield consumer
+
+    @requires_dependencies(["confluent_kafka"], extras="kafka")
+    def generate_messages(self) -> Generator[Any, None, None]:
+        from confluent_kafka import KafkaError, KafkaException
+
+        messages_consumed = 0
+        max_empty_polls = 10
+        empty_polls = 0
+        num_messages_to_consume = self.index_config.num_messages_to_consume
+        with self.get_consumer() as consumer:
+            while messages_consumed < num_messages_to_consume and empty_polls < max_empty_polls:
+                msg = consumer.poll(timeout=self.index_config.timeout)
+                if msg is None:
+                    logger.debug("No Kafka messages found")
+                    empty_polls += 1
+                    continue
+                if msg.error():
+                    if msg.error().code() == KafkaError._PARTITION_EOF:
+                        logger.info(
+                            "Reached end of partition for topic %s [%d] at offset %d"
+                            % (msg.topic(), msg.partition(), msg.offset())
+                        )
+                        break
+                    else:
+                        raise KafkaException(msg.error())
+                try:
+                    empty_polls = 0
+                    messages_consumed += 1
+                    yield msg
+                finally:
+                    consumer.commit(asynchronous=False)
+
+    def generate_file_data(self, msg) -> FileData:
+        msg_content = msg.value().decode("utf8")
+        identifier = f"{msg.topic()}_{msg.partition()}_{msg.offset()}"
+        additional_metadata = {
+            "topic": msg.topic(),
+            "partition": msg.partition(),
+            "offset": msg.offset(),
+            "content": msg_content,
+        }
+        filename = f"{identifier}.txt"
+        return FileData(
+            identifier=identifier,
+            connector_type=self.connector_type,
+            source_identifiers=SourceIdentifiers(
+                filename=filename,
+                fullpath=filename,
+            ),
+            metadata=FileDataSourceMetadata(
+                date_processed=str(time()),
+            ),
+            additional_metadata=additional_metadata,
+            display_name=filename,
+        )
+
+    def run(self) -> Generator[FileData, None, None]:
+        for message in self.generate_messages():
+            yield self.generate_file_data(message)
+
+    async def run_async(self, file_data: FileData, **kwargs: Any) -> DownloadResponse:
+        raise NotImplementedError()
+
+    def precheck(self):
+        try:
+            with self.get_consumer() as consumer:
+                cluster_meta = consumer.list_topics(timeout=self.index_config.timeout)
+                current_topics = [
+                    topic for topic in cluster_meta.topics if topic != "__consumer_offsets"
+                ]
+                logger.info(f"successfully checked available topics: {current_topics}")
+        except Exception as e:
+            logger.error(f"failed to validate connection: {e}", exc_info=True)
+            raise SourceConnectionError(f"failed to validate connection: {e}")
+
+
+class KafkaDownloaderConfig(DownloaderConfig):
+    pass
+
+
+@dataclass
+class KafkaDownloader(Downloader, ABC):
+    connection_config: KafkaConnectionConfig
+    download_config: KafkaDownloaderConfig = field(default_factory=KafkaDownloaderConfig)
+    version: Optional[str] = None
+    source_url: Optional[str] = None
+
+    def run(self, file_data: FileData, **kwargs: Any) -> DownloadResponse:
+        source_identifiers = file_data.source_identifiers
+        if source_identifiers is None:
+            raise ValueError("FileData is missing source_identifiers")
+
+        # Build the download path using source_identifiers
+        download_path = Path(self.download_dir) / source_identifiers.relative_path
+        download_path.parent.mkdir(parents=True, exist_ok=True)
+
+        try:
+            content = file_data.additional_metadata["content"]
+            with open(download_path, "w") as file:
+                file.write(content)
+        except Exception as e:
+            logger.error(f"Failed to download file {file_data.identifier}: {e}")
+            raise SourceConnectionNetworkError(f"failed to download file {file_data.identifier}")
+
+        return self.generate_download_response(file_data=file_data, download_path=download_path)
+
+
+class KafkaUploaderConfig(UploaderConfig):
+    batch_size: int = Field(default=100, description="Batch size")
+    topic: str = Field(description="which topic to write to")
+    timeout: Optional[float] = Field(
+        default=10.0, description="Timeout in seconds to flush batch of messages"
+    )
+
+
+@dataclass
+class KafkaUploader(Uploader, ABC):
+    connection_config: KafkaConnectionConfig
+    upload_config: KafkaUploaderConfig
+
+    def precheck(self):
+        try:
+            with self.connection_config.get_consumer() as consumer:
+                cluster_meta = consumer.list_topics(timeout=self.upload_config.timeout)
+                current_topics = [
+                    topic for topic in cluster_meta.topics if topic != "__consumer_offsets"
+                ]
+                logger.info(f"successfully checked available topics: {current_topics}")
+        except Exception as e:
+            logger.error(f"failed to validate connection: {e}", exc_info=True)
+            raise DestinationConnectionError(f"failed to validate connection: {e}")
+
+    def produce_batch(self, elements: list[dict]) -> None:
+        from confluent_kafka.error import KafkaException
+
+        producer = self.connection_config.get_producer()
+        failed_producer = False
+
+        def acked(err, msg):
+            if err is not None:
+                logger.error("Failed to deliver message: %s: %s" % (str(msg), str(err)))
+
+        for element in elements:
+            producer.produce(
+                topic=self.upload_config.topic,
+                value=json.dumps(element),
+                callback=acked,
+            )
+
+        producer.flush(timeout=self.upload_config.timeout)
+        if failed_producer:
+            raise KafkaException("failed to produce all messages in batch")
+
+    def run(self, path: Path, file_data: FileData, **kwargs: Any) -> None:
+        with path.open("r") as elements_file:
+            elements = json.load(elements_file)
+        for element_batch in batch_generator(elements, batch_size=self.upload_config.batch_size):
+            self.produce_batch(elements=element_batch)

unstructured_ingest/v2/processes/connectors/kafka/local.py

@@ -0,0 +1,103 @@
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+from pydantic import Field, Secret
+
+from unstructured_ingest.v2.processes.connector_registry import (
+    DestinationRegistryEntry,
+    SourceRegistryEntry,
+)
+from unstructured_ingest.v2.processes.connectors.kafka.kafka import (
+    KafkaAccessConfig,
+    KafkaConnectionConfig,
+    KafkaDownloader,
+    KafkaDownloaderConfig,
+    KafkaIndexer,
+    KafkaIndexerConfig,
+    KafkaUploader,
+    KafkaUploaderConfig,
+)
+
+if TYPE_CHECKING:
+    pass
+
+CONNECTOR_TYPE = "kafka-local"
+
+
+class LocalKafkaAccessConfig(KafkaAccessConfig):
+    pass
+
+
+class LocalKafkaConnectionConfig(KafkaConnectionConfig):
+    access_config: Secret[LocalKafkaAccessConfig] = Field(
+        default=LocalKafkaAccessConfig(), validate_default=True
+    )
+
+    def get_consumer_configuration(self) -> dict:
+        bootstrap = self.bootstrap_server
+        port = self.port
+
+        conf = {
+            "bootstrap.servers": f"{bootstrap}:{port}",
+            "group.id": "default_group_id",
+            "enable.auto.commit": "false",
+            "auto.offset.reset": "earliest",
+        }
+        return conf
+
+    def get_producer_configuration(self) -> dict:
+        bootstrap = self.bootstrap_server
+        port = self.port
+
+        conf = {
+            "bootstrap.servers": f"{bootstrap}:{port}",
+        }
+        return conf
+
+
+class LocalKafkaIndexerConfig(KafkaIndexerConfig):
+    pass
+
+
+@dataclass
+class LocalKafkaIndexer(KafkaIndexer):
+    connection_config: LocalKafkaConnectionConfig
+    index_config: LocalKafkaIndexerConfig
+    connector_type: str = CONNECTOR_TYPE
+
+
+class LocalKafkaDownloaderConfig(KafkaDownloaderConfig):
+    pass
+
+
+@dataclass
+class LocalKafkaDownloader(KafkaDownloader):
+    connection_config: LocalKafkaConnectionConfig
+    download_config: LocalKafkaDownloaderConfig
+    connector_type: str = CONNECTOR_TYPE
+
+
+class LocalKafkaUploaderConfig(KafkaUploaderConfig):
+    pass
+
+
+@dataclass
+class LocalKafkaUploader(KafkaUploader):
+    connection_config: LocalKafkaConnectionConfig
+    upload_config: LocalKafkaUploaderConfig
+    connector_type: str = CONNECTOR_TYPE
+
+
+kafka_local_source_entry = SourceRegistryEntry(
+    connection_config=LocalKafkaConnectionConfig,
+    indexer=LocalKafkaIndexer,
+    indexer_config=LocalKafkaIndexerConfig,
+    downloader=LocalKafkaDownloader,
+    downloader_config=LocalKafkaDownloaderConfig,
+)
+
+kafka_local_destination_entry = DestinationRegistryEntry(
+    connection_config=LocalKafkaConnectionConfig,
+    uploader=LocalKafkaUploader,
+    uploader_config=LocalKafkaUploaderConfig,
+)

unstructured_ingest/v2/processes/connectors/onedrive.py

@@ -9,7 +9,11 @@ from typing import TYPE_CHECKING, Any, Generator, Optional
 from dateutil import parser
 from pydantic import Field, Secret
 
-from unstructured_ingest.error import SourceConnectionError, SourceConnectionNetworkError
+from unstructured_ingest.error import (
+    DestinationConnectionError,
+    SourceConnectionError,
+    SourceConnectionNetworkError,
+)
 from unstructured_ingest.utils.dep_check import requires_dependencies
 from unstructured_ingest.v2.interfaces import (
     AccessConfig,
@@ -22,16 +26,19 @@ from unstructured_ingest.v2.interfaces import (
     Indexer,
     IndexerConfig,
     SourceIdentifiers,
-    download_responses,
+    Uploader,
+    UploaderConfig,
 )
 from unstructured_ingest.v2.logger import logger
 from unstructured_ingest.v2.processes.connector_registry import (
+    DestinationRegistryEntry,
     SourceRegistryEntry,
 )
 
 if TYPE_CHECKING:
     from office365.graph_client import GraphClient
     from office365.onedrive.driveitems.driveItem import DriveItem
+    from office365.onedrive.drives.drive import Drive
 
 CONNECTOR_TYPE = "onedrive"
 MAX_MB_SIZE = 512_000_000
@@ -55,6 +62,11 @@ class OnedriveConnectionConfig(ConnectionConfig):
     )
     access_config: Secret[OnedriveAccessConfig]
 
+    def get_drive(self) -> "Drive":
+        client = self.get_client()
+        drive = client.users[self.user_pname].drive
+        return drive
+
     @requires_dependencies(["msal"], extras="onedrive")
     def get_token(self):
         from msal import ConfidentialClientApplication
@@ -100,7 +112,6 @@ class OnedriveIndexer(Indexer):
                 raise SourceConnectionError(
                     "{} ({})".format(error, token_resp.get("error_description"))
                 )
-            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}")
@@ -208,7 +219,7 @@ class OnedriveDownloader(Downloader):
         return self.download_dir / Path(rel_path)
 
     @SourceConnectionError.wrap
-    def run(self, file_data: FileData, **kwargs: Any) -> download_responses:
+    def run(self, file_data: FileData, **kwargs: Any) -> DownloadResponse:
         file = self._fetch_file(file_data=file_data)
         fsize = file.get_property("size", 0)
         download_path = self.get_download_path(file_data=file_data)
@@ -221,7 +232,150 @@ class OnedriveDownloader(Downloader):
         else:
             with download_path.open(mode="wb") as f:
                 file.download(f).execute_query()
-        return DownloadResponse(file_data=file_data, path=download_path)
+        return self.generate_download_response(file_data=file_data, download_path=download_path)
+
+
+class OnedriveUploaderConfig(UploaderConfig):
+    remote_url: str = Field(
+        description="URL of the destination in OneDrive, e.g., 'onedrive://Documents/Folder'"
+    )
+    prefix: str = "onedrive://"
+
+    @property
+    def root_folder(self) -> str:
+        url = (
+            self.remote_url.replace(self.prefix, "", 1)
+            if self.remote_url.startswith(self.prefix)
+            else self.remote_url
+        )
+        return url.split("/")[0]
+
+    @property
+    def url(self) -> str:
+        url = (
+            self.remote_url.replace(self.prefix, "", 1)
+            if self.remote_url.startswith(self.prefix)
+            else self.remote_url
+        )
+        return url
+
+
+@dataclass
+class OnedriveUploader(Uploader):
+    connection_config: OnedriveConnectionConfig
+    upload_config: OnedriveUploaderConfig
+    connector_type: str = CONNECTOR_TYPE
+
+    @requires_dependencies(["office365"], extras="onedrive")
+    def precheck(self) -> None:
+        from office365.runtime.client_request_exception import ClientRequestException
+
+        try:
+            token_resp: dict = self.connection_config.get_token()
+            if error := token_resp.get("error"):
+                raise SourceConnectionError(
+                    "{} ({})".format(error, token_resp.get("error_description"))
+                )
+            drive = self.connection_config.get_drive()
+            root = drive.root
+            root_folder = self.upload_config.root_folder
+            folder = root.get_by_path(root_folder)
+            try:
+                folder.get().execute_query()
+            except ClientRequestException as e:
+                if e.message != "The resource could not be found.":
+                    raise e
+                folder = root.create_folder(root_folder).execute_query()
+                logger.info(f"successfully created folder: {folder.name}")
+        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, path: Path, file_data: FileData, **kwargs: Any) -> None:
+        drive = self.connection_config.get_drive()
+
+        # Use the remote_url from upload_config as the base destination folder
+        base_destination_folder = self.upload_config.url
+
+        # Use the file's relative path to maintain directory structure, if needed
+        if file_data.source_identifiers and file_data.source_identifiers.rel_path:
+            # Combine the base destination folder with the file's relative path
+            destination_path = Path(base_destination_folder) / Path(
+                file_data.source_identifiers.rel_path
+            )
+        else:
+            # If no relative path is provided, upload directly to the base destination folder
+            destination_path = Path(base_destination_folder) / path.name
+
+        destination_folder = destination_path.parent
+        file_name = destination_path.name
+
+        # Convert destination folder to a string suitable for OneDrive API
+        destination_folder_str = str(destination_folder).replace("\\", "/")
+
+        # Resolve the destination folder in OneDrive, creating it if necessary
+        try:
+            # Attempt to get the folder
+            folder = drive.root.get_by_path(destination_folder_str)
+            folder.get().execute_query()
+        except Exception:
+            # Folder doesn't exist, create it recursively
+            current_folder = drive.root
+            for part in destination_folder.parts:
+                # Use filter to find the folder by name
+                folders = (
+                    current_folder.children.filter(f"name eq '{part}' and folder ne null")
+                    .get()
+                    .execute_query()
+                )
+                if folders:
+                    current_folder = folders[0]
+                else:
+                    # Folder doesn't exist, create it
+                    current_folder = current_folder.create_folder(part).execute_query()
+            folder = current_folder
+
+        # Check the size of the file
+        file_size = path.stat().st_size
+
+        if file_size < MAX_MB_SIZE:
+            # Use simple upload for small files
+            with path.open("rb") as local_file:
+                content = local_file.read()
+                logger.info(f"Uploading {path} to {destination_path} using simple upload")
+                try:
+                    uploaded_file = folder.upload(file_name, content).execute_query()
+                    if not uploaded_file or uploaded_file.name != file_name:
+                        raise DestinationConnectionError(f"Upload failed for file '{file_name}'")
+                    # Log details about the uploaded file
+                    logger.info(
+                        f"Uploaded file '{uploaded_file.name}' with ID '{uploaded_file.id}'"
+                    )
+                except Exception as e:
+                    logger.error(f"Failed to upload file '{file_name}': {e}", exc_info=True)
+                    raise DestinationConnectionError(
+                        f"Failed to upload file '{file_name}': {e}"
+                    ) from e
+        else:
+            # Use resumable upload for large files
+            destination_fullpath = f"{destination_folder_str}/{file_name}"
+            destination_drive_item = drive.root.item_with_path(destination_fullpath)
+
+            logger.info(f"Uploading {path} to {destination_fullpath} using resumable upload")
+            try:
+                uploaded_file = destination_drive_item.resumable_upload(
+                    source_path=str(path)
+                ).execute_query()
+                # Validate the upload
+                if not uploaded_file or uploaded_file.name != file_name:
+                    raise DestinationConnectionError(f"Upload failed for file '{file_name}'")
+                # Log details about the uploaded file
+                logger.info(f"Uploaded file {uploaded_file.name} with ID {uploaded_file.id}")
+            except Exception as e:
+                logger.error(f"Failed to upload file '{file_name}' using resumable upload: {e}")
+                raise DestinationConnectionError(
+                    f"Failed to upload file '{file_name}' using resumable upload: {e}"
+                ) from e
 
 
 onedrive_source_entry = SourceRegistryEntry(
@@ -231,3 +385,9 @@ onedrive_source_entry = SourceRegistryEntry(
     downloader_config=OnedriveDownloaderConfig,
     downloader=OnedriveDownloader,
 )
+
+onedrive_destination_entry = DestinationRegistryEntry(
+    connection_config=OnedriveConnectionConfig,
+    uploader=OnedriveUploader,
+    uploader_config=OnedriveUploaderConfig,
+)

unstructured_ingest/v2/processes/connectors/outlook.py

@@ -15,10 +15,10 @@ from unstructured_ingest.v2.interfaces import (
     ConnectionConfig,
     Downloader,
     DownloaderConfig,
+    DownloadResponse,
     FileData,
     Indexer,
     IndexerConfig,
-    download_responses,
 )
 from unstructured_ingest.v2.interfaces.file_data import FileDataSourceMetadata, SourceIdentifiers
 from unstructured_ingest.v2.processes.connector_registry import SourceRegistryEntry
@@ -191,7 +191,7 @@ class OutlookDownloader(Downloader):
     connection_config: OutlookConnectionConfig
     download_config: OutlookDownloaderConfig = field(default_factory=OutlookDownloaderConfig)
 
-    def run(self, file_data: FileData, **kwargs: Any) -> download_responses:
+    def run(self, file_data: FileData, **kwargs: Any) -> DownloadResponse:
         # NOTE: Indexer should provide source identifiers required to generate the download path
         download_path = self.get_download_path(file_data)
         if download_path is None: