unstructured-ingest 0.0.0__py3-none-any.whl

unstructured_ingest/v2/processes/connectors/google_drive.py

@@ -0,0 +1,335 @@
+import io
+import os
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Generator, Optional, Union
+
+from dateutil import parser
+from unstructured.documents.elements import DataSourceMetadata
+from unstructured.file_utils.google_filetype import GOOGLE_DRIVE_EXPORT_TYPES
+
+from unstructured_ingest.enhanced_dataclass import enhanced_field
+from unstructured_ingest.error import SourceConnectionNetworkError
+from unstructured_ingest.utils.dep_check import requires_dependencies
+from unstructured_ingest.utils.string_and_date_utils import json_to_dict
+from unstructured_ingest.v2.interfaces import (
+    AccessConfig,
+    ConnectionConfig,
+    Downloader,
+    DownloaderConfig,
+    FileData,
+    Indexer,
+    IndexerConfig,
+    SourceIdentifiers,
+    download_responses,
+)
+from unstructured_ingest.v2.logger import logger
+from unstructured_ingest.v2.processes.connector_registry import (
+    SourceRegistryEntry,
+)
+
+CONNECTOR_TYPE = "google_drive"
+
+if TYPE_CHECKING:
+    from googleapiclient.discovery import Resource as GoogleAPIResource
+    from googleapiclient.http import MediaIoBaseDownload
+
+
+@dataclass
+class GoogleDriveAccessConfig(AccessConfig):
+    service_account_key: Union[str, dict]
+
+
+@dataclass
+class GoogleDriveConnectionConfig(ConnectionConfig):
+    drive_id: str
+    access_config: GoogleDriveAccessConfig = enhanced_field(sensitive=True)
+
+    @requires_dependencies(["googleapiclient"], extras="google-drive")
+    def get_files_service(self) -> "GoogleAPIResource":
+        from google.auth import default, exceptions
+        from google.oauth2 import service_account
+        from googleapiclient.discovery import build
+        from googleapiclient.errors import HttpError
+
+        # Service account key can be a dict or a file path(str)
+        # But the dict may come in as a string
+        if isinstance(self.access_config.service_account_key, str):
+            key_path = json_to_dict(self.access_config.service_account_key)
+        elif isinstance(self.access_config.service_account_key, dict):
+            key_path = self.access_config.service_account_key
+        else:
+            raise TypeError(
+                f"access_config.service_account_key must be "
+                f"str or dict, got: {type(self.access_config.service_account_key)}"
+            )
+
+        try:
+            if isinstance(key_path, dict):
+                creds = service_account.Credentials.from_service_account_info(key_path)
+            elif isinstance(key_path, str):
+                os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = key_path
+                creds, _ = default()
+            else:
+                raise ValueError(
+                    f"key path not recognized as a dictionary or a file path: "
+                    f"[{type(key_path)}] {key_path}",
+                )
+            service = build("drive", "v3", credentials=creds)
+            return service.files()
+
+        except HttpError as exc:
+            raise ValueError(f"{exc.reason}")
+        except exceptions.DefaultCredentialsError:
+            raise ValueError("The provided API key is invalid.")
+
+
+@dataclass
+class GoogleDriveIndexerConfig(IndexerConfig):
+    extensions: Optional[list[str]] = None
+    recursive: bool = False
+
+    def __post_init__(self):
+        # Strip leading period of extension
+        if self.extensions is not None:
+            self.extensions = [e[1:] if e.startswith(".") else e for e in self.extensions]
+
+
+@dataclass
+class GoogleDriveIndexer(Indexer):
+    connection_config: GoogleDriveConnectionConfig
+    index_config: GoogleDriveIndexerConfig
+    fields: list[str] = field(
+        default_factory=lambda: [
+            "id",
+            "name",
+            "mimeType",
+            "fileExtension",
+            "md5Checksum",
+            "sha1Checksum",
+            "sha256Checksum",
+            "headRevisionId",
+            "permissions",
+            "createdTime",
+            "modifiedTime",
+            "version",
+            "originalFilename",
+            "capabilities",
+            "permissionIds",
+            "webViewLink",
+            "webContentLink",
+        ]
+    )
+
+    @staticmethod
+    def is_dir(record: dict) -> bool:
+        return record.get("mimeType") == "application/vnd.google-apps.folder"
+
+    @staticmethod
+    def map_file_data(f: dict) -> FileData:
+        file_id = f["id"]
+        filename = f.pop("name")
+        url = f.pop("webContentLink", None)
+        version = f.pop("version", None)
+        permissions = f.pop("permissions", None)
+        date_created_str = f.pop("createdTime", None)
+        date_created_dt = parser.parse(date_created_str) if date_created_str else None
+        date_modified_str = f.pop("modifiedTime", None)
+        parent_path = f.pop("parent_path", None)
+        parent_root_path = f.pop("parent_root_path", None)
+        date_modified_dt = parser.parse(date_modified_str) if date_modified_str else None
+        if (
+            parent_path
+            and isinstance(parent_path, str)
+            and parent_root_path
+            and isinstance(parent_root_path, str)
+        ):
+            fullpath = f"{parent_path}/{filename}"
+            rel_path = fullpath.replace(parent_root_path, "")
+            source_identifiers = SourceIdentifiers(
+                filename=filename, fullpath=fullpath, rel_path=rel_path
+            )
+        else:
+            source_identifiers = SourceIdentifiers(fullpath=filename, filename=filename)
+        return FileData(
+            connector_type=CONNECTOR_TYPE,
+            identifier=file_id,
+            source_identifiers=source_identifiers,
+            metadata=DataSourceMetadata(
+                url=url,
+                version=version,
+                date_created=str(date_created_dt.timestamp()),
+                date_modified=str(date_modified_dt.timestamp()),
+                permissions_data=permissions,
+                record_locator={
+                    "file_id": file_id,
+                },
+            ),
+            additional_metadata=f,
+        )
+
+    def get_paginated_results(
+        self,
+        files_client,
+        object_id: str,
+        extensions: Optional[list[str]] = None,
+        recursive: bool = False,
+        previous_path: Optional[str] = None,
+    ) -> list[dict]:
+
+        fields_input = "nextPageToken, files({})".format(",".join(self.fields))
+        q = f"'{object_id}' in parents"
+        # Filter by extension but still include any directories
+        if extensions:
+            ext_filter = " or ".join([f"fileExtension = '{e}'" for e in extensions])
+            q = f"{q} and ({ext_filter} or mimeType = 'application/vnd.google-apps.folder')"
+        logger.debug(f"Query used when indexing: {q}")
+        logger.debug("response fields limited to: {}".format(", ".join(self.fields)))
+        done = False
+        page_token = None
+        files_response = []
+        while not done:
+            response: dict = files_client.list(
+                spaces="drive",
+                fields=fields_input,
+                corpora="user",
+                pageToken=page_token,
+                q=q,
+            ).execute()
+            if files := response.get("files", []):
+                fs = [f for f in files if not self.is_dir(record=f)]
+                for r in fs:
+                    r["parent_path"] = previous_path
+                dirs = [f for f in files if self.is_dir(record=f)]
+                files_response.extend(fs)
+                if recursive:
+                    for d in dirs:
+                        dir_id = d["id"]
+                        dir_name = d["name"]
+                        files_response.extend(
+                            self.get_paginated_results(
+                                files_client=files_client,
+                                object_id=dir_id,
+                                extensions=extensions,
+                                recursive=recursive,
+                                previous_path=f"{previous_path}/{dir_name}",
+                            )
+                        )
+            page_token = response.get("nextPageToken")
+            if page_token is None:
+                done = True
+        for r in files_response:
+            r["parent_root_path"] = previous_path
+        return files_response
+
+    def get_root_info(self, files_client, object_id: str) -> dict:
+        return files_client.get(fileId=object_id, fields=",".join(self.fields)).execute()
+
+    def get_files(
+        self,
+        files_client,
+        object_id: str,
+        recursive: bool = False,
+        extensions: Optional[list[str]] = None,
+    ) -> list[FileData]:
+        root_info = self.get_root_info(files_client=files_client, object_id=object_id)
+        if not self.is_dir(root_info):
+            data = [self.map_file_data(root_info)]
+        else:
+
+            file_contents = self.get_paginated_results(
+                files_client=files_client,
+                object_id=object_id,
+                extensions=extensions,
+                recursive=recursive,
+                previous_path=root_info["name"],
+            )
+            data = [self.map_file_data(f=f) for f in file_contents]
+        for d in data:
+            d.metadata.record_locator["drive_id"]: object_id
+        return data
+
+    def run(self, **kwargs: Any) -> Generator[FileData, None, None]:
+        for f in self.get_files(
+            files_client=self.connection_config.get_files_service(),
+            object_id=self.connection_config.drive_id,
+            recursive=self.index_config.recursive,
+            extensions=self.index_config.extensions,
+        ):
+            yield f
+
+
+@dataclass
+class GoogleDriveDownloaderConfig(DownloaderConfig):
+    pass
+
+
+@dataclass
+class GoogleDriveDownloader(Downloader):
+    connection_config: GoogleDriveConnectionConfig
+    download_config: GoogleDriveDownloaderConfig = field(
+        default_factory=lambda: GoogleDriveDownloaderConfig()
+    )
+    connector_type: str = CONNECTOR_TYPE
+
+    def get_download_path(self, file_data: FileData) -> Path:
+        rel_path = file_data.source_identifiers.relative_path
+        rel_path = rel_path[1:] if rel_path.startswith("/") else rel_path
+        return self.download_dir / Path(rel_path)
+
+    @SourceConnectionNetworkError.wrap
+    def _get_content(self, downloader: "MediaIoBaseDownload") -> bool:
+        downloaded = False
+        while downloaded is False:
+            _, downloaded = downloader.next_chunk()
+        return downloaded
+
+    def _write_file(self, file_data: FileData, file_contents: io.BytesIO):
+        download_path = self.get_download_path(file_data=file_data)
+        download_path.parent.mkdir(parents=True, exist_ok=True)
+        logger.debug(f"writing {file_data.source_identifiers.fullpath} to {download_path}")
+        with open(download_path, "wb") as handler:
+            handler.write(file_contents.getbuffer())
+        return self.generate_download_response(file_data=file_data, download_path=download_path)
+
+    @requires_dependencies(["googleapiclient"], extras="google-drive")
+    def run(self, file_data: FileData, **kwargs: Any) -> download_responses:
+        from googleapiclient.http import MediaIoBaseDownload
+
+        logger.debug(f"fetching file: {file_data.source_identifiers.fullpath}")
+        mime_type = file_data.additional_metadata["mimeType"]
+        record_id = file_data.identifier
+        files_client = self.connection_config.get_files_service()
+        if mime_type.startswith("application/vnd.google-apps"):
+            export_mime = GOOGLE_DRIVE_EXPORT_TYPES.get(
+                self.meta.get("mimeType"),  # type: ignore
+            )
+            if not export_mime:
+                raise TypeError(
+                    f"File not supported. Name: {file_data.source_identifiers.filename} "
+                    f"ID: {record_id} "
+                    f"MimeType: {mime_type}"
+                )
+
+            request = files_client.export_media(
+                fileId=record_id,
+                mimeType=export_mime,
+            )
+        else:
+            request = files_client.get_media(fileId=record_id)
+
+        file_contents = io.BytesIO()
+        downloader = MediaIoBaseDownload(file_contents, request)
+        downloaded = self._get_content(downloader=downloader)
+        if not downloaded or not file_contents:
+            return []
+        return self._write_file(file_data=file_data, file_contents=file_contents)
+
+
+google_drive_source_entry = SourceRegistryEntry(
+    connection_config=GoogleDriveConnectionConfig,
+    indexer_config=GoogleDriveIndexerConfig,
+    indexer=GoogleDriveIndexer,
+    downloader_config=GoogleDriveDownloaderConfig,
+    downloader=GoogleDriveDownloader,
+)

