unstructured-ingest 0.2.0__py3-none-any.whl → 0.2.2__py3-none-any.whl

unstructured_ingest/v2/processes/connectors/gitlab.py

@@ -0,0 +1,267 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Generator, Optional
+from urllib.parse import urlparse
+
+from pydantic import Field, Secret, model_validator
+
+from unstructured_ingest.error import SourceConnectionError
+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,
+)
+from unstructured_ingest.v2.logger import logger
+from unstructured_ingest.v2.processes.connector_registry import SourceRegistryEntry
+
+CONNECTOR_TYPE = "gitlab"
+if TYPE_CHECKING:
+    from gitlab import Gitlab
+    from gitlab.v4.objects.projects import Project
+
+
+class GitLabAccessConfig(AccessConfig):
+    access_token: Optional[str] = Field(
+        default=None,
+        description="Optional personal access token for authenticating with the GitLab API.",
+    )
+
+
+class GitLabConnectionConfig(ConnectionConfig):
+    access_config: Secret[GitLabAccessConfig] = Field(
+        default_factory=GitLabAccessConfig,
+        validate_default=True,
+        description="Secret configuration for accessing the GitLab API by authentication token.",
+    )
+    url: str = Field(description="The full URL to the GitLab project or repository.")
+    base_url: str = Field(
+        default="https://gitlab.com",
+        description="The base URL for the GitLab instance (default is GitLab's public domain).",
+    )
+    repo_path: str = Field(
+        default=None,
+        init=False,
+        repr=False,
+        description="The normalized path extracted from the repository URL.",
+    )
+
+    @model_validator(mode="after")
+    def set_repo_path(self):
+        """
+        Parses the provided GitLab URL to extract the `base_url` and `repo_path`,
+        ensuring both are properly formatted for use.
+
+        If the URL contains a scheme (e.g., 'https') and a network location (e.g., 'gitlab.com'),
+        the `base_url` is set accordingly. The repository path is extracted and normalized
+        by removing any leading slashes.
+
+        Notes:
+            - If the URL contains both a scheme and network location, the `base_url` is
+              extracted directly from the URL.
+            - The `repo_path` is adjusted to remove any leading slashes.
+            - This method assumes that the URL follows GitLab's structure
+              (e.g., 'https://gitlab.com/owner/repo').
+        """
+        parsed_gh_url = urlparse(self.url)
+
+        if parsed_gh_url.scheme and parsed_gh_url.netloc:
+            self.base_url = f"{parsed_gh_url.scheme}://{parsed_gh_url.netloc}"
+        self.repo_path = parsed_gh_url.path.lstrip("/")
+
+        return self
+
+    @SourceConnectionError.wrap
+    @requires_dependencies(["gitlab"], extras="gitlab")
+    def get_client(self) -> "Gitlab":
+        from gitlab import Gitlab
+
+        logger.info(f"Connection to GitLab: {self.base_url!r}")
+        gitlab = Gitlab(
+            self.base_url, private_token=self.access_config.get_secret_value().access_token
+        )
+        return gitlab
+
+    def get_project(self) -> "Project":
+        """Retrieves the specified GitLab project using the configured base URL and access token.
+
+        Returns:
+            Project: A GitLab `Project` object representing the specified repository.
+
+        Raises:
+            SourceConnectionError: If the GitLab API connection fails.
+            gitlab.exceptions.GitlabGetError: If the project is not found.
+        """
+        gitlab = self.get_client()
+
+        logger.info(f"Accessing Project: '{self.repo_path}'")
+        project = gitlab.projects.get(self.repo_path)
+
+        logger.info(f"Successfully accessed project '{self.repo_path}'")
+        return project
+
+
+class GitLabIndexerConfig(IndexerConfig):
+    path: Path = Field(
+        default="/", description=("Path to the location in the repository that will be processed.")
+    )
+    recursive: bool = Field(
+        default=True,
+        description=(
+            "Flag to control recursive operations when indexing. "
+            "If True, the indexer will traverse directories recursively."
+        ),
+    )
+    git_branch: Optional[str] = Field(
+        default=None,
+        description="The name of the branch to interact with.",
+    )
+
+
+@dataclass
+class GitLabIndexer(Indexer):
+    connection_config: GitLabConnectionConfig
+    index_config: GitLabIndexerConfig
+
+    def precheck(self) -> None:
+        """Validates the connection to the GitLab instance by authenticating or
+        accessing the project.
+
+        This method ensures that the GitLab credentials and configuration are correct by
+        either authenticating or attempting to fetch the specified project.
+
+        Raises:
+            SourceConnectionError: If the connection or authentication with GitLab fails.
+        """
+
+        try:
+            gitlab = self.connection_config.get_client()
+            if self.connection_config.access_config.get_secret_value().access_token is not None:
+                gitlab.auth()
+            else:
+                gitlab.projects.get(self.connection_config.repo_path)
+
+        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]:
+        """Iterates over the GitLab repository tree and yields file metadata as `FileData` objects.
+
+        This method fetches the repository tree for the specified branch and iterates
+        over its contents. For each file (blob), it generates a `FileData` object containing
+        the file's metadata, path, and permissions.
+
+        Args:
+            **kwargs (Any): Additional keyword arguments (if required).
+
+        Yields:
+            FileData: A generator that yields `FileData` objects representing each file (blob)
+            in the repository.
+        """
+        project = self.connection_config.get_project()
+
+        ref = self.index_config.git_branch or project.default_branch
+
+        files = project.repository_tree(
+            path=str(self.index_config.path),
+            ref=ref,
+            recursive=self.index_config.recursive,
+            iterator=True,
+            all=True,
+        )
+
+        for file in files:
+            relative_path = str(Path(file["path"]).relative_to(self.index_config.path))
+            if file["type"] == "blob":
+                record_locator = {
+                    "file_path": file["path"],
+                    "ref": ref,
+                }
+
+                yield FileData(
+                    identifier=file["id"],
+                    connector_type=CONNECTOR_TYPE,
+                    source_identifiers=SourceIdentifiers(
+                        fullpath=file["path"],
+                        filename=Path(file["path"]).name,
+                        rel_path=relative_path,
+                    ),
+                    metadata=FileDataSourceMetadata(
+                        url=file["id"],
+                        record_locator=record_locator,
+                        permissions_data=[{"mode": file["mode"]}],
+                    ),
+                    additional_metadata={},
+                )
+
+
+class GitLabDownloaderConfig(DownloaderConfig):
+    pass
+
+
+@dataclass
+class GitLabDownloader(Downloader):
+    connection_config: GitLabConnectionConfig
+    download_config: GitLabDownloaderConfig
+
+    def run(self, file_data: FileData, **kwargs: Any) -> DownloadResponse:
+        """Downloads a file from the repository and returns a `DownloadResponse`.
+
+        Args:
+            file_data (FileData): Metadata about the file to be downloaded.
+            **kwargs (Any): Additional arguments (if required).
+
+        Returns:
+            DownloadResponse: A response object containing the download details.
+        """
+        download_path = self.get_download_path(file_data=file_data)
+        if download_path is None:
+            logger.error(
+                "Generated download path is None, source_identifiers might be missing"
+                "from FileData."
+            )
+            raise ValueError("Generated invalid download path.")
+
+        self._download_file(file_data, download_path)
+        return self.generate_download_response(file_data=file_data, download_path=download_path)
+
+    def _download_file(self, file_data: FileData, download_path: Path) -> None:
+        # NOTE: Indexer should supply the record locator in metadata
+        if (
+            file_data.metadata.record_locator is None
+            or "ref" not in file_data.metadata.record_locator
+            or "file_path" not in file_data.metadata.record_locator
+        ):
+            logger.error(
+                f"Invalid record locator in metadata: {file_data.metadata.record_locator}."
+                "Keys 'ref' and 'path' must be present."
+            )
+            raise ValueError("Invalid record locator.")
+
+        ref = file_data.metadata.record_locator["ref"]
+        path = file_data.metadata.record_locator["file_path"]
+
+        project_file = self.connection_config.get_project().files.get(file_path=path, ref=ref)
+        download_path.parent.mkdir(exist_ok=True, parents=True)
+
+        with open(download_path, "wb") as file:
+            file.write(project_file.decode())
+
+
+gitlab_source_entry = SourceRegistryEntry(
+    connection_config=GitLabConnectionConfig,
+    indexer_config=GitLabIndexerConfig,
+    indexer=GitLabIndexer,
+    downloader_config=GitLabDownloaderConfig,
+    downloader=GitLabDownloader,
+)

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

