unstructured-ingest 0.0.0__py3-none-any.whl

unstructured_ingest/connector/google_drive.py

@@ -0,0 +1,349 @@
+import io
+import json
+import os
+import typing as t
+from dataclasses import dataclass, field
+from datetime import datetime
+from mimetypes import guess_extension
+from pathlib import Path
+
+from unstructured.file_utils.google_filetype import GOOGLE_DRIVE_EXPORT_TYPES
+
+from unstructured_ingest.enhanced_dataclass import enhanced_field
+from unstructured_ingest.error import SourceConnectionError, SourceConnectionNetworkError
+from unstructured_ingest.interfaces import (
+    AccessConfig,
+    BaseConnectorConfig,
+    BaseSessionHandle,
+    BaseSingleIngestDoc,
+    BaseSourceConnector,
+    ConfigSessionHandleMixin,
+    IngestDocCleanupMixin,
+    IngestDocSessionHandleMixin,
+    SourceConnectorCleanupMixin,
+    SourceMetadata,
+)
+from unstructured_ingest.logger import logger
+from unstructured_ingest.utils.dep_check import requires_dependencies
+from unstructured_ingest.utils.string_and_date_utils import json_to_dict
+
+if t.TYPE_CHECKING:
+    from googleapiclient.discovery import Resource as GoogleAPIResource
+    from googleapiclient.http import MediaIoBaseDownload
+
+FILE_FORMAT = "{id}-{name}{ext}"
+DIRECTORY_FORMAT = "{id}-{name}"
+
+
+@dataclass
+class GoogleDriveSessionHandle(BaseSessionHandle):
+    service: "GoogleAPIResource"
+
+
+@requires_dependencies(["googleapiclient"], extras="google-drive")
+def create_service_account_object(key_path: t.Union[str, dict], id=None):
+    """
+    Creates a service object for interacting with Google Drive.
+
+    Providing a drive id enforces a key validation process.
+
+    Args:
+        key_path: Path to Google Drive service account json file. (or the actual json)
+        id: ID of a file on Google Drive. File has to be either publicly accessible or accessible
+            to the service account.
+
+    Returns:
+        Service account object
+    """
+    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
+    key_path = json_to_dict(key_path)
+
+    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)
+
+        if id:
+            service.files().list(
+                spaces="drive",
+                fields="files(id)",
+                pageToken=None,
+                corpora="user",
+                q=f"'{id}' in parents",
+            ).execute()
+
+    except HttpError as exc:
+        raise ValueError(f"{exc.reason}")
+    except exceptions.DefaultCredentialsError:
+        raise ValueError("The provided API key is invalid.")
+
+    return service
+
+
+@dataclass
+class GoogleDriveAccessConfig(AccessConfig):
+    service_account_key: t.Union[str, dict] = enhanced_field(sensitive=True)
+
+
+@dataclass
+class SimpleGoogleDriveConfig(ConfigSessionHandleMixin, BaseConnectorConfig):
+    """Connector config where drive_id is the id of the document to process or
+    the folder to process all documents from."""
+
+    # Google Drive Specific Options
+    drive_id: str
+    access_config: GoogleDriveAccessConfig
+    extension: t.Optional[str] = None
+    recursive: bool = False
+
+    def create_session_handle(
+        self,
+    ) -> GoogleDriveSessionHandle:
+        service = create_service_account_object(self.access_config.service_account_key)
+        return GoogleDriveSessionHandle(service=service)
+
+
+@dataclass
+class GoogleDriveIngestDoc(IngestDocSessionHandleMixin, IngestDocCleanupMixin, BaseSingleIngestDoc):
+    connector_config: SimpleGoogleDriveConfig
+    meta: t.Dict[str, str] = field(default_factory=dict)
+    registry_name: str = "google_drive"
+
+    @property
+    def filename(self):
+        return Path(self.meta.get("download_filepath")).resolve()  # type: ignore
+
+    @property
+    def _output_filename(self):
+        return Path(f"{self.meta.get('output_filepath')}.json").resolve()
+
+    @property
+    def record_locator(self) -> t.Optional[t.Dict[str, t.Any]]:
+        return {
+            "drive_id": self.connector_config.drive_id,
+            "file_id": self.meta["id"],
+        }
+
+    @requires_dependencies(["googleapiclient"], extras="google-drive")
+    def update_source_metadata(self):
+        from googleapiclient.errors import HttpError
+
+        try:
+            file_obj = (
+                self.session_handle.service.files()
+                .get(
+                    fileId=self.meta["id"],
+                    fields="id, createdTime, modifiedTime, version, webContentLink",
+                )
+                .execute()
+            )
+        except HttpError as e:
+            if e.status_code == 404:
+                logger.error(f"File {self.meta['name']} not found")
+                self.source_metadata = SourceMetadata(
+                    exists=True,
+                )
+                return
+            raise
+
+        date_created = None
+        if dc := file_obj.get("createdTime", ""):
+            date_created = datetime.strptime(
+                dc,
+                "%Y-%m-%dT%H:%M:%S.%fZ",
+            ).isoformat()
+
+        date_modified = None
+        if dm := file_obj.get("modifiedTime", ""):
+            date_modified = datetime.strptime(
+                dm,
+                "%Y-%m-%dT%H:%M:%S.%fZ",
+            ).isoformat()
+
+        self.source_metadata = SourceMetadata(
+            date_created=date_created,
+            date_modified=date_modified,
+            version=file_obj.get("version", ""),
+            source_url=file_obj.get("webContentLink", ""),
+            exists=True,
+        )
+
+    @SourceConnectionNetworkError.wrap
+    def _run_downloader(self, downloader: "MediaIoBaseDownload") -> bool:
+        downloaded = False
+        while downloaded is False:
+            _, downloaded = downloader.next_chunk()
+        return downloaded
+
+    @requires_dependencies(["googleapiclient"], extras="google-drive")
+    @SourceConnectionError.wrap
+    @BaseSingleIngestDoc.skip_if_file_exists
+    def get_file(self):
+        from googleapiclient.http import MediaIoBaseDownload
+
+        if self.meta.get("mimeType", "").startswith("application/vnd.google-apps"):
+            export_mime = GOOGLE_DRIVE_EXPORT_TYPES.get(
+                self.meta.get("mimeType"),  # type: ignore
+            )
+            if not export_mime:
+                logger.info(
+                    f"File not supported. Name: {self.meta.get('name')} "
+                    f"ID: {self.meta.get('id')} "
+                    f"MimeType: {self.meta.get('mimeType')}",
+                )
+                return
+
+            request = self.session_handle.service.files().export_media(
+                fileId=self.meta.get("id"),
+                mimeType=export_mime,
+            )
+        else:
+            request = self.session_handle.service.files().get_media(fileId=self.meta.get("id"))
+        file = io.BytesIO()
+        downloader = MediaIoBaseDownload(file, request)
+        self.update_source_metadata()
+        downloaded = self._run_downloader(downloader=downloader)
+
+        saved = False
+        if downloaded and file:
+            dir_ = Path(self.meta["download_dir"])
+            if dir_:
+                if not dir_.is_dir():
+                    logger.debug(f"Creating directory: {self.meta.get('download_dir')}")
+
+                    if dir_:
+                        dir_.mkdir(parents=True, exist_ok=True)
+
+                with open(self.filename, "wb") as handler:
+                    handler.write(file.getbuffer())
+                    saved = True
+                    logger.debug(f"File downloaded: {self.filename}.")
+        if not saved:
+            logger.error(f"Error while downloading and saving file: {self.filename}.")
+
+    def write_result(self):
+        """Write the structured json result for this doc. result must be json serializable."""
+        if self.read_config.download_only:
+            return
+        self._output_filename.parent.mkdir(parents=True, exist_ok=True)
+        with open(self._output_filename, "w") as output_f:
+            output_f.write(json.dumps(self.isd_elems_no_filename, ensure_ascii=False, indent=2))
+        logger.info(f"Wrote {self._output_filename}")
+
+
+@dataclass
+class GoogleDriveSourceConnector(SourceConnectorCleanupMixin, BaseSourceConnector):
+    """Objects of this class support fetching documents from Google Drive"""
+
+    connector_config: SimpleGoogleDriveConfig
+
+    def _list_objects(self, drive_id, recursive=False):
+        files = []
+        service = self.connector_config.create_session_handle().service
+
+        def traverse(drive_id, download_dir, output_dir, recursive=False):
+            page_token = None
+            while True:
+                response = (
+                    service.files()
+                    .list(
+                        spaces="drive",
+                        fields="nextPageToken, files(id, name, mimeType)",
+                        pageToken=page_token,
+                        corpora="user",
+                        q=f"'{drive_id}' in parents",
+                    )
+                    .execute()
+                )
+
+                for meta in response.get("files", []):
+                    if meta.get("mimeType") == "application/vnd.google-apps.folder":
+                        dir_ = DIRECTORY_FORMAT.format(name=meta.get("name"), id=meta.get("id"))
+                        if recursive:
+                            download_sub_dir = (download_dir / dir_).resolve()
+                            output_sub_dir = (output_dir / dir_).resolve()
+                            traverse(meta.get("id"), download_sub_dir, output_sub_dir, True)
+                    else:
+                        ext = ""
+                        if not Path(meta.get("name")).suffixes:
+                            guess = guess_extension(meta.get("mimeType"))
+                            ext = guess if guess else ext
+
+                        if meta.get("mimeType", "").startswith("application/vnd.google-apps"):
+                            export_mime = GOOGLE_DRIVE_EXPORT_TYPES.get(meta.get("mimeType"))
+                            if not export_mime:
+                                logger.info(
+                                    f"File {meta.get('name')} has an "
+                                    f"unsupported MimeType {meta.get('mimeType')}",
+                                )
+                                continue
+
+                            if not ext:
+                                guess = guess_extension(export_mime)
+                                ext = guess if guess else ext
+
+                        # TODO (Habeeb): Consider filtering at the query level.
+                        if (
+                            self.connector_config.extension
+                            and self.connector_config.extension != ext
+                        ):  # noqa: SIM102
+                            logger.debug(
+                                f"File {meta.get('name')} does not match "
+                                f"the file type {self.connector_config.extension}",
+                            )
+                            continue
+
+                        name = FILE_FORMAT.format(name=meta.get("name"), id=meta.get("id"), ext=ext)
+                        meta["download_dir"] = str(download_dir)
+                        meta["download_filepath"] = (download_dir / name).resolve().as_posix()
+                        meta["output_dir"] = str(output_dir)
+                        meta["output_filepath"] = (output_dir / name).resolve().as_posix()
+                        files.append(meta)
+
+                page_token = response.get("nextPageToken", None)
+                if page_token is None:
+                    break
+
+        traverse(
+            drive_id,
+            Path(self.read_config.download_dir),
+            Path(self.processor_config.output_dir),
+            recursive,
+        )
+        return files
+
+    def initialize(self):
+        pass
+
+    def check_connection(self):
+        try:
+            self.connector_config.create_session_handle().service
+        except Exception as e:
+            logger.error(f"failed to validate connection: {e}", exc_info=True)
+            raise SourceConnectionError(f"failed to validate connection: {e}")
+
+    def get_ingest_docs(self):
+        files = self._list_objects(self.connector_config.drive_id, self.connector_config.recursive)
+        return [
+            GoogleDriveIngestDoc(
+                connector_config=self.connector_config,
+                processor_config=self.processor_config,
+                read_config=self.read_config,
+                meta=file,
+            )
+            for file in files
+        ]