unstructured_ingest/v2/processes/connectors/local.py

@@ -0,0 +1,204 @@
+import glob
+import itertools
+import shutil
+from dataclasses import dataclass, field
+from pathlib import Path
+from time import time
+from typing import Any, Generator, Optional
+
+from unstructured.documents.elements import DataSourceMetadata
+
+from unstructured_ingest.v2.interfaces import (
+    AccessConfig,
+    ConnectionConfig,
+    Downloader,
+    DownloaderConfig,
+    DownloadResponse,
+    FileData,
+    Indexer,
+    IndexerConfig,
+    SourceIdentifiers,
+    UploadContent,
+    Uploader,
+    UploaderConfig,
+)
+from unstructured_ingest.v2.logger import logger
+from unstructured_ingest.v2.processes.connector_registry import (
+    DestinationRegistryEntry,
+    SourceRegistryEntry,
+)
+
+CONNECTOR_TYPE = "local"
+
+
+@dataclass
+class LocalAccessConfig(AccessConfig):
+    pass
+
+
+@dataclass
+class LocalConnectionConfig(ConnectionConfig):
+    access_config: LocalAccessConfig = field(default_factory=lambda: LocalAccessConfig())
+
+
+@dataclass
+class LocalIndexerConfig(IndexerConfig):
+    input_path: str
+    recursive: bool = False
+    file_glob: Optional[list[str]] = None
+
+    @property
+    def path(self) -> Path:
+        return Path(self.input_path).resolve()
+
+
+@dataclass
+class LocalIndexer(Indexer):
+    index_config: LocalIndexerConfig
+    connection_config: LocalConnectionConfig = field(
+        default_factory=lambda: LocalConnectionConfig()
+    )
+    connector_type: str = CONNECTOR_TYPE
+
+    def list_files(self) -> list[Path]:
+        input_path = self.index_config.path
+        if input_path.is_file():
+            return [Path(s) for s in glob.glob(f"{self.index_config.path}")]
+        glob_fn = input_path.rglob if self.index_config.recursive else input_path.glob
+        if not self.index_config.file_glob:
+            return list(glob_fn("*"))
+        return list(
+            itertools.chain.from_iterable(
+                glob_fn(pattern) for pattern in self.index_config.file_glob
+            )
+        )
+
+    def get_file_metadata(self, path: Path) -> DataSourceMetadata:
+        stats = path.stat()
+        try:
+            date_modified = str(stats.st_mtime)
+        except Exception as e:
+            logger.warning(f"Couldn't detect date modified: {e}")
+            date_modified = None
+
+        try:
+            date_created = str(stats.st_birthtime)
+        except Exception as e:
+            logger.warning(f"Couldn't detect date created: {e}")
+            date_created = None
+
+        try:
+            mode = stats.st_mode
+            permissions_data = [{"mode": mode}]
+        except Exception as e:
+            logger.warning(f"Couldn't detect file mode: {e}")
+            permissions_data = None
+        return DataSourceMetadata(
+            date_modified=date_modified,
+            date_created=date_created,
+            date_processed=str(time()),
+            permissions_data=permissions_data,
+            record_locator={"path": str(path.resolve())},
+        )
+
+    def run(self, **kwargs: Any) -> Generator[FileData, None, None]:
+        for file_path in self.list_files():
+            file_data = FileData(
+                identifier=str(file_path.resolve()),
+                connector_type=CONNECTOR_TYPE,
+                source_identifiers=SourceIdentifiers(
+                    fullpath=str(file_path.resolve()),
+                    filename=file_path.name,
+                    rel_path=(
+                        str(file_path.resolve()).replace(str(self.index_config.path.resolve()), "")[
+                            1:
+                        ]
+                        if not self.index_config.path.is_file()
+                        else self.index_config.path.name
+                    ),
+                ),
+                metadata=self.get_file_metadata(path=file_path),
+            )
+            yield file_data
+
+
+@dataclass
+class LocalDownloaderConfig(DownloaderConfig):
+    pass
+
+
+@dataclass
+class LocalDownloader(Downloader):
+    connector_type: str = CONNECTOR_TYPE
+    connection_config: LocalConnectionConfig = field(
+        default_factory=lambda: LocalConnectionConfig()
+    )
+    download_config: LocalDownloaderConfig = field(default_factory=lambda: LocalDownloaderConfig())
+
+    def get_download_path(self, file_data: FileData) -> Path:
+        return Path(file_data.source_identifiers.fullpath)
+
+    def run(self, file_data: FileData, **kwargs: Any) -> DownloadResponse:
+        return DownloadResponse(
+            file_data=file_data, path=Path(file_data.source_identifiers.fullpath)
+        )
+
+
+@dataclass
+class LocalUploaderConfig(UploaderConfig):
+    output_dir: str = field(default="structured-output")
+
+    @property
+    def output_path(self) -> Path:
+        return Path(self.output_dir).resolve()
+
+    def __post_init__(self):
+        if self.output_path.exists() and self.output_path.is_file():
+            raise ValueError("output path already exists as a file")
+
+
+@dataclass
+class LocalUploader(Uploader):
+    connector_type: str = CONNECTOR_TYPE
+    upload_config: LocalUploaderConfig = field(default_factory=lambda: LocalUploaderConfig())
+    connection_config: LocalConnectionConfig = field(
+        default_factory=lambda: LocalConnectionConfig()
+    )
+
+    def is_async(self) -> bool:
+        return False
+
+    def run(self, contents: list[UploadContent], **kwargs: Any) -> None:
+        self.upload_config.output_path.mkdir(parents=True, exist_ok=True)
+        for content in contents:
+            if source_identifiers := content.file_data.source_identifiers:
+                identifiers = source_identifiers
+                rel_path = (
+                    identifiers.relative_path[1:]
+                    if identifiers.relative_path.startswith("/")
+                    else identifiers.relative_path
+                )
+                new_path = self.upload_config.output_path / Path(rel_path)
+                final_path = str(new_path).replace(
+                    identifiers.filename, f"{identifiers.filename}.json"
+                )
+            else:
+                final_path = self.upload_config.output_path / Path(
+                    f"{content.file_data.identifier}.json"
+                )
+            Path(final_path).parent.mkdir(parents=True, exist_ok=True)
+            logger.debug(f"copying file from {content.path} to {final_path}")
+            shutil.copy(src=str(content.path), dst=str(final_path))
+
+
+local_source_entry = SourceRegistryEntry(
+    indexer=LocalIndexer,
+    indexer_config=LocalIndexerConfig,
+    downloader=LocalDownloader,
+    downloader_config=LocalDownloaderConfig,
+    connection_config=LocalConnectionConfig,
+)
+
+local_destination_entry = DestinationRegistryEntry(
+    uploader=LocalUploader, uploader_config=LocalUploaderConfig
+)