@@ -0,0 +1,13 @@
+from __future__ import annotations
+
+from unstructured_ingest.v2.processes.connector_registry import (
+    add_source_entry,
+)
+
+from .cloud import CONNECTOR_TYPE as CLOUD_CONNECTOR
+from .cloud import kafka_cloud_source_entry
+from .local import CONNECTOR_TYPE as LOCAL_CONNECTOR
+from .local import kafka_local_source_entry
+
+add_source_entry(source_type=LOCAL_CONNECTOR, entry=kafka_local_source_entry)
+add_source_entry(source_type=CLOUD_CONNECTOR, entry=kafka_cloud_source_entry)

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

@@ -0,0 +1,82 @@
+import socket
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Optional
+
+from pydantic import Field, Secret, SecretStr
+
+from unstructured_ingest.v2.processes.connector_registry import SourceRegistryEntry
+from unstructured_ingest.v2.processes.connectors.kafka.kafka import (
+    KafkaAccessConfig,
+    KafkaConnectionConfig,
+    KafkaDownloader,
+    KafkaDownloaderConfig,
+    KafkaIndexer,
+    KafkaIndexerConfig,
+)
+
+if TYPE_CHECKING:
+    pass
+
+CONNECTOR_TYPE = "kafka-cloud"
+
+
+class CloudKafkaAccessConfig(KafkaAccessConfig):
+    api_key: Optional[SecretStr] = Field(
+        description="Kafka API key to connect at the server", alias="kafka_api_key", default=None
+    )
+    secret: Optional[SecretStr] = Field(description="", default=None)
+
+
+class CloudKafkaConnectionConfig(KafkaConnectionConfig):
+    access_config: Secret[CloudKafkaAccessConfig]
+
+    def get_consumer_configuration(self) -> dict:
+        bootstrap = self.bootstrap_server
+        port = self.port
+        access_config = self.access_config.get_secret_value()
+
+        conf = {
+            "bootstrap.servers": f"{bootstrap}:{port}",
+            "client.id": socket.gethostname(),
+            "group.id": "default_group_id",
+            "enable.auto.commit": "false",
+            "auto.offset.reset": "earliest",
+            "message.max.bytes": 10485760,
+            "sasl.username": access_config.api_key,
+            "sasl.password": access_config.secret,
+            "sasl.mechanism": "PLAIN",
+            "security.protocol": "SASL_SSL",
+        }
+
+        return conf
+
+
+class CloudKafkaIndexerConfig(KafkaIndexerConfig):
+    pass
+
+
+@dataclass
+class CloudKafkaIndexer(KafkaIndexer):
+    connection_config: CloudKafkaConnectionConfig
+    index_config: CloudKafkaIndexerConfig
+    connector_type: str = CONNECTOR_TYPE
+
+
+class CloudKafkaDownloaderConfig(KafkaDownloaderConfig):
+    pass
+
+
+@dataclass
+class CloudKafkaDownloader(KafkaDownloader):
+    connection_config: CloudKafkaConnectionConfig
+    download_config: CloudKafkaDownloaderConfig
+    connector_type: str = CONNECTOR_TYPE
+
+
+kafka_cloud_source_entry = SourceRegistryEntry(
+    connection_config=CloudKafkaConnectionConfig,
+    indexer=CloudKafkaIndexer,
+    indexer_config=CloudKafkaIndexerConfig,
+    downloader=CloudKafkaDownloader,
+    downloader_config=CloudKafkaDownloaderConfig,
+)

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

