unstructured-ingest 0.0.0__py3-none-any.whl

unstructured_ingest/v2/processes/connectors/salesforce.py

@@ -0,0 +1,293 @@
+"""
+Salesforce Connector
+Able to download Account, Case, Campaign, EmailMessage, Lead
+Salesforce returns everything as a list of json.
+This saves each entry as a separate file to be partitioned.
+Using JWT authorization
+https://developer.salesforce.com/docs/atlas.en-us.sfdx_dev.meta/sfdx_dev/sfdx_dev_auth_key_and_cert.htm
+https://developer.salesforce.com/docs/atlas.en-us.sfdx_dev.meta/sfdx_dev/sfdx_dev_auth_connected_app.htm
+"""
+
+import json
+from collections import OrderedDict
+from dataclasses import dataclass, field
+from email.utils import formatdate
+from pathlib import Path
+from string import Template
+from textwrap import dedent
+from typing import TYPE_CHECKING, Any, Generator, Type
+
+from dateutil import parser
+from unstructured.documents.elements import DataSourceMetadata
+
+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.v2.interfaces import (
+    AccessConfig,
+    ConnectionConfig,
+    Downloader,
+    DownloaderConfig,
+    DownloadResponse,
+    FileData,
+    Indexer,
+    IndexerConfig,
+    SourceIdentifiers,
+)
+from unstructured_ingest.v2.logger import logger
+from unstructured_ingest.v2.processes.connector_registry import (
+    SourceRegistryEntry,
+)
+
+
+class MissingCategoryError(Exception):
+    """There are no categories with that name."""
+
+
+CONNECTOR_TYPE = "salesforce"
+
+if TYPE_CHECKING:
+    from simple_salesforce import Salesforce
+
+SALESFORCE_API_VERSION = "57.0"
+
+# TODO: Add more categories as needed
+ACCEPTED_CATEGORIES: list[str] = ["Account", "Case", "Campaign", "EmailMessage", "Lead"]
+
+# Generic minimal email template used only
+# to process EmailMessage records as .eml files
+EMAIL_TEMPLATE = Template(
+    """MIME-Version: 1.0
+Date: $date
+Message-ID: $message_identifier
+Subject: $subject
+From: $from_email
+To: $to_email
+Content-Type: multipart/alternative; boundary="00000000000095c9b205eff92630"
+--00000000000095c9b205eff92630
+Content-Type: text/plain; charset="UTF-8"
+$textbody
+--00000000000095c9b205eff92630
+Content-Type: text/html; charset="UTF-8"
+$htmlbody
+--00000000000095c9b205eff92630--
+""",
+)
+
+
+@dataclass
+class SalesforceAccessConfig(AccessConfig):
+    consumer_key: str
+    private_key: str
+
+    @requires_dependencies(["cryptography"])
+    def get_private_key_value_and_type(self) -> tuple[str, Type]:
+        from cryptography.hazmat.primitives import serialization
+
+        try:
+            serialization.load_pem_private_key(data=self.private_key.encode("utf-8"), password=None)
+        except ValueError:
+            pass
+        else:
+            return self.private_key, str
+
+        if Path(self.private_key).is_file():
+            return self.private_key, Path
+
+        raise ValueError("private_key does not contain PEM private key or path")
+
+
+@dataclass
+class SalesforceConnectionConfig(ConnectionConfig):
+    username: str
+    access_config: SalesforceAccessConfig = enhanced_field(sensitive=True)
+
+    @requires_dependencies(["simple_salesforce"], extras="salesforce")
+    def get_client(self) -> "Salesforce":
+        from simple_salesforce import Salesforce
+
+        pkey_value, pkey_type = self.access_config.get_private_key_value_and_type()
+
+        return Salesforce(
+            username=self.username,
+            consumer_key=self.access_config.consumer_key,
+            privatekey_file=pkey_value if pkey_type is Path else None,
+            privatekey=pkey_value if pkey_type is str else None,
+            version=SALESFORCE_API_VERSION,
+        )
+
+
+@dataclass
+class SalesforceIndexerConfig(IndexerConfig):
+    categories: list[str]
+
+
+@dataclass
+class SalesforceIndexer(Indexer):
+    connection_config: SalesforceConnectionConfig
+    index_config: SalesforceIndexerConfig
+
+    def __post_init__(self):
+        for record_type in self.index_config.categories:
+            if record_type not in ACCEPTED_CATEGORIES:
+                raise ValueError(f"{record_type} not currently an accepted Salesforce category")
+
+    def get_file_extension(self, record_type) -> str:
+        if record_type == "EmailMessage":
+            extension = ".eml"
+        elif record_type in ["Account", "Lead", "Case", "Campaign"]:
+            extension = ".xml"
+        else:
+            raise MissingCategoryError(
+                f"There are no categories with the name: {record_type}",
+            )
+        return extension
+
+    @requires_dependencies(["simple_salesforce"], extras="salesforce")
+    def list_files(self) -> list[FileData]:
+        """Get Salesforce Ids for the records.
+        Send them to next phase where each doc gets downloaded into the
+        appropriate format for partitioning.
+        """
+        from simple_salesforce.exceptions import SalesforceMalformedRequest
+
+        client = self.connection_config.get_client()
+
+        files_list = []
+        for record_type in self.index_config.categories:
+            try:
+                # Get ids from Salesforce
+                records = client.query_all_iter(
+                    f"select Id, SystemModstamp, CreatedDate, LastModifiedDate from {record_type}",
+                )
+                for record in records:
+                    record_with_extension = record["Id"] + self.get_file_extension(
+                        record["attributes"]["type"]
+                    )
+                    files_list.append(
+                        FileData(
+                            connector_type=CONNECTOR_TYPE,
+                            identifier=record["Id"],
+                            source_identifiers=SourceIdentifiers(
+                                filename=record_with_extension,
+                                fullpath=f"{record['attributes']['type']}/{record_with_extension}",
+                            ),
+                            metadata=DataSourceMetadata(
+                                url=record["attributes"]["url"],
+                                version=str(parser.parse(record["SystemModstamp"]).timestamp()),
+                                date_created=str(parser.parse(record["CreatedDate"]).timestamp()),
+                                date_modified=str(
+                                    parser.parse(record["LastModifiedDate"]).timestamp()
+                                ),
+                                record_locator={"id": record["Id"]},
+                            ),
+                            additional_metadata={"record_type": record["attributes"]["type"]},
+                        )
+                    )
+            except SalesforceMalformedRequest as e:
+                raise SalesforceMalformedRequest(f"Problem with Salesforce query: {e}")
+
+        return files_list
+
+    def run(self, **kwargs: Any) -> Generator[FileData, None, None]:
+        for f in self.list_files():
+            yield f
+
+
+@dataclass
+class SalesforceDownloaderConfig(DownloaderConfig):
+    pass
+
+
+@dataclass
+class SalesforceDownloader(Downloader):
+    connection_config: SalesforceConnectionConfig
+    download_config: SalesforceDownloaderConfig = field(
+        default_factory=lambda: SalesforceDownloaderConfig()
+    )
+    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)
+
+    def _xml_for_record(self, record: OrderedDict) -> str:
+        """Creates partitionable xml file from a record"""
+        import xml.etree.ElementTree as ET
+
+        def create_xml_doc(data, parent, prefix=""):
+            for key, value in data.items():
+                if isinstance(value, OrderedDict):
+                    create_xml_doc(value, parent, prefix=f"{prefix}{key}.")
+                else:
+                    item = ET.Element("item")
+                    item.text = f"{prefix}{key}: {value}"
+                    parent.append(item)
+
+        root = ET.Element("root")
+        create_xml_doc(record, root)
+
+        xml_string = ET.tostring(root, encoding="utf-8", xml_declaration=True).decode()
+        return xml_string
+
+    def _eml_for_record(self, email_json: dict[str, Any]) -> str:
+        """Recreates standard expected .eml format using template."""
+        eml = EMAIL_TEMPLATE.substitute(
+            date=formatdate(parser.parse(email_json.get("MessageDate")).timestamp()),
+            message_identifier=email_json.get("MessageIdentifier"),
+            subject=email_json.get("Subject"),
+            from_email=email_json.get("FromAddress"),
+            to_email=email_json.get("ToAddress"),
+            textbody=email_json.get("TextBody"),
+            htmlbody=email_json.get("HtmlBody"),
+        )
+        return dedent(eml)
+
+    @SourceConnectionNetworkError.wrap
+    def _get_response(self, file_data: FileData) -> OrderedDict:
+        client = self.connection_config.get_client()
+        return client.query(
+            f"select FIELDS(STANDARD) from {file_data.additional_metadata['record_type']} where Id='{file_data.identifier}'",  # noqa: E501
+        )
+
+    def get_record(self, file_data: FileData) -> OrderedDict:
+        # Get record from Salesforce based on id
+        response = self._get_response(file_data)
+        logger.debug(f"response was returned for salesforce record id: {file_data.identifier}")
+        records = response["records"]
+        if not records:
+            raise ValueError(
+                f"No record found with record id {file_data.identifier}: {json.dumps(response)}"
+            )
+        record_json = records[0]
+        return record_json
+
+    def run(self, file_data: FileData, **kwargs: Any) -> DownloadResponse:
+        record = self.get_record(file_data)
+
+        try:
+            if file_data.additional_metadata["record_type"] == "EmailMessage":
+                document = self._eml_for_record(record)
+            else:
+                document = self._xml_for_record(record)
+            download_path = self.get_download_path(file_data=file_data)
+            download_path.parent.mkdir(parents=True, exist_ok=True)
+
+            with open(download_path, "w") as page_file:
+                page_file.write(document)
+
+        except Exception as e:
+            logger.error(f"failed to download file {file_data.identifier}: {e}", exc_info=True)
+            raise SourceConnectionNetworkError(f"failed to download file {file_data.identifier}")
+
+        return self.generate_download_response(file_data=file_data, download_path=download_path)
+
+
+salesforce_source_entry = SourceRegistryEntry(
+    connection_config=SalesforceConnectionConfig,
+    indexer_config=SalesforceIndexerConfig,
+    indexer=SalesforceIndexer,
+    downloader_config=SalesforceDownloaderConfig,
+    downloader=SalesforceDownloader,
+)