arcade-mcp/toolkits/salesforce/arcade_salesforce/models.py

import asyncio
import os
from dataclasses import dataclass
from typing import Any, cast

import httpx

from arcade_salesforce.constants import MAX_CONCURRENT_REQUESTS, SALESFORCE_API_VERSION
from arcade_salesforce.enums import SalesforceObject
from arcade_salesforce.exceptions import (
    BadRequestError,
    ResourceNotFoundError,
    SalesforceToolExecutionError,
)
from arcade_salesforce.utils import (
    build_soql_query,
    clean_contact_data,
    clean_lead_data,
    clean_note_data,
    clean_object_data,
    clean_opportunity_data,
    clean_task_data,
    expand_associations,
    get_ids_referenced,
    get_object_type,
    remove_none_values,
)


@dataclass
class SalesforceClient:
    auth_token: str
    org_subdomain: str | None = None
    api_version: str = SALESFORCE_API_VERSION
    max_concurrent_requests: int = MAX_CONCURRENT_REQUESTS

    # Internal state properties
    _state_object_fields: dict[SalesforceObject, list[str]] | None = None
    _state_is_person_account_enabled: bool | None = None
    _semaphore: asyncio.Semaphore | None = None

    def __post_init__(self) -> None:
        if self.org_subdomain is None:
            self.org_subdomain = os.getenv("SALESFORCE_ORG_SUBDOMAIN")
        if self.org_subdomain is None:
            raise ValueError(
                "Either `org_subdomain` argument or `SALESFORCE_ORG_SUBDOMAIN` env var must be set"
            )

        if self._state_object_fields is None:
            self._state_object_fields = {}

        if self._semaphore is None:
            self._semaphore = asyncio.Semaphore(self.max_concurrent_requests)

    @property
    def _base_url(self) -> str:
        return f"https://{self.org_subdomain}.my.salesforce.com/services/data/{self.api_version}"

    @property
    def object_fields(self) -> dict[SalesforceObject, list[str]]:
        return cast(dict, self._state_object_fields)

    def _endpoint_url(self, endpoint: str) -> str:
        return f"{self._base_url}/{endpoint.lstrip('/')}"

    def _build_headers(self, headers: dict | None = None) -> dict:
        headers = headers or {}
        headers["Authorization"] = f"Bearer {self.auth_token}"
        return headers

    def _raise_salesforce_error(self, response: httpx.Response) -> None:
        """Raise a ToolExecutionError based on the Salesforce API response status code."""
        errors = [error.get("message") for error in response.json()]
        if response.status_code == 404:
            raise ResourceNotFoundError(errors)
        elif response.status_code == 400:
            raise BadRequestError(errors)
        raise SalesforceToolExecutionError(errors)

    async def get(
        self,
        endpoint: str,
        params: dict | None = None,
        headers: dict | None = None,
    ) -> dict:
        """Make a GET request to the Salesforce API.

        Args:
            endpoint: The Salesforce API endpoint to call.
            params: The query parameters to include in the request.
            headers: The headers to include in the request.

        Returns:
            The JSON-loaded representation of the response from the Salesforce API.
        """
        kwargs: dict[str, Any] = {
            "url": self._endpoint_url(endpoint),
            "headers": self._build_headers(headers),
        }
        if params:
            kwargs["params"] = params

        async with self._semaphore, httpx.AsyncClient() as client:  # type: ignore[union-attr]
            response = await client.get(**kwargs)

        if response.status_code >= 300:
            self._raise_salesforce_error(response)

        return cast(dict, response.json())

    async def post(
        self,
        endpoint: str,
        data: dict | None = None,
        json_data: dict | None = None,
        headers: dict | None = None,
    ) -> dict:
        """Make a POST request to the Salesforce API.

        Args:
            endpoint: The Salesforce API endpoint to call.
            data: The data for the request. (provide data or json_data, not both)
            json_data: The JSON data for the request. (provide data or json_data, not both)
            headers: The headers to include in the request.

        Returns:
            The JSON-loaded representation of the response from the Salesforce API.
        """
        kwargs: dict[str, Any] = {
            "url": self._endpoint_url(endpoint),
            "headers": self._build_headers(headers),
        }
        if data:
            kwargs["data"] = data
        if json_data:
            kwargs["json"] = json_data

        async with self._semaphore, httpx.AsyncClient() as client:  # type: ignore[union-attr]
            response = await client.post(**kwargs)

        if response.status_code >= 300:
            self._raise_salesforce_error(response)

        return cast(dict, response.json())

    async def get_object_fields(self, object_type: SalesforceObject) -> list[str]:
        """Get the fields available for a Salesforce object.

        Args:
            object_type: The Salesforce object to get the fields for.

        Returns:
            The list of fields available for a Salesforce object.
        """
        if object_type not in self.object_fields:
            response = await self._describe_object(object_type)
            self.object_fields[object_type] = [field["name"] for field in response["fields"]]

        return self.object_fields[object_type]

    async def _describe_object(self, object_type: SalesforceObject) -> dict:
        return await self.get(f"sobjects/{object_type.value}/describe/")

    async def _get_related_objects(
        self,
        child_object_type: SalesforceObject,
        parent_object_type: SalesforceObject,
        parent_object_id: str,
        limit: int | None = 10,
    ) -> list[dict]:
        """Get the objects that are associated with another Salesforce object.

        Args:
            child_object_type: The type of child object to retrieve.
            parent_object_type: The type of parent object.
            parent_object_id: The ID of the parent object.
            limit: The maximum number of related objects to retrieve.

        Returns:
            The list of related objects.
        """
        try:
            response = await self.get(
                f"sobjects/{parent_object_type.value}/{parent_object_id}/{child_object_type.plural.lower()}",
                params={"limit": limit},
            )
            return cast(list, response["records"])
        except ResourceNotFoundError:
            return []

    async def get_object_by_id(self, object_id: str) -> dict | None:
        """Get a Salesforce object by its ID.

        Args:
            object_id: The ID of the Salesforce object to retrieve.

        Returns:
            The Salesforce object.
        """
        try:
            response = await self.get(f"sobjects/{object_id}")
            return clean_object_data(response)
        except ResourceNotFoundError:
            if "User" not in object_id:
                return await self.get_object_by_id(f"User/{object_id}")
            return None

    async def get_account(self, account_id: str) -> dict[str, Any] | None:
        """Get an account by its ID.

        Args:
            account_id: The ID of the account to retrieve.

        Returns:
            The account.
        """
        try:
            return cast(dict, await self.get(f"sobjects/Account/{account_id}"))
        except ResourceNotFoundError:
            return None

    async def get_account_contacts(self, account_id: str, limit: int | None = 10) -> list[dict]:
        """Get the contacts associated with an account.

        Args:
            account_id: The ID of the account to retrieve the contacts for.
            limit: The maximum number of contacts to retrieve.

        Returns:
            The list of contacts with cleaned and standardized dictionaries.
        """
        contacts = await self._get_related_objects(
            SalesforceObject.CONTACT, SalesforceObject.ACCOUNT, account_id, limit
        )

        return [
            clean_contact_data(contact)
            for contact in await asyncio.gather(*[
                self.enrich_contact_with_notes(contact, limit) for contact in contacts
            ])
        ]

    async def enrich_contact_with_notes(self, contact: dict, limit: int | None = 10) -> dict:
        """Get the notes associated with a contact and add to the contact dictionary.

        Args:
            contact: The contact to retrieve the notes for.
            limit: The maximum number of notes to retrieve.

        Returns:
            The contact with the notes added to the dictionary in the key `Notes`.
        """
        notes = await self.get_notes(contact["Id"], limit)
        if notes:
            contact["Notes"] = notes
        return contact

    async def get_account_leads(self, account_id: str, limit: int | None = 10) -> list[dict]:
        """Get the leads associated with an account.

        Args:
            account_id: The ID of the account to retrieve the leads for.
            limit: The maximum number of leads to retrieve.

        Returns:
            The list of leads with cleaned and standardized dictionaries.
        """
        leads = await self._get_related_objects(
            SalesforceObject.LEAD, SalesforceObject.ACCOUNT, account_id, limit
        )
        return [clean_lead_data(lead) for lead in leads]

    async def get_notes(self, parent_id: str, limit: int | None = 10) -> list[dict]:
        """Get the notes associated with a Salesforce object.

        Args:
            parent_id: The ID of the Salesforce object to retrieve the notes for.
            limit: The maximum number of notes to retrieve.

        Returns:
            The list of notes with cleaned and standardized dictionaries.
        """
        query = build_soql_query(
            "SELECT Id, Title, Body, OwnerId, CreatedById, CreatedDate "
            "FROM Note "
            "WHERE ParentId = '{parent_id}' "
            "LIMIT {limit}",
            parent_id=parent_id,
            limit=limit,
        )
        response = await self.get("query", params={"q": query})
        notes = response["records"]
        return [clean_note_data(note) for note in notes]

    # TODO: Add support for retrieving Currency, when enabled in the org account.
    # If not enabled and we try to retrieve it, we get a 400 error.
    # More inf: https://developer.salesforce.com/docs/atlas.en-us.254.0.object_reference.meta/object_reference/sforce_api_objects_opportunity.htm#i1455437
    async def get_account_opportunities(
        self,
        account_id: str,
        limit: int | None = 10,
    ) -> list[dict]:
        """Get the opportunities associated with an account.

        Args:
            account_id: The ID of the account to retrieve the opportunities for.
            limit: The maximum number of opportunities to retrieve.

        Returns:
            The list of opportunities with cleaned and standardized dictionaries.
        """
        query = build_soql_query(
            "SELECT Id, Name, Type, StageName, OwnerId, CreatedById, LastModifiedById, "
            "Description, Amount, Probability, ExpectedRevenue, CloseDate, ContactId "
            "FROM Opportunity "
            "WHERE AccountId = '{account_id}' "
            "LIMIT {limit}",
            account_id=account_id,
            limit=limit,
        )
        response = await self.get("query", params={"q": query})
        opportunities = response["records"]
        return [clean_opportunity_data(opportunity) for opportunity in opportunities]

    async def get_account_tasks(
        self,
        account_id: str,
        limit: int | None = 10,
    ) -> list[dict]:
        """Get the tasks associated with an account.

        Args:
            account_id: The ID of the account to retrieve the tasks for.
            limit: The maximum number of tasks to retrieve.

        Returns:
            The list of tasks with cleaned and standardized dictionaries.
        """
        tasks = await self._get_related_objects(
            SalesforceObject.TASK, SalesforceObject.ACCOUNT, account_id, limit
        )
        return [clean_task_data(task) for task in tasks]

    async def enrich_account(
        self,
        account_id: str | None = None,
        account_data: dict[str, Any] | None = None,
        limit_per_association: int | None = 10,
    ) -> dict:
        """Enrich account dictionary with contacts, leads, opportunities, and tasks.

        Provide `account_id` or `account_data`, but not both.

        Args:
            account_id: The ID of the account to retrieve the associations for.
            account_data: The account data to enrich.
            limit_per_association: The maximum number of associations to retrieve.

        Returns:
            The account with the associations added to the dictionary.
        """
        if (account_id and account_data) or (not account_id and not account_data):
            raise ValueError("Must provide exactly one of `account_id` or `account_data`")

        if account_data is None:
            account_data = await self.get_account(cast(str, account_id))

            if not account_data:
                raise ResourceNotFoundError([f"Account not found with ID: {account_id}"])

        if not account_id:
            account_id = cast(str, account_data["Id"])

        associations = await asyncio.gather(
            self.get_account_contacts(account_id, limit=limit_per_association),
            self.get_account_leads(account_id, limit=limit_per_association),
            self.get_account_opportunities(account_id, limit=limit_per_association),
            self.get_account_tasks(account_id, limit=limit_per_association),
        )

        for association in associations:
            for item in association:
                try:
                    obj_type = SalesforceObject(get_object_type(item)).plural
                except ValueError:
                    obj_type = get_object_type(item) + "s"

                if obj_type not in account_data:
                    account_data[obj_type] = []
                account_data[obj_type].append(item)

        return await self.expand_account_associations(account_data)

    async def expand_account_associations(self, account: dict) -> dict:
        """Expand with metadata about objects referenced by ID in an account dictionary.

        This method will, for example, expand an `OwnerId` or `ContactId` referenced in an account
        dictionary with the owner or contact name, for example.

        Args:
            account: The account dictionary to expand.

        Returns:
            The account with the associations expanded.
        """
        objects_by_id = {
            obj["Id"]: obj
            for obj_type in SalesforceObject
            for obj in account.get(obj_type.plural, [])
        }
        objects_by_id[account["Id"]] = account

        referenced_ids = get_ids_referenced(
            account,
            *[account.get(obj_type.plural, []) for obj_type in SalesforceObject],
        )

        missing_referenced_ids = [ref for ref in referenced_ids if ref not in objects_by_id]

        if missing_referenced_ids:
            missing_objects = await asyncio.gather(*[
                self.get_object_by_id(missing_id) for missing_id in missing_referenced_ids
            ])
            objects_by_id.update({obj["Id"]: obj for obj in missing_objects if obj is not None})

        account = expand_associations(account, objects_by_id)

        for object_type in SalesforceObject:
            if object_type.plural not in account:
                continue

            expanded_items = []

            for item in account[object_type.plural]:
                if "AccountId" in item:
                    del item["AccountId"]

                expanded_items.append(expand_associations(item, objects_by_id))

            if object_type == SalesforceObject.CONTACT:
                for contact in expanded_items:
                    if "Notes" in contact:
                        contact["Notes"] = [
                            expand_associations(note, objects_by_id) for note in contact["Notes"]
                        ]

            account[object_type.plural] = expanded_items

        return account

    async def create_contact(
        self,
        account_id: str,
        last_name: str,
        first_name: str | None = None,
        email: str | None = None,
        phone: str | None = None,
        mobile_phone: str | None = None,
        title: str | None = None,
        department: str | None = None,
        description: str | None = None,
    ) -> dict:
        """Create a contact in Salesforce.

        Args:
            account_id: The ID of the account to associate the contact with.
            last_name: The last name of the contact.
            first_name: The first name of the contact.
            email: The email of the contact.
            phone: The phone number of the contact.
            mobile_phone: The mobile phone number of the contact.
            title: The title of the contact.
            department: The department of the contact.
            description: The description of the contact.

        Returns:
            The created contact.
        """
        data = {
            "AccountId": account_id,
            "FirstName": first_name,
            "LastName": last_name,
            "Email": email,
            "Phone": phone,
            "MobilePhone": mobile_phone,
            "Title": title,
            "Department": department,
            "Description": description,
        }

        return await self.post(
            f"sobjects/{SalesforceObject.CONTACT.value}",
            json_data=remove_none_values(data),
        )