update search endpoint to include associations

Author: kevin-induro

sources/hubspot/__init__.py

@@ -50,9 +50,9 @@
 from .settings import (
     ALL_OBJECTS,
     ARCHIVED_PARAM,
+    CRM_OBJECT_ASSOCIATIONS,
     CRM_OBJECT_ENDPOINTS,
     CRM_PIPELINES_ENDPOINT,
-    CRM_SEARCH_OBJECT_ENDPOINTS,
     ENTITY_PROPERTIES,
     LAST_MODIFIED_PROPERTY,
     HUBSPOT_CREATION_DATE,
@@ -95,7 +95,13 @@ def fetch_data_for_properties(
     logger.info(f"Fetching data for {object_type}.")
     # The Hubspot API expects a comma separated string as properties
     joined_props = ",".join(sorted(props))
-    params: Dict[str, Any] = {"properties": joined_props, "limit": 100}
+    associations = CRM_OBJECT_ASSOCIATIONS[object_type]
+    joined_associations = ",".join(associations)
+    params: Dict[str, Any] = {
+        "associations": joined_associations,
+        "properties": joined_props,
+        "limit": 100,
+    }
     context: Optional[Dict[str, Any]] = (
         {SOFT_DELETE_KEY: False} if soft_delete else None
     )
@@ -120,8 +126,9 @@ def fetch_data_for_properties(
 
         try:
             yield from search_data(
-                CRM_SEARCH_OBJECT_ENDPOINTS[object_type],
+                CRM_OBJECT_ENDPOINTS[object_type],
                 api_key,
+                associations=associations,
                 params=search_params,
                 context=context,
             )
@@ -228,7 +235,7 @@ def crm_object_history(
     # This is especially relevant for columns of type "number" in Hubspot
     # that are returned as strings by the API
     for batch in fetch_property_history(
-        CRM_SEARCH_OBJECT_ENDPOINTS[object_type],
+        CRM_OBJECT_ENDPOINTS[object_type],
         api_key,
         ",".join(sorted(props_to_type.keys())),
     ):

sources/hubspot/helpers.py

@@ -8,7 +8,12 @@
 from dlt.common.schema.typing import TColumnSchema
 from dlt.sources.helpers import requests
 
-from .settings import OBJECT_TYPE_PLURAL, HS_TO_DLT_TYPE
+from .settings import (
+    CRM_ASSOCIATIONS_ENDPOINT,
+    CRM_SEARCH_ENDPOINT,
+    OBJECT_TYPE_PLURAL,
+    HS_TO_DLT_TYPE,
+)
 
 BASE_URL = "https://api.hubapi.com/"
 
@@ -62,7 +67,7 @@ def search_pagination(
     if _next:
         after = _next["after"]
         # Get the next page response
-        r = requests.post(url, headers=headers, json={**params, after: after})
+        r = requests.post(url, headers=headers, json={**params, "after": after})
         return r.json()  # type: ignore
     else:
         return None
@@ -149,11 +154,12 @@ def fetch_property_history(
 def search_data(
     endpoint: str,
     api_key: str,
+    associations: Optional[List[str]] = None,
     params: Optional[Dict[str, Any]] = None,
     context: Optional[Dict[str, Any]] = None,
 ) -> Iterator[List[Dict[str, Any]]]:
     # Construct the URL and headers for the API request
-    url = get_url(endpoint)
+    url = get_url(CRM_SEARCH_ENDPOINT.format(crm_endpoint=endpoint))
     headers = _get_headers(api_key)
 
     # Make the API request
@@ -168,7 +174,9 @@ def search_data(
         # Yield the properties of each result in the API response
         while _data is not None:
             if "results" in _data:
-                yield _data_to_objects(_data, headers, context)
+                yield _data_to_objects(
+                    _data, endpoint, headers, associations=associations, context=context
+                )
 
             # Follow pagination links if they exist
             _data = search_pagination(url, _data, headers, params)
@@ -220,15 +228,17 @@ def fetch_data(
     # Yield the properties of each result in the API response
     while _data is not None:
         if "results" in _data:
-            yield _data_to_objects(_data, headers, context)
+            yield _data_to_objects(_data, endpoint, headers, context=context)
 
         # Follow pagination links if they exist
         _data = pagination(_data, headers)
 
 
 def _data_to_objects(
     data: Any,
+    endpoint: str,
     headers: Dict[str, str],
+    associations: Optional[List[str]] = None,
     context: Optional[Dict[str, Any]] = None,
 ) -> List[Dict[str, Any]]:
     _objects: List[Dict[str, Any]] = []
@@ -240,19 +250,36 @@ def _data_to_objects(
         if "associations" in _result:
             for association in _result["associations"]:
                 __data = _result["associations"][association]
-
-                __values = extract_association_data(_obj, __data, association, headers)
-
-                # remove duplicates from list of dicts
-                __values = [dict(t) for t in {tuple(d.items()) for d in __values}]
-
-                _obj[association] = __values
+                _add_association_data(__data, association, headers, _obj)
+        elif associations is not None:
+            for association in associations:
+                __endpoint = get_url(
+                    CRM_ASSOCIATIONS_ENDPOINT.format(
+                        crm_endpoint=endpoint,
+                        object_id=_result["id"],
+                        association=association,
+                    )
+                )
+                r = requests.get(__endpoint, headers=headers, params={"limit": 500})
+                __data = r.json()
+                _add_association_data(__data, association, headers, _obj)
         if context:
             _obj.update(context)
         _objects.append(_obj)
     return _objects
 
 
+def _add_association_data(
+    data: Any, association: str, headers: Dict[str, str], obj: Any
+) -> None:
+    __values = extract_association_data(obj, data, association, headers)
+
+    # remove duplicates from list of dicts
+    __values = [dict(t) for t in {tuple(d.items()) for d in __values}]
+
+    obj[association] = __values
+
+
 def _get_property_names_types(
     api_key: str, object_type: str
 ) -> Dict[str, Union[str, None]]:

sources/hubspot/settings.py

@@ -6,28 +6,8 @@
 HUBSPOT_CREATION_DATE = pendulum.datetime(year=2006, month=6, day=1)
 STARTDATE = pendulum.datetime(year=2024, month=2, day=10)
 
-CRM_CONTACTS_SEARCH_ENDPOINT = "/crm/v3/objects/contacts/search"
-CRM_COMPANIES_SEARCH_ENDPOINT = "/crm/v3/objects/companies/search"
-CRM_DEALS_SEARCH_ENDPOINT = "/crm/v3/objects/deals/search"
-CRM_PRODUCTS_SEARCH_ENDPOINT = "/crm/v3/objects/products/search"
-CRM_TICKETS_SEARCH_ENDPOINT = "/crm/v3/objects/tickets/search"
-CRM_QUOTES_SEARCH_ENDPOINT = "/crm/v3/objects/quotes/search"
-
-CRM_SEARCH_OBJECT_ENDPOINTS = {
-    "contact": CRM_CONTACTS_SEARCH_ENDPOINT,
-    "company": CRM_COMPANIES_SEARCH_ENDPOINT,
-    "deal": CRM_DEALS_SEARCH_ENDPOINT,
-    "product": CRM_PRODUCTS_SEARCH_ENDPOINT,
-    "ticket": CRM_TICKETS_SEARCH_ENDPOINT,
-    "quote": CRM_QUOTES_SEARCH_ENDPOINT,
-}
-
-CRM_CONTACTS_ENDPOINT = (
-    "/crm/v3/objects/contacts?associations=deals,products,tickets,quotes"
-)
-CRM_COMPANIES_ENDPOINT = (
-    "/crm/v3/objects/companies?associations=contacts,deals,products,tickets,quotes"
-)
+CRM_CONTACTS_ENDPOINT = "/crm/v3/objects/contacts"
+CRM_COMPANIES_ENDPOINT = "/crm/v3/objects/companies"
 CRM_DEALS_ENDPOINT = "/crm/v3/objects/deals"
 CRM_PRODUCTS_ENDPOINT = "/crm/v3/objects/products"
 CRM_TICKETS_ENDPOINT = "/crm/v3/objects/tickets"
@@ -36,6 +16,9 @@
 CRM_PROPERTIES_ENDPOINT = "/crm/v3/properties/{objectType}/{property_name}"
 CRM_PIPELINES_ENDPOINT = "/crm/v3/pipelines/{objectType}"
 
+CRM_SEARCH_ENDPOINT = "{crm_endpoint}/search"
+CRM_ASSOCIATIONS_ENDPOINT = "{crm_endpoint}/{object_id}/associations/{association}"
+
 CRM_OBJECT_ENDPOINTS = {
     "contact": CRM_CONTACTS_ENDPOINT,
     "company": CRM_COMPANIES_ENDPOINT,
@@ -46,6 +29,16 @@
     "owner": CRM_OWNERS_ENDPOINT,
 }
 
+CRM_OBJECT_ASSOCIATIONS = {
+    "contact": ["deals", "products", "tickets", "quotes"],
+    "company": ["contacts", "deals", "products", "tickets", "quotes"],
+    "deal": [],
+    "product": [],
+    "ticket": [],
+    "quote": [],
+    "owner": [],
+}
+
 WEB_ANALYTICS_EVENTS_ENDPOINT = "/events/v3/events?objectType={objectType}&objectId={objectId}&occurredAfter={occurredAfter}&occurredBefore={occurredBefore}&sort=-occurredAt"
 
 OBJECT_TYPE_SINGULAR = {