unstructured_ingest/connector/hubspot.py

@@ -0,0 +1,278 @@
+import typing as t
+from dataclasses import dataclass
+from enum import Enum
+from functools import reduce
+from pathlib import Path
+
+from unstructured_ingest.enhanced_dataclass import enhanced_field
+from unstructured_ingest.error import SourceConnectionError
+from unstructured_ingest.interfaces import (
+    AccessConfig,
+    BaseConnectorConfig,
+    BaseSessionHandle,
+    BaseSingleIngestDoc,
+    BaseSourceConnector,
+    ConfigSessionHandleMixin,
+    IngestDocCleanupMixin,
+    IngestDocSessionHandleMixin,
+    SourceConnectorCleanupMixin,
+    SourceMetadata,
+)
+from unstructured_ingest.logger import logger
+from unstructured_ingest.utils.dep_check import requires_dependencies
+
+if t.TYPE_CHECKING:
+    from hubspot import HubSpot
+
+CONTENT_TAG = "content"
+
+
+class HubSpotObjectTypes(Enum):
+    CALLS = "calls"
+    COMMUNICATIONS = "communications"
+    EMAILS = "emails"
+    NOTES = "notes"
+    PRODUCTS = "products"
+    TICKETS = "tickets"
+
+
+@dataclass
+class HubSpotSessionHandle(BaseSessionHandle):
+    service: "HubSpot"
+
+
+@dataclass
+class HubSpotAccessConfig(AccessConfig):
+    api_token: str = enhanced_field(repr=False, sensitive=True)
+
+
+@dataclass
+class SimpleHubSpotConfig(ConfigSessionHandleMixin, BaseConnectorConfig):
+    access_config: HubSpotAccessConfig
+    params: t.Optional[str] = None
+    properties: t.Optional[dict] = None
+    object_types: t.Optional[t.List[str]] = None
+    custom_properties: t.Optional[t.Dict[str, t.List[str]]] = None
+
+    @requires_dependencies(["hubspot"], extras="hubspot")
+    def create_session_handle(self) -> HubSpotSessionHandle:
+        from hubspot import HubSpot
+
+        service = HubSpot(access_token=self.access_config.api_token)
+        return HubSpotSessionHandle(service=service)
+
+
+@dataclass
+class HubSpotIngestDoc(IngestDocSessionHandleMixin, IngestDocCleanupMixin, BaseSingleIngestDoc):
+    connector_config: SimpleHubSpotConfig
+    object_id: str
+    object_type: str
+    content_properties: t.List[str]
+    registry_name: str = "hubspot"
+
+    def __post_init__(self):
+        self._add_custom_properties()
+
+    @property
+    def filename(self):
+        return (
+            Path(self.read_config.download_dir)
+            / f"{self.object_type}/{self.object_id}.txt"  # type: ignore
+        ).resolve()
+
+    @property
+    def _output_filename(self):
+        return (
+            Path(self.processor_config.output_dir)
+            / f"{self.object_type}/{self.object_id}.json"  # type: ignore
+        ).resolve()
+
+    @property
+    def record_locator(self) -> t.Optional[t.Dict[str, t.Any]]:
+        return {
+            f"{self.registry_name}_id": self.object_id,
+        }
+
+    @property
+    def version(self) -> t.Optional[str]:
+        return None
+
+    @property
+    def source_url(self) -> t.Optional[str]:
+        return None
+
+    def _add_custom_properties(self):
+        if (self.connector_config.custom_properties is not None) and (
+            (cprops := self.connector_config.custom_properties.get(self.object_type)) is not None
+        ):
+            self.content_properties += cprops
+
+    def _join_object_properties(self, obj) -> str:
+        return "\n".join(
+            [
+                obj.properties[cprop]
+                for cprop in self.content_properties
+                if (obj.properties.get(cprop) is not None)
+            ],
+        )
+
+    def _resolve_getter(self):
+        method_path = ""
+        if self.object_type in [
+            HubSpotObjectTypes.CALLS.value,
+            HubSpotObjectTypes.COMMUNICATIONS.value,
+            HubSpotObjectTypes.EMAILS.value,
+            HubSpotObjectTypes.NOTES.value,
+        ]:
+            method_path = f"crm.objects.{self.object_type}.basic_api.get_by_id"
+        if self.object_type in [
+            HubSpotObjectTypes.PRODUCTS.value,
+            HubSpotObjectTypes.TICKETS.value,
+        ]:
+            method_path = f"crm.{self.object_type}.basic_api.get_by_id"
+
+        method = reduce(getattr, method_path.split("."), self.session_handle.service)
+        return method
+
+    @requires_dependencies(["hubspot"], extras="hubspot")
+    def _fetch_obj(self, check_only=False):
+        from hubspot.crm.objects.exceptions import NotFoundException
+
+        get_by_id_method = self._resolve_getter()
+        try:
+            response = get_by_id_method(
+                self.object_id,
+                properties=([] if check_only else self.content_properties),
+            )
+        except NotFoundException as e:
+            logger.error(e)
+            return None
+        return response
+
+    def update_source_metadata(self, **kwargs) -> None:
+        obj = kwargs.get("object", self._fetch_obj(check_only=True))  # type: ignore
+        if obj is None:
+            self.source_metadata = SourceMetadata(
+                exists=False,
+            )
+            return
+        self.source_metadata = SourceMetadata(
+            date_created=obj.created_at.isoformat(),
+            date_modified=obj.updated_at.isoformat(),
+            exists=True,
+        )
+
+    @SourceConnectionError.wrap
+    @BaseSingleIngestDoc.skip_if_file_exists
+    def get_file(self):
+        obj = self._fetch_obj()
+        if obj is None:
+            raise ValueError(
+                f"Failed to retrieve object {self.registry_name}",
+                f"with ID {self.object_id}",
+            )
+        self.update_source_metadata(object=obj)
+        output = self._join_object_properties(obj)
+        self.filename.parent.mkdir(parents=True, exist_ok=True)
+        with open(self.filename, "w", encoding="utf8") as f:
+            f.write(output)
+        return
+
+
+@dataclass
+class HubSpotSourceConnector(SourceConnectorCleanupMixin, BaseSourceConnector):
+    connector_config: SimpleHubSpotConfig
+
+    def initialize(self):
+        self.hubspot = self.connector_config.create_session_handle().service
+
+    def check_connection(self):
+        return self.connector_config.create_session_handle().service
+
+    @requires_dependencies(["hubspot"], extras="hubspot")
+    def _list_objects(self, get_page_method, object_type: str, content_properties: t.List[str]):
+        try:
+            objects = get_page_method()
+        except Exception as e:
+            logger.error(e)
+            logger.error(
+                f"Failed to retrieve {object_type}, omitting processing...",
+            )
+            return []
+        return [
+            HubSpotIngestDoc(
+                connector_config=self.connector_config,
+                processor_config=self.processor_config,
+                read_config=self.read_config,
+                object_id=obj.id,
+                object_type=object_type,
+                content_properties=content_properties,
+            )
+            for obj in objects.results
+        ]
+
+    def _get_calls(self) -> t.List[HubSpotIngestDoc]:
+        return self._list_objects(
+            self.hubspot.crm.objects.calls.basic_api.get_page,
+            HubSpotObjectTypes.CALLS.value,
+            ["hs_call_title", "hs_call_body"],
+        )
+
+    def _get_communications(self) -> t.List[HubSpotIngestDoc]:
+        return self._list_objects(
+            self.hubspot.crm.objects.communications.basic_api.get_page,
+            HubSpotObjectTypes.COMMUNICATIONS.value,
+            ["hs_communication_body"],
+        )
+
+    def _get_emails(self) -> t.List[HubSpotIngestDoc]:
+        return self._list_objects(
+            self.hubspot.crm.objects.emails.basic_api.get_page,
+            HubSpotObjectTypes.EMAILS.value,
+            ["hs_email_subject", "hs_email_text"],
+        )
+
+    def _get_notes(self) -> t.List[HubSpotIngestDoc]:
+        return self._list_objects(
+            self.hubspot.crm.objects.notes.basic_api.get_page,
+            HubSpotObjectTypes.NOTES.value,
+            ["hs_note_body"],
+        )
+
+    def _get_products(self) -> t.List[HubSpotIngestDoc]:
+        return self._list_objects(
+            self.hubspot.crm.products.basic_api.get_page,
+            HubSpotObjectTypes.PRODUCTS.value,
+            ["description"],
+        )
+
+    def _get_tickets(self) -> t.List[HubSpotIngestDoc]:
+        return self._list_objects(
+            self.hubspot.crm.tickets.basic_api.get_page,
+            HubSpotObjectTypes.TICKETS.value,
+            ["subject", "content"],
+        )
+
+    def get_ingest_docs(self):
+        obj_method_resolver = {
+            HubSpotObjectTypes.CALLS.value: self._get_calls,
+            HubSpotObjectTypes.COMMUNICATIONS.value: self._get_communications,
+            HubSpotObjectTypes.EMAILS.value: self._get_emails,
+            HubSpotObjectTypes.NOTES.value: self._get_notes,
+            HubSpotObjectTypes.PRODUCTS.value: self._get_products,
+            HubSpotObjectTypes.TICKETS.value: self._get_tickets,
+        }
+
+        if self.connector_config.object_types is not None:
+            obj_method_resolver = {
+                obj_name: obj_method_resolver.get(obj_name)  # type: ignore
+                for obj_name in self.connector_config.object_types
+            }
+
+        ingest_docs: t.List[HubSpotIngestDoc] = []
+        for obj_name, obj_method in obj_method_resolver.items():
+            logger.info(f"Retrieving - {obj_name}")
+            results: t.List[HubSpotIngestDoc] = obj_method()  # type: ignore
+            ingest_docs += results  # type: ignore
+
+        return ingest_docs