unstructured-ingest 0.0.0__py3-none-any.whl

unstructured_ingest/connector/kafka.py

@@ -0,0 +1,294 @@
+import base64
+import json
+import socket
+import typing as t
+from dataclasses import dataclass
+from pathlib import Path
+
+from unstructured_ingest.enhanced_dataclass import enhanced_field
+from unstructured_ingest.error import DestinationConnectionError, SourceConnectionError
+from unstructured_ingest.interfaces import (
+    AccessConfig,
+    BaseConnectorConfig,
+    BaseDestinationConnector,
+    BaseIngestDoc,
+    BaseSingleIngestDoc,
+    BaseSourceConnector,
+    ConfigSessionHandleMixin,
+    IngestDocCleanupMixin,
+    IngestDocSessionHandleMixin,
+    SourceConnectorCleanupMixin,
+    WriteConfig,
+)
+from unstructured_ingest.logger import logger
+from unstructured_ingest.utils.data_prep import batch_generator
+from unstructured_ingest.utils.dep_check import requires_dependencies
+
+if t.TYPE_CHECKING:
+    from confluent_kafka import Consumer, Producer
+
+
+@dataclass
+class KafkaAccessConfig(AccessConfig):
+    kafka_api_key: t.Optional[str] = enhanced_field(sensitive=True)
+    secret: t.Optional[str] = enhanced_field(sensitive=True)
+
+
+@dataclass
+class SimpleKafkaConfig(ConfigSessionHandleMixin, BaseConnectorConfig):
+    bootstrap_server: str
+    port: str
+    topic: str
+    access_config: KafkaAccessConfig
+    confluent: t.Optional[bool] = True
+    num_messages_to_consume: t.Optional[int] = 1
+    timeout: t.Optional[float] = 1.0
+
+
+@dataclass
+class KafkaIngestDoc(IngestDocCleanupMixin, BaseSingleIngestDoc):
+    """Class encapsulating fetching a message and writing processed results."""
+
+    connector_config: SimpleKafkaConfig
+    raw_content: str
+    raw_filename: str
+    registry_name: str = "kafka"
+
+    def _tmp_download_file(self):
+        topic_file = self.connector_config.topic + "-" + self.raw_filename
+        return Path(self.read_config.download_dir) / topic_file
+
+    @property
+    def version(self) -> t.Optional[str]:
+        return None
+
+    @property
+    def source_url(self) -> t.Optional[str]:
+        return None
+
+    @property
+    def filename(self):
+        """The filename of the file created"""
+        return self._tmp_download_file()
+
+    def _create_full_tmp_dir_path(self):
+        self._tmp_download_file().parent.mkdir(parents=True, exist_ok=True)
+
+    @property
+    def _output_filename(self):
+        """Create filename document id combined with a hash of the query to uniquely identify
+        the output file."""
+        output_file = self.connector_config.topic + ".json"
+        return Path(self.processor_config.output_dir) / output_file
+
+    @SourceConnectionError.wrap
+    @BaseSingleIngestDoc.skip_if_file_exists
+    def get_file(self):
+        self._create_full_tmp_dir_path()
+
+        pdf_data = base64.b64decode(self.raw_content)
+
+        with open(self.filename, "wb") as file:
+            file.write(pdf_data)
+
+
+@dataclass
+class KafkaSourceConnector(SourceConnectorCleanupMixin, BaseSourceConnector):
+    """Source connector for Kafka.
+    Main job is to consume from a Kafka topic and create instances of
+    KakfaIngestDoc.
+    Note that messages have the format of:
+    <filename>: the name of the file (with correct file extension)
+    <content>: base64 encoded (whether was binary or not)
+    """
+
+    connector_config: SimpleKafkaConfig
+    _consumer: t.Optional["Consumer"] = None
+
+    def check_connection(self):
+        try:
+            self.kafka_consumer
+        except Exception as e:
+            logger.error(f"failed to validate connection: {e}", exc_info=True)
+            raise SourceConnectionError(f"failed to validate connection: {e}")
+
+    def initialize(self):
+        topic = self.connector_config.topic
+        logger.info(f"Subscribing to topic: {topic}")
+        self.kafka_consumer.subscribe([topic])
+
+    @property
+    def kafka_consumer(self):
+        if self._consumer is None:
+            self._consumer = self.create_consumer()
+        return self._consumer
+
+    @requires_dependencies(["confluent_kafka"], extras="kafka")
+    def create_consumer(self) -> "Consumer":
+        from confluent_kafka import Consumer
+
+        is_confluent = self.connector_config.confluent
+        bootstrap = self.connector_config.bootstrap_server
+        port = self.connector_config.port
+
+        conf = {
+            "bootstrap.servers": f"{bootstrap}:{port}",
+            "client.id": socket.gethostname(),
+            "group.id": "your_group_id",
+            "enable.auto.commit": "false",
+            "auto.offset.reset": "earliest",
+            "message.max.bytes": 10485760,
+        }
+
+        if is_confluent:
+            kafka_api_key = self.connector_config.access_config.kafka_api_key
+            secret = self.connector_config.access_config.secret
+            conf["sasl.mechanism"] = "PLAIN"
+            conf["security.protocol"] = "SASL_SSL"
+            conf["sasl.username"] = kafka_api_key
+            conf["sasl.password"] = secret
+
+        consumer = Consumer(conf)
+        logger.debug(f"Kafka Consumer connected to bootstrap: {bootstrap}")
+        return consumer
+
+    @SourceConnectionError.wrap
+    def get_ingest_docs(self):
+        from confluent_kafka import KafkaError
+
+        consumer = self.kafka_consumer
+        running = True
+
+        collected = []
+        num_messages_to_consume = self.connector_config.num_messages_to_consume
+        logger.info(f"Config set for blocking on {num_messages_to_consume} messages")
+        # Consume specified number of messages
+        while running:
+            msg = consumer.poll(timeout=self.connector_config.timeout)
+            if msg is None:
+                logger.debug("No Kafka messages found")
+                continue
+            if msg.error():
+                if msg.error().code() == KafkaError._PARTITION_EOF:
+                    # End of partition event
+                    logger.error(
+                        "%% %s [%d] reached end at offset %d\n"
+                        % (msg.topic(), msg.partition(), msg.offset())
+                    )
+            else:
+                collected.append(json.loads(msg.value().decode("utf8")))
+                if len(collected) >= num_messages_to_consume:
+                    logger.debug(f"Found {len(collected)} messages, stopping")
+                    consumer.commit(asynchronous=False)
+                    break
+
+        return [
+            KafkaIngestDoc(
+                connector_config=self.connector_config,
+                processor_config=self.processor_config,
+                read_config=self.read_config,
+                raw_filename=msg["filename"],
+                raw_content=msg["content"],
+            )
+            for msg in collected
+        ]
+
+
+@dataclass
+class KafkaWriteConfig(WriteConfig):
+    batch_size: int = 4
+
+
+@dataclass
+class KafkaDestinationConnector(IngestDocSessionHandleMixin, BaseDestinationConnector):
+    """Connector to write BaseIngestDoc types to Kafka
+    Writes messages to Kafka in the format:
+    "type"<type>
+    "text":<the partitioned text>
+    "filename":<name of the upstream file>
+    """
+
+    write_config: KafkaWriteConfig
+    connector_config: SimpleKafkaConfig
+    _producer: t.Optional["Producer"] = None
+
+    @property
+    def kafka_producer(self):
+        if self._producer is None:
+            self._producer = self.create_producer()
+        return self._producer
+
+    def initialize(self):
+        pass
+
+    @requires_dependencies(["confluent_kafka"], extras="kafka")
+    def create_producer(self) -> "Producer":
+        from confluent_kafka import Producer
+
+        is_confluent = self.connector_config.confluent
+        bootstrap = self.connector_config.bootstrap_server
+        port = self.connector_config.port
+
+        conf = {
+            "bootstrap.servers": f"{bootstrap}:{port}",
+            "client.id": socket.gethostname(),
+        }
+
+        if is_confluent:
+            api_key = self.connector_config.access_config.kafka_api_key
+            secret = self.connector_config.access_config.secret
+            conf["sasl.mechanism"] = "PLAIN"
+            conf["security.protocol"] = "SASL_SSL"
+            conf["sasl.username"] = api_key
+            conf["sasl.password"] = secret
+
+        producer = Producer(conf)
+        logger.debug(f"Connected to bootstrap: {bootstrap}")
+        return producer
+
+    def check_connection(self):
+        try:
+            self.kafka_producer
+        except Exception as e:
+            logger.error(f"failed to validate connection: {e}", exc_info=True)
+            raise DestinationConnectionError(f"failed to validate connection: {e}")
+
+    @DestinationConnectionError.wrap
+    def upload_msg(self, batch) -> int:
+        logger.debug(f"Uploading batch: {batch}")
+        topic = self.connector_config.topic
+        producer = self.kafka_producer
+        uploaded = 0
+        for i in range(len(batch)):
+            filename = f'{batch[i].pop("filename")}'
+            producer.produce(topic, key=filename, value=str(batch[i]))
+            uploaded += 1
+        return uploaded
+
+    @DestinationConnectionError.wrap
+    def write_dict(self, *args, dict_list: t.List[t.Dict[str, t.Any]], **kwargs) -> None:
+        logger.info(f"Writing {len(dict_list)} documents to Kafka")
+        num_uploaded = 0
+
+        for chunk in batch_generator(dict_list, self.write_config.batch_size):
+            num_uploaded += self.upload_msg(chunk)  # noqa: E203
+
+        producer = self.kafka_producer
+        producer.flush()
+        logger.info(f"Uploaded {num_uploaded} documents to Kafka")
+
+    def write(self, docs: t.List[BaseIngestDoc]) -> None:
+        content_list: t.List[t.Dict[str, t.Any]] = []
+        for doc in docs:
+            local_path = doc._output_filename
+            with open(local_path) as json_file:
+                dict_content = json.load(json_file)
+                for content in dict_content:
+                    content_list.append(
+                        {
+                            "type": content["type"],
+                            "text": content["text"],
+                            "filename": content["metadata"]["filename"],
+                        }
+                    )
+        self.write_dict(dict_list=content_list)