@@ -0,0 +1,196 @@
+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 Secret
+
+from unstructured_ingest.error import (
+    SourceConnectionError,
+    SourceConnectionNetworkError,
+)
+from unstructured_ingest.utils.dep_check import requires_dependencies
+from unstructured_ingest.v2.interfaces import (
+    AccessConfig,
+    ConnectionConfig,
+    Downloader,
+    DownloaderConfig,
+    FileData,
+    FileDataSourceMetadata,
+    Indexer,
+    IndexerConfig,
+    SourceIdentifiers,
+    download_responses,
+)
+from unstructured_ingest.v2.logger import logger
+from unstructured_ingest.v2.processes.connector_registry import SourceRegistryEntry
+
+if TYPE_CHECKING:
+    from confluent_kafka import Consumer
+
+CONNECTOR_TYPE = "kafka"
+
+
+class KafkaAccessConfig(AccessConfig, ABC):
+    pass
+
+
+class KafkaConnectionConfig(ConnectionConfig, ABC):
+    access_config: Secret[KafkaAccessConfig]
+    timeout: Optional[float] = 1.0
+    bootstrap_server: str
+    port: int
+
+    @abstractmethod
+    def get_consumer_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()
+
+
+class KafkaIndexerConfig(IndexerConfig):
+    topic: str
+    num_messages_to_consume: Optional[int] = 100
+
+    def update_consumer(self, consumer: "Consumer") -> None:
+        consumer.subscribe([self.topic])
+
+
+@dataclass
+class KafkaIndexer(Indexer):
+    connection_config: KafkaConnectionConfig
+    index_config: KafkaIndexerConfig
+    connector_type: str = CONNECTOR_TYPE
+
+    @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.connection_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) -> download_responses:
+        raise NotImplementedError()
+
+    def precheck(self):
+        try:
+            with self.get_consumer() as consumer:
+                cluster_meta = consumer.list_topics(timeout=self.connection_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):
+    connection_config: KafkaConnectionConfig
+    download_config: KafkaDownloaderConfig = field(default_factory=KafkaDownloaderConfig)
+    connector_type: str = CONNECTOR_TYPE
+    version: Optional[str] = None
+    source_url: Optional[str] = None
+
+    def run(self, file_data: FileData, **kwargs: Any) -> download_responses:
+        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)
+
+
+kafka_source_entry = SourceRegistryEntry(
+    connection_config=KafkaConnectionConfig,
+    indexer=KafkaIndexer,
+    indexer_config=KafkaIndexerConfig,
+    downloader=KafkaDownloader,
+    downloader_config=KafkaDownloaderConfig,
+)

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

@@ -0,0 +1,75 @@
+import socket
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+from pydantic import Field, Secret
+
+from unstructured_ingest.v2.processes.connector_registry import SourceRegistryEntry
+from unstructured_ingest.v2.processes.connectors.kafka.kafka import (
+    KafkaAccessConfig,
+    KafkaConnectionConfig,
+    KafkaDownloader,
+    KafkaDownloaderConfig,
+    KafkaIndexer,
+    KafkaIndexerConfig,
+)
+
+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}",
+            "client.id": socket.gethostname(),
+            "group.id": "default_group_id",
+            "enable.auto.commit": "false",
+            "auto.offset.reset": "earliest",
+            "message.max.bytes": 10485760,
+        }
+        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
+
+
+kafka_local_source_entry = SourceRegistryEntry(
+    connection_config=LocalKafkaConnectionConfig,
+    indexer=LocalKafkaIndexer,
+    indexer_config=LocalKafkaIndexerConfig,
+    downloader=LocalKafkaDownloader,
+    downloader_config=LocalKafkaDownloaderConfig,
+)