source-salesforce-native: speed up building resources with scatter-gather & improve access token handling

Previously, all resources in enabled_bindings were built in series. Since the fields for each resource must be fetched from Salesforce before they're built, building resources was a very slow process. Using the scatter-gather technique to build multiple resources concurrently significantly speeds up the process.

I also figured out how to store instance_url in the config's credentials. We no longer need to make a duplicate request outside of the standard OAuth method to fetch the instance url.

Also around OAuth, I decided to use SalesforceTokenSource to subclass TokenSource and override the default access token expiration duration of 0 seconds. This lets us reuse our cached access token instead of always requesting a new one when we make an HTTP request. This also gets around intermittent failures when exchanging the same refresh token for an access token mutiple times within a small time window.

Author: Alex-Bair

source-salesforce-native/source_salesforce_native/models.py

@@ -15,7 +15,9 @@
 from estuary_cdk.capture.common import (
     ConnectorState as GenericConnectorState,
 )
-
+from estuary_cdk.http import (
+    TokenSource,
+)
 
 EARLIEST_VALID_DATE_IN_SALESFORCE = datetime(1700, 1, 1, tzinfo=UTC)
 
@@ -46,6 +48,7 @@
     ),
     accessTokenResponseMap={
         "refresh_token": "/refresh_token",
+        "instance_url": "/instance_url",
     },
 )
 
@@ -54,10 +57,45 @@ def update_oauth_spec(is_sandbox: bool):
     OAUTH2_SPEC.accessTokenUrlTemplate = f"https://{'test' if is_sandbox else 'login'}.salesforce.com/services/oauth2/token"
 
 
+# The access token response does not contain any field indicating if/when the access token we receive expires.
+# The default TokenSource.AccessTokenResponse.expires_in is 0, causing every request to fetch a new access token.
+# The actual expires_in value depends on the session settings in the user's Salesforce account. The default seems
+# to be 2 hours. More importantly, each time the access token is used, it's validity is extended. Meaning that as long
+# as the access token is actively used, it shouldn't expire. Setting an expires_in of 1 hour should be a 
+# safe fallback in case a user changes all enabled bindings' intervals to an hour or greater.
+class SalesforceTokenSource(TokenSource):
+    class AccessTokenResponse(TokenSource.AccessTokenResponse):
+        expires_in: int = 1 * 60 * 60
+
+
+class SalesforceOAuth2Credentials(BaseOAuth2Credentials):
+    instance_url: str = Field(
+        title="Instance URL",
+    )
+
+    @staticmethod
+    def for_provider(provider: str) -> type["SalesforceOAuth2Credentials"]:
+        """
+        Builds an OAuth2Credentials model for the given OAuth2 `provider`.
+        This routine is only available in Pydantic V2 environments.
+        """
+        from pydantic import ConfigDict
+
+        class _OAuth2Credentials(SalesforceOAuth2Credentials):
+            model_config = ConfigDict(
+                json_schema_extra={"x-oauth2-provider": provider},
+                title="OAuth",
+            )
+
+            def _you_must_build_oauth2_credentials_for_a_provider(self): ...
+
+        return _OAuth2Credentials
+
+
 if TYPE_CHECKING:
-    OAuth2Credentials = BaseOAuth2Credentials
+    OAuth2Credentials = SalesforceOAuth2Credentials
 else:
-    OAuth2Credentials = BaseOAuth2Credentials.for_provider(OAUTH2_SPEC.provider)
+    OAuth2Credentials = SalesforceOAuth2Credentials.for_provider(OAUTH2_SPEC.provider)
 
 
 def default_start_date():
@@ -97,10 +135,6 @@ class Advanced(BaseModel):
 ConnectorState = GenericConnectorState[ResourceState]
 
 
-class AccessTokenResponse(BaseModel, extra="allow"):
-    instance_url: str
-
-
 class PartialSObject(BaseModel, extra="allow"):
     name: str
     queryable: bool

source-salesforce-native/source_salesforce_native/resources.py

@@ -1,11 +1,12 @@
+import asyncio
 from datetime import datetime, timedelta, UTC
 import functools
 from logging import Logger
 from typing import Any
 
 from estuary_cdk.flow import CaptureBinding
 from estuary_cdk.capture import common, Task
-from estuary_cdk.http import HTTPMixin, TokenSource
+from estuary_cdk.http import HTTPMixin
 
 from .supported_standard_objects import SUPPORTED_STANDARD_OBJECTS, COMMON_CUSTOM_OBJECT_DETAILS
 
@@ -17,7 +18,7 @@
     EndpointConfig,
     ResourceConfig,
     ResourceState,
-    AccessTokenResponse,
+    SalesforceTokenSource,
     GlobalDescribeObjectsResponse,
     SOAP_TYPES_NOT_SUPPORTED_BY_BULK_API,
     FieldDetailsDict,
@@ -35,22 +36,7 @@
 
 
 CUSTOM_OBJECT_SUFFIX = '__c'