unstructured_ingest/connector/local.py

@@ -0,0 +1,139 @@
+import fnmatch
+import glob
+import os
+import typing as t
+from dataclasses import dataclass
+from datetime import datetime
+from pathlib import Path
+
+from unstructured_ingest.interfaces import (
+    BaseConnectorConfig,
+    BaseSingleIngestDoc,
+    BaseSourceConnector,
+    SourceMetadata,
+)
+from unstructured_ingest.logger import logger
+
+
+@dataclass
+class SimpleLocalConfig(BaseConnectorConfig):
+    # Local specific options
+    input_path: str
+    recursive: bool = False
+    file_glob: t.Optional[t.List[str]] = None
+
+    def __post_init__(self):
+        if os.path.isfile(self.input_path):
+            self.input_path_is_file = True
+        else:
+            self.input_path_is_file = False
+
+
+@dataclass
+class LocalIngestDoc(BaseSingleIngestDoc):
+    """Class encapsulating fetching a doc and writing processed results (but not
+    doing the processing!).
+    """
+
+    connector_config: SimpleLocalConfig
+    path: str
+    registry_name: str = "local"
+
+    @property
+    def base_filename(self) -> t.Optional[str]:
+        download_path = Path(self.connector_config.input_path).resolve()
+        full_path = Path(self.filename).resolve()
+        if download_path.is_file():
+            download_path = download_path.parent
+        relative_path = full_path.relative_to(download_path)
+        return str(relative_path)
+
+    @property
+    def filename(self):
+        """The filename of the local file to be processed"""
+        return Path(self.path)
+
+    def cleanup_file(self):
+        """Not applicable to local file system"""
+
+    def get_file(self):
+        """Not applicable to local file system"""
+
+    def update_source_metadata(self, **kwargs) -> None:
+        try:
+            out = os.lstat(self.path)
+            self._source_metadata = SourceMetadata(
+                exists=True,
+                date_created=str(datetime.fromtimestamp(out.st_ctime)),
+                date_modified=str(datetime.fromtimestamp(out.st_mtime)),
+                permissions_data=[{"mode": out.st_mode}],
+                source_url=self.path,
+            )
+        except FileNotFoundError:
+            self._source_metadata = SourceMetadata(exists=False)
+
+    @property
+    def _output_filename(self) -> Path:
+        """Returns output filename for the doc
+        If input path argument is a file itself, it returns the filename of the doc.
+        If input path argument is a folder, it returns the relative path of the doc.
+        """
+        input_path = Path(self.connector_config.input_path)
+        basename = (
+            f"{self.base_filename}.json"
+            if input_path.is_file()
+            else f"{Path(self.path).relative_to(input_path)}.json"
+        )
+        return Path(self.processor_config.output_dir) / basename
+
+
+@dataclass
+class LocalSourceConnector(BaseSourceConnector):
+    """Objects of this class support fetching document(s) from local file system"""
+
+    def check_connection(self):
+        pass
+
+    connector_config: SimpleLocalConfig
+
+    def __post_init__(self):
+        self.ingest_doc_cls: t.Type[LocalIngestDoc] = LocalIngestDoc
+
+    def cleanup(self, cur_dir=None):
+        """Not applicable to local file system"""
+
+    def initialize(self):
+        """Not applicable to local file system"""
+
+    def _list_files(self):
+        if self.connector_config.input_path_is_file:
+            return glob.glob(f"{self.connector_config.input_path}")
+        elif self.connector_config.recursive:
+            return glob.glob(
+                f"{self.connector_config.input_path}/**",
+                recursive=self.connector_config.recursive,
+            )
+        else:
+            return glob.glob(f"{self.connector_config.input_path}/*")
+
+    def does_path_match_glob(self, path: str) -> bool:
+        if self.connector_config.file_glob is None:
+            return True
+        patterns = self.connector_config.file_glob
+        for pattern in patterns:
+            if fnmatch.filter([path], pattern):
+                return True
+        logger.debug(f"The file {path!r} is discarded as it does not match any given glob.")
+        return False
+
+    def get_ingest_docs(self):
+        return [
+            self.ingest_doc_cls(
+                connector_config=self.connector_config,
+                processor_config=self.processor_config,
+                read_config=self.read_config,
+                path=file,
+            )
+            for file in self._list_files()
+            if os.path.isfile(file) and self.does_path_match_glob(file)
+        ]