unstructured-ingest 0.6.4__py3-none-any.whl → 0.7.0__py3-none-any.whl

unstructured_ingest/connector/kafka.py

@@ -1,293 +0,0 @@
-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

@@ -1,139 +0,0 @@
-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)
-        ]