generalize batch execution, add use_batch flag

fix SAT

Author: eugene-kulak

airbyte-integrations/connectors/source-facebook-marketing/acceptance-test-config.yml

@@ -13,6 +13,7 @@ tests:
     - config_path: "secrets/config.json"
   basic_read:
     - config_path: "secrets/config.json"
+      empty_streams: ["videos"]
   incremental:
     - config_path: "secrets/config.json"
       configured_catalog_path: "integration_tests/configured_catalog_without_insights.json"

airbyte-integrations/connectors/source-facebook-marketing/source_facebook_marketing/api.py

@@ -4,7 +4,6 @@
 
 import json
 import logging
-from datetime import datetime
 from time import sleep
 from typing import Tuple
 

airbyte-integrations/connectors/source-facebook-marketing/source_facebook_marketing/streams/common.py

@@ -24,13 +24,6 @@ class JobException(Exception):
     """Scheduled job failed"""
 
 
-def batch(iterable: Sequence, size: int = 1) -> Iterable:
-    """Split sequence in chunks"""
-    total_size = len(iterable)
-    for ndx in range(0, total_size, size):
-        yield iterable[ndx : min(ndx + size, total_size)]
-
-
 def retry_pattern(backoff_type, exception, **wait_gen_kwargs):
     def log_retry_attempt(details):
         _, exc, _ = sys.exc_info()

airbyte-integrations/connectors/source-facebook-marketing/source_facebook_marketing/streams/streams.py

@@ -4,41 +4,26 @@
 
 import base64
 import logging
-import urllib.parse as urlparse
 from abc import ABC
 from datetime import datetime
-from typing import Any, Iterable, Iterator, List, Mapping, MutableMapping, Sequence, TYPE_CHECKING, Optional
+from typing import Any, Iterable, Iterator, List, Mapping, MutableMapping, TYPE_CHECKING, Optional
 
 import pendulum
 import requests
 from airbyte_cdk.models import SyncMode
 from airbyte_cdk.sources.streams import Stream
 from airbyte_cdk.sources.utils.transform import TransformConfig, TypeTransformer
 from cached_property import cached_property
+from facebook_business.adobjects.abstractobject import AbstractObject
 from facebook_business.api import FacebookAdsApiBatch, FacebookRequest, FacebookResponse
-from .common import batch, deep_merge
+from .common import deep_merge
 
 if TYPE_CHECKING:
     from source_facebook_marketing.api import API
 
 logger = logging.getLogger("airbyte")
 
 
-def remove_params_from_url(url: str, params: List[str]) -> str:
-    """
-    Parses a URL and removes the query parameters specified in params
-    :param url: URL
-    :param params: list of query parameters
-    :return: URL with params removed
-    """
-    parsed = urlparse.urlparse(url)
-    query = urlparse.parse_qs(parsed.query, keep_blank_values=True)
-    filtered = dict((k, v) for k, v in query.items() if k not in params)
-    return urlparse.urlunparse(
-        [parsed.scheme, parsed.netloc, parsed.path, parsed.params, urlparse.urlencode(filtered, doseq=True), parsed.fragment]
-    )
-
-
 def fetch_thumbnail_data_url(url: str) -> Optional[str]:
     """Request thumbnail image and return it embedded into the data-link"""
     try:
@@ -47,8 +32,10 @@ def fetch_thumbnail_data_url(url: str) -> Optional[str]:
             _type = response.headers["content-type"]
             data = base64.b64encode(response.content)
             return f"data:{_type};base64,{data.decode('ascii')}"
-    except requests.exceptions.RequestException:
-        pass
+        else:
+            logger.warning(f"Got {repr(response)} while requesting thumbnail image.")
+    except requests.exceptions.RequestException as exc:
+        logger.warning(f"Got {str(exc)} while requesting thumbnail image.")
     return None
 
 
@@ -59,10 +46,13 @@ class FBMarketingStream(Stream, ABC):
     transformer: TypeTransformer = TypeTransformer(TransformConfig.DefaultSchemaNormalization)
 
     page_size = 100
+    use_batch = False
 
     enable_deleted = False
     entity_prefix = None
 
+    MAX_BATCH_SIZE = 50
+
     def __init__(self, api: 'API', include_deleted: bool = False, **kwargs):
         super().__init__(**kwargs)
         self._api = api
@@ -73,26 +63,36 @@ def fields(self) -> List[str]:
         """List of fields that we want to query, for now just all properties from stream's schema"""
         return list(self.get_json_schema().get("properties", {}).keys())
 
-    def execute_in_batch(self, requests: Iterable[FacebookRequest]) -> Sequence[MutableMapping[str, Any]]:
+    def _execute_batch(self, batch):
+        """Execute batch, retry in case of failures"""
+        while batch:
+            batch = batch.execute()
+            if batch:
+                logger.info("Retry failed requests in batch")
+
+    def execute_in_batch(self, pending_requests: Iterable[FacebookRequest]) -> Iterable[MutableMapping[str, Any]]:
         """Execute list of requests in batches"""
         records = []
 
         def success(response: FacebookResponse):
             records.append(response.json())
 
         def failure(response: FacebookResponse):