-
-
-async def _fetch_instance_url(log: Logger, http: HTTPMixin, config: EndpointConfig) -> str:
-    url = OAUTH2_SPEC.accessTokenUrlTemplate
-    body = {
-        "grant_type": "refresh_token",
-        "client_id": config.credentials.client_id,
-        "client_secret": config.credentials.client_secret,
-        "refresh_token": config.credentials.refresh_token,
-    }
-
-    response = AccessTokenResponse.model_validate_json(
-        await http.request(log, url, method="POST", form=body, _with_token=False)
-    )
-
-    return response.instance_url
+BUILD_RESOURCE_SEMAPHORE_LIMIT = 15
 
 
 async def _fetch_queryable_objects(log: Logger, http: HTTPMixin, instance_url: str) -> list[str]:
@@ -272,45 +258,56 @@ async def enabled_resources(
     log: Logger, http: HTTPMixin, config: EndpointConfig, bindings: list[common._ResolvableBinding]
 ) -> list[common.Resource]:
     update_oauth_spec(config.is_sandbox)
-    http.token_source = TokenSource(oauth_spec=OAUTH2_SPEC, credentials=config.credentials)
-
-    instance_url = await _fetch_instance_url(log, http, config)
+    http.token_source = SalesforceTokenSource(oauth_spec=OAUTH2_SPEC, credentials=config.credentials)
 
     enabled_binding_names: list[str] = []
 
     for binding in bindings:
         path: list[str] = binding.resourceConfig.path()
         enabled_binding_names.append(path[0])
 
-    bulk_job_manager = BulkJobManager(http, log, instance_url)
-    rest_query_manager = RestQueryManager(http, log, instance_url)
-    resources: list[common.Resource] = []
-
-    for name in enabled_binding_names:
-        r = await _object_to_resource(log, http, config, bulk_job_manager, rest_query_manager, instance_url, name, True)
-        if r:
-            resources.append(r)
+    bulk_job_manager = BulkJobManager(http, log, config.credentials.instance_url)
+    rest_query_manager = RestQueryManager(http, log, config.credentials.instance_url)
+
+    # If we concurrently send multiple requests that exchange the same refresh token for an access token,
+    # some of those requests intermittently fail. 
+    # https://help.salesforce.com/s/articleView?id=xcloud.remoteaccess_oauth_refresh_token_flow.htm&type=5#:~:text=Avoid%20sending%20simultaneous%20requests%20that%20contain%20the%20same%20refresh%20token.%20If%20your%20client%20sends%20identical%20requests%20at%20the%20same%20time%2C%20some%20of%20the%20requests%20fail%20intermittently%20and%20the%20Status%20column%20in%20the%20Login%20History%20displays%20Failed%3A%20Token%20request%20is%20already%20being%20processed.
+    # To avoid this, we make a noop request to set the token_source's access token before using the scatter-gather
+    # technique to make multiple requests concurrently. This prevents the first BUILD_RESOURCE_SEMAPHORE_LIMIT
+    # requests from all exchanging the same access token and encountering that intermittent error.
+    await _fetch_queryable_objects(log, http, config.credentials.instance_url)
+
+    semaphore = asyncio.Semaphore(BUILD_RESOURCE_SEMAPHORE_LIMIT)
+    async def build_resource(name: str) -> common.Resource | None:
+        async with semaphore:
+            return await _object_to_resource(
+                log, http, config, bulk_job_manager, rest_query_manager, config.credentials.instance_url, name, True
+            )
+
+    task_results = await asyncio.gather(
+        *(
+            build_resource(name) for name in enabled_binding_names
+        )
+    )
 
-    return resources
+    return [resource for resource in task_results if resource is not None]
 
 
 # all_resources returns resources for all possible supported bindings.
 async def all_resources(
     log: Logger, http: HTTPMixin, config: EndpointConfig
 ) -> list[common.Resource]:
     update_oauth_spec(config.is_sandbox)
-    http.token_source = TokenSource(oauth_spec=OAUTH2_SPEC, credentials=config.credentials)
-
-    instance_url = await _fetch_instance_url(log, http, config)
+    http.token_source = SalesforceTokenSource(oauth_spec=OAUTH2_SPEC, credentials=config.credentials)
 
-    queryable_object_names = await _fetch_queryable_objects(log, http, instance_url)
+    queryable_object_names = await _fetch_queryable_objects(log, http, config.credentials.instance_url)
 
-    bulk_job_manager = BulkJobManager(http, log, instance_url)
-    rest_query_manager = RestQueryManager(http, log, instance_url)
+    bulk_job_manager = BulkJobManager(http, log, config.credentials.instance_url)
+    rest_query_manager = RestQueryManager(http, log, config.credentials.instance_url)
     resources: list[common.Resource] = []
 
     for name in queryable_object_names:
-        r = await _object_to_resource(log, http, config, bulk_job_manager, rest_query_manager, instance_url, name)
+        r = await _object_to_resource(log, http, config, bulk_job_manager, rest_query_manager, config.credentials.instance_url, name)
         if r:
             resources.append(r)