unstructured_ingest/v2/processes/connectors/mongodb.py

@@ -0,0 +1,138 @@
+import json
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Optional
+
+from unstructured.__version__ import __version__ as unstructured_version
+
+from unstructured_ingest.enhanced_dataclass import enhanced_field
+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,
+    FileData,
+    UploadContent,
+    Uploader,
+    UploaderConfig,
+    UploadStager,
+    UploadStagerConfig,
+)
+from unstructured_ingest.v2.logger import logger
+from unstructured_ingest.v2.processes.connector_registry import (
+    DestinationRegistryEntry,
+)
+
+if TYPE_CHECKING:
+    from pymongo import MongoClient
+
+CONNECTOR_TYPE = "mongodb"
+SERVER_API_VERSION = "1"
+
+
+@dataclass
+class MongoDBAccessConfig(AccessConfig):
+    uri: Optional[str] = None
+
+
+@dataclass
+class MongoDBConnectionConfig(ConnectionConfig):
+    access_config: MongoDBAccessConfig = enhanced_field(
+        sensitive=True, default_factory=MongoDBAccessConfig
+    )
+    host: Optional[str] = None
+    database: Optional[str] = None
+    collection: Optional[str] = None
+    port: int = 27017
+    batch_size: int = 100
+    connector_type: str = CONNECTOR_TYPE
+
+
+@dataclass
+class MongoDBUploadStagerConfig(UploadStagerConfig):
+    pass
+
+
+@dataclass
+class MongoDBUploadStager(UploadStager):
+    upload_stager_config: MongoDBUploadStagerConfig = field(
+        default_factory=lambda: MongoDBUploadStagerConfig()
+    )
+
+    def run(
+        self,
+        elements_filepath: Path,
+        file_data: FileData,
+        output_dir: Path,
+        output_filename: str,
+        **kwargs: Any,
+    ) -> Path:
+        with open(elements_filepath) as elements_file:
+            elements_contents = json.load(elements_file)
+
+        output_path = Path(output_dir) / Path(f"{output_filename}.json")
+        with open(output_path, "w") as output_file:
+            json.dump(elements_contents, output_file)
+        return output_path
+
+
+@dataclass
+class MongoDBUploaderConfig(UploaderConfig):
+    batch_size: int = 100
+
+
+@dataclass
+class MongoDBUploader(Uploader):
+    upload_config: MongoDBUploaderConfig
+    connection_config: MongoDBConnectionConfig
+    client: Optional["MongoClient"] = field(init=False)
+    connector_type: str = CONNECTOR_TYPE
+
+    def __post_init__(self):
+        self.client = self.create_client()
+
+    @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
+
+        if self.connection_config.access_config.uri:
+            return MongoClient(
+                self.connection_config.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, contents: list[UploadContent], **kwargs: Any) -> None:
+        elements_dict = []
+        for content in contents:
+            with open(content.path) as elements_file:
+                elements = json.load(elements_file)
+                elements_dict.extend(elements)
+
+        logger.info(
+            f"writing {len(elements_dict)} objects to destination "
+            f"db, {self.connection_config.database}, "
+            f"collection {self.connection_config.collection} "
+            f"at {self.connection_config.host}",
+        )
+        db = self.client[self.connection_config.database]
+        collection = db[self.connection_config.collection]
+        for chunk in batch_generator(elements_dict, self.upload_config.batch_size):
+            collection.insert_many(chunk)
+
+
+mongodb_destination_entry = DestinationRegistryEntry(
+    connection_config=MongoDBConnectionConfig,
+    uploader=MongoDBUploader,
+    uploader_config=MongoDBUploaderConfig,
+    upload_stager=MongoDBUploadStager,
+    upload_stager_config=MongoDBUploadStagerConfig,
+)