-            logger.info(f"Request failed with response: {response.body()}")
+            # FIXME: stop sync or retry
+            logger.warning(f"Request failed with response: {response.body()}")
 
         api_batch: FacebookAdsApiBatch = self._api.api.new_batch()
-        for request in requests:
+        for request in pending_requests:
             api_batch.add_request(request, success=success, failure=failure)
+            if len(api_batch) == self.MAX_BATCH_SIZE:
+                self._execute_batch(api_batch)
 
-        while api_batch:
-            api_batch = api_batch.execute()
-            if api_batch:
-                logger.info("Retry failed requests in batch")
+            yield from records
+            api_batch: FacebookAdsApiBatch = self._api.api.new_batch()
+            records = []
 
-        return records
+        self._execute_batch(api_batch)
+        yield from records
 
     def read_records(
         self,
@@ -102,19 +102,23 @@ def read_records(
         stream_state: Mapping[str, Any] = None,
     ) -> Iterable[Mapping[str, Any]]:
         """Main read method used by CDK"""
-        for record in self._read_records(params=self.request_params(stream_state=stream_state)):
-            yield self._extend_record(record, fields=self.fields)
+        records_iter = self._read_records(params=self.request_params(stream_state=stream_state))
+        loaded_records_iter = (record.api_get(fields=self.fields, pending=self.use_batch) for record in records_iter)
+        if self.use_batch:
+            loaded_records_iter = self.execute_in_batch(loaded_records_iter)
+
+        for record in loaded_records_iter:
+            if isinstance(record, AbstractObject):
+                yield record.export_all_data()
+            else:
+                yield record
 
     def _read_records(self, params: Mapping[str, Any]) -> Iterable:
         """Wrapper around query to backoff errors.
         We have default implementation because we still can override read_records so this method is not mandatory.
         """
         return []
 
-    def _extend_record(self, obj: Any, **kwargs):
-        """Wrapper around api_get to backoff errors"""
-        return obj.api_get(**kwargs).export_all_data()
-
     def request_params(self, **kwargs) -> MutableMapping[str, Any]:
         """Parameters that should be passed to query_records method"""
         params = {"limit": self.page_size}
@@ -205,39 +209,30 @@ class AdCreatives(FBMarketingStream):
     """
 
     entity_prefix = "adcreative"
-    batch_size = 50
+    use_batch = True
 
     def __init__(self, fetch_thumbnail_images: bool = False, **kwargs):
         super().__init__(**kwargs)
         self._fetch_thumbnail_images = fetch_thumbnail_images
 
+    @cached_property
+    def fields(self) -> List[str]:
+        """ Remove "thumbnail_data_url" field because it is computed field and it's not a field that we can request from Facebook
+        """
+        return [f for f in super().fields if f != "thumbnail_data_url"]
+
     def read_records(
         self,
         sync_mode: SyncMode,
         cursor_field: List[str] = None,
         stream_slice: Mapping[str, Any] = None,
         stream_state: Mapping[str, Any] = None,
     ) -> Iterable[Mapping[str, Any]]:
-        """Read records using batch API"""
-        records = self._read_records(params=self.request_params(stream_state=stream_state))
-        # "thumbnail_data_url" is a field in our stream's schema because we
-        # output it (see fix_thumbnail_urls below), but it's not a field that
-        # we can request from Facebook
-        request_fields = [f for f in self.fields if f != "thumbnail_data_url"]
-        requests = [record.api_get(fields=request_fields, pending=True) for record in records]
-        for requests_batch in batch(requests, size=self.batch_size):
-            for record in self.execute_in_batch(requests_batch):
-                yield self.fix_thumbnail_urls(record)
-
-    def fix_thumbnail_urls(self, record: MutableMapping[str, Any]) -> MutableMapping[str, Any]:
-        """Cleans and, if enabled, fetches thumbnail URLs for each creative."""
-        # The thumbnail_url contains some extra query parameters that don't affect the validity of the URL, but break SAT
-        thumbnail_url = record.get("thumbnail_url")
-        if thumbnail_url:
-            record["thumbnail_url"] = remove_params_from_url(thumbnail_url, ["_nc_hash", "d"])
+        """Read with super method and append thumbnail_data_url if enabled"""
+        for record in super().read_records(sync_mode, cursor_field, stream_slice, stream_state):
             if self._fetch_thumbnail_images:
-                record["thumbnail_data_url"] = fetch_thumbnail_data_url(thumbnail_url)
-        return record
+                record["thumbnail_data_url"] = fetch_thumbnail_data_url(record.get("thumbnail_url"))
+            yield record
 
     def _read_records(self, params: Mapping[str, Any]) -> Iterator:
         return self._api.account.get_ad_creatives(params=params)