#12486 and #49 from alpha-beta-issues fixes

davydov-d · davydov-d · commit c37fe3bb82e0 · 2022-05-17T14:32:20.000+03:00
diff --git a/airbyte-integrations/connectors/source-google-ads/source_google_ads/custom_query_stream.py b/airbyte-integrations/connectors/source-google-ads/source_google_ads/custom_query_stream.py
@@ -72,7 +72,7 @@ def get_json_schema(self) -> Dict[str, Any]:
                 # Represents protobuf message and could be anything, set custom
                 # attribute "protobuf_message" to convert it to a string (or
                 # array of strings) later.
-                # https://developers.google.com/google-ads/api/reference/rpc/v8/GoogleAdsFieldDataTypeEnum.GoogleAdsFieldDataType?hl=en#message
+                # https://developers.google.com/google-ads/api/reference/rpc/v9/GoogleAdsFieldDataTypeEnum.GoogleAdsFieldDataType?hl=en#message
                 if node.is_repeated:
                     output_type = ["array", "null"]
                 else:
@@ -95,7 +95,8 @@ def get_json_schema(self) -> Dict[str, Any]:
     WHERE_EXPR = re.compile("where.*", flags=RE_FLAGS)
     # list of keywords that can come after WHERE clause,
     # according to https://developers.google.com/google-ads/api/docs/query/grammar
-    KEYWORDS_EXPR = re.compile("(order by|limit|parameters)", flags=RE_FLAGS)
+    # each whitespace matters!
+    KEYWORDS_EXPR = re.compile("(order by| limit | parameters )", flags=RE_FLAGS)
 
     @staticmethod
     def get_query_fields(query: str) -> List[str]:
diff --git a/airbyte-integrations/connectors/source-google-ads/source_google_ads/google_ads.py b/airbyte-integrations/connectors/source-google-ads/source_google_ads/google_ads.py
@@ -4,15 +4,17 @@
 
 
 from enum import Enum
-from typing import Any, Iterator, List, Mapping
+from typing import Any, Iterator, List, Mapping, MutableMapping
 
 import pendulum
 from google.ads.googleads.client import GoogleAdsClient
-from google.ads.googleads.v8.services.types.google_ads_service import GoogleAdsRow, SearchGoogleAdsResponse
+from google.ads.googleads.v9.services.types.google_ads_service import GoogleAdsRow, SearchGoogleAdsResponse
 from proto.marshal.collections import Repeated, RepeatedComposite
 
+
 REPORT_MAPPING = {
     "accounts": "customer",
+    "service_accounts": "customer",
     "ad_group_ads": "ad_group_ad",
     "ad_group_ad_labels": "ad_group_ad_label",
     "ad_groups": "ad_group",
@@ -34,13 +36,11 @@
 class GoogleAds:
     DEFAULT_PAGE_SIZE = 1000
 
-    def __init__(self, credentials: Mapping[str, Any], customer_id: str):
+    def __init__(self, credentials: MutableMapping[str, Any]):
         # `google-ads` library version `14.0.0` and higher requires an additional required parameter `use_proto_plus`.
         # More details can be found here: https://developers.google.com/google-ads/api/docs/client-libs/python/protobuf-messages
         credentials["use_proto_plus"] = True
-
         self.client = GoogleAdsClient.load_from_dict(credentials)
-        self.customer_ids = customer_id.split(",")
         self.ga_service = self.client.get_service("GoogleAdsService")
 
     def send_request(self, query: str, customer_id: str) -> Iterator[SearchGoogleAdsResponse]:
diff --git a/airbyte-integrations/connectors/source-google-ads/source_google_ads/models.py b/airbyte-integrations/connectors/source-google-ads/source_google_ads/models.py
@@ -0,0 +1,26 @@
+from dataclasses import dataclass
+from typing import Any, Iterable, Mapping, Union
+from pendulum import timezone
+from pendulum.tz.timezone import Timezone
+
+
+@dataclass
+class Customer:
+    id: str
+    time_zone: Union[timezone, str] = "local"
+    is_manager_account: bool = False
+
+    @classmethod
+    def from_accounts(cls, accounts: Iterable[Iterable[Mapping[str, Any]]]):
+        data_objects = []
+        for account_list in accounts:
+            for account in account_list:
+                time_zone_name = account.get("customer.time_zone")
+                tz = Timezone(time_zone_name) if time_zone_name else "local"
+
+                data_objects.append(cls(
+                    id=str(account["customer.id"]),
+                    time_zone=tz,
+                    is_manager_account=bool(account.get("customer.manager"))
+                ))
+        return data_objects
diff --git a/airbyte-integrations/connectors/source-google-ads/source_google_ads/schemas/service_accounts.json b/airbyte-integrations/connectors/source-google-ads/source_google_ads/schemas/service_accounts.json
@@ -0,0 +1,69 @@
+{
+  "$schema": "http://json-schema.org/draft-07/schema#",
+  "type": "object",
+  "properties": {
+    "customer.auto_tagging_enabled": {
+      "type": ["null", "boolean"]
+    },
+    "customer.call_reporting_setting.call_conversion_action": {
+      "type": ["null", "string"]
+    },
+    "customer.call_reporting_setting.call_conversion_reporting_enabled": {
+      "type": ["null", "boolean"]
+    },
+    "customer.call_reporting_setting.call_reporting_enabled": {
+      "type": ["null", "boolean"]
+    },
+    "customer.conversion_tracking_setting.conversion_tracking_id": {
+      "type": ["null", "integer"]
+    },
+    "customer.conversion_tracking_setting.cross_account_conversion_tracking_id": {
+      "type": ["null", "integer"]
+    },
+    "customer.currency_code": {
+      "type": ["null", "string"]
+    },
+    "customer.descriptive_name": {
+      "type": ["null", "string"]
+    },
+    "customer.final_url_suffix": {
+      "type": ["null", "string"]
+    },
+    "customer.has_partners_badge": {
+      "type": ["null", "boolean"]
+    },
+    "customer.id": {
+      "type": ["null", "integer"]
+    },
+    "customer.manager": {
+      "type": ["null", "boolean"]
+    },
+    "customer.optimization_score": {
+      "type": ["null", "number"]
+    },
+    "customer.optimization_score_weight": {
+      "type": ["null", "number"]
+    },
+    "customer.pay_per_conversion_eligibility_failure_reasons": {
+      "type": ["null", "array"],
+      "items": {
+        "type": "string"
+      }
+    },
+    "customer.remarketing_setting.google_global_site_tag": {
+      "type": ["null", "string"]
+    },
+    "customer.resource_name": {
+      "type": ["null", "string"]
+    },
+    "customer.test_account": {
+      "type": ["null", "boolean"]
+    },
+    "customer.time_zone": {
+      "type": ["null", "string"]
+    },
+    "customer.tracking_url_template": {
+      "type": ["null", "string"]
+    }
+  }
+}
diff --git a/airbyte-integrations/connectors/source-google-ads/source_google_ads/source.py b/airbyte-integrations/connectors/source-google-ads/source_google_ads/source.py
@@ -4,21 +4,22 @@
 
 
 import traceback
-from typing import Any, List, Mapping, Tuple, Union
+from typing import Any, Iterable, List, Mapping, MutableMapping, Tuple
 
 from airbyte_cdk import AirbyteLogger
 from airbyte_cdk.models import SyncMode
 from airbyte_cdk.sources import AbstractSource
 from airbyte_cdk.sources.streams import Stream
 from google.ads.googleads.errors import GoogleAdsException
 from pendulum import parse, timezone, today
-from pendulum.tz.timezone import Timezone
 
 from .custom_query_stream import CustomQuery
 from .google_ads import GoogleAds
+from .models import Customer
 from .streams import (
     AccountPerformanceReport,
     Accounts,
+    ServiceAccounts,
     AdGroupAdLabels,
     AdGroupAdReport,
     AdGroupAds,
@@ -38,7 +39,7 @@
 
 class SourceGoogleAds(AbstractSource):
     @staticmethod
-    def get_credentials(config: Mapping[str, Any]) -> Mapping[str, Any]:
+    def get_credentials(config: Mapping[str, Any]) -> MutableMapping[str, Any]:
         credentials = config["credentials"]
         # use_proto_plus is set to True, because setting to False returned wrong value types, which breakes the backward compatibility.
         # For more info read the related PR's description: https://github.com/airbytehq/airbyte/pull/9996
@@ -50,36 +51,25 @@ def get_credentials(config: Mapping[str, Any]) -> Mapping[str, Any]:
         return credentials
 
     @staticmethod
-    def get_incremental_stream_config(google_api: GoogleAds, config: Mapping[str, Any], tz: Union[timezone, str] = "local"):
+    def get_incremental_stream_config(google_api: GoogleAds, config: Mapping[str, Any], customers: List[Customer]):
         true_end_date = None
         configured_end_date = config.get("end_date")
         if configured_end_date is not None:
             true_end_date = min(today(), parse(configured_end_date)).to_date_string()
         incremental_stream_config = dict(
             api=google_api,
+            customers=customers,
             conversion_window_days=config["conversion_window_days"],
             start_date=config["start_date"],
-            time_zone=tz,
             end_date=true_end_date,
         )
         return incremental_stream_config
 
-    def get_account_info(self, google_api: GoogleAds, config: Mapping[str, Any]) -> dict:
-        incremental_stream_config = self.get_incremental_stream_config(google_api, config)
-        accounts_streams = Accounts(**incremental_stream_config)
-        for stream_slice in accounts_streams.stream_slices(sync_mode=SyncMode.full_refresh):
-            return next(accounts_streams.read_records(sync_mode=SyncMode.full_refresh, stream_slice=stream_slice), {})
-
-    @staticmethod
-    def get_time_zone(account: dict) -> Union[timezone, str]:
-        time_zone_name = account.get("customer.time_zone")
-        if time_zone_name:
-            return Timezone(time_zone_name)
-        return "local"
-
-    @staticmethod
-    def is_manager_account(account: dict) -> bool:
-        return bool(account.get("customer.manager"))
+    def get_account_info(self, google_api: GoogleAds, config: Mapping[str, Any]) -> Iterable[Iterable[Mapping[str, Any]]]:
+        dummy_customers = [Customer(id=_id) for _id in config["customer_id"].split(",")]
+        accounts_stream = ServiceAccounts(google_api, customers=dummy_customers)
+        for slice_ in accounts_stream.stream_slices():
+            yield accounts_stream.read_records(sync_mode=SyncMode.full_refresh, stream_slice=slice_)
 
     @staticmethod
     def is_metrics_in_custom_query(query: str) -> bool:
@@ -92,42 +82,42 @@ def is_metrics_in_custom_query(query: str) -> bool:
     def check_connection(self, logger: AirbyteLogger, config: Mapping[str, Any]) -> Tuple[bool, any]:
         try:
             logger.info("Checking the config")
-            google_api = GoogleAds(credentials=self.get_credentials(config), customer_id=config["customer_id"])
-            account_info = self.get_account_info(google_api, config)
-            is_manager_account = self.is_manager_account(account_info)
+            google_api = GoogleAds(credentials=self.get_credentials(config))
 
+            accounts = self.get_account_info(google_api, config)
+            customers = Customer.from_accounts(accounts)
             # Check custom query request validity by sending metric request with non-existant time window
-            for query in config.get("custom_queries", []):
-                query = query.get("query")
-
-                if is_manager_account and self.is_metrics_in_custom_query(query):
-                    return False, f"Metrics are not available for manager account. Check fields in your custom query: {query}"
-                if CustomQuery.cursor_field in query:
-                    return False, f"Custom query should not contain {CustomQuery.cursor_field}"
-
-                req_q = CustomQuery.insert_segments_date_expr(query, "1980-01-01", "1980-01-01")
-                for customer_id in google_api.customer_ids:
-                    google_api.send_request(req_q, customer_id=customer_id)
+            for customer in customers:
+                for query in config.get("custom_queries", []):
+                    query = query.get("query")
+                    if customer.is_manager_account and self.is_metrics_in_custom_query(query):
+                        return False, f"Metrics are not available for manager account. Check fields in your custom query: {query}"
+                    if CustomQuery.cursor_field in query:
+                        return False, f"Custom query should not contain {CustomQuery.cursor_field}"
+                    req_q = CustomQuery.insert_segments_date_expr(query, "1980-01-01", "1980-01-01")
+                    response = google_api.send_request(req_q, customer_id=customer.id)
+                    # iterate over the response otherwise exceptions will not be raised!
+                    for _ in response:
+                        pass
             return True, None
         except GoogleAdsException as exception:
             error_messages = ", ".join([error.message for error in exception.failure.errors])
             logger.error(traceback.format_exc())
             return False, f"Unable to connect to Google Ads API with the provided credentials - {error_messages}"
 
     def streams(self, config: Mapping[str, Any]) -> List[Stream]:
-        google_api = GoogleAds(credentials=self.get_credentials(config), customer_id=config["customer_id"])
-        account_info = self.get_account_info(google_api, config)
-        time_zone = self.get_time_zone(account_info)
-        incremental_stream_config = self.get_incremental_stream_config(google_api, config, tz=time_zone)
-
+        google_api = GoogleAds(credentials=self.get_credentials(config))
+        accounts = self.get_account_info(google_api, config)
+        customers = Customer.from_accounts(accounts)
+        incremental_stream_config = self.get_incremental_stream_config(google_api, config, customers)
         streams = [
             AdGroupAds(**incremental_stream_config),
-            AdGroupAdLabels(google_api),
+            AdGroupAdLabels(google_api, customers=customers),
             AdGroups(**incremental_stream_config),
-            AdGroupLabels(google_api),
+            AdGroupLabels(google_api, customers=customers),
             Accounts(**incremental_stream_config),
             Campaigns(**incremental_stream_config),
-            CampaignLabels(google_api),
+            CampaignLabels(google_api, customers=customers),
             ClickView(**incremental_stream_config),
         ]
         custom_query_streams = [
@@ -137,7 +127,9 @@ def streams(self, config: Mapping[str, Any]) -> List[Stream]:
         streams.extend(custom_query_streams)
 
         # Metrics streams cannot be requested for a manager account.
-        if not self.is_manager_account(account_info):
+        non_manager_accounts = [customer for customer in customers if not customer.is_manager_account]
+        if non_manager_accounts:
+            incremental_stream_config["customers"] = non_manager_accounts
             streams.extend(
                 [
                     UserLocationReport(**incremental_stream_config),
diff --git a/airbyte-integrations/connectors/source-google-ads/source_google_ads/spec.json b/airbyte-integrations/connectors/source-google-ads/source_google_ads/spec.json
@@ -90,7 +90,7 @@
             "query": {
               "type": "string",
               "title": "Custom Query",
-              "description": "A custom defined GAQL query for building the report. Should not contain segments.date expression because it is used by incremental streams. See Google's <a href=\"https://developers.google.com/google-ads/api/fields/v8/overview_query_builder\">query builder</a> for more information.",
+              "description": "A custom defined GAQL query for building the report. Should not contain segments.date expression because it is used by incremental streams. See Google's <a href=\"https://developers.google.com/google-ads/api/fields/v9/overview_query_builder\">query builder</a> for more information.",
               "examples": [
                 "SELECT segments.ad_destination_type, campaign.advertising_channel_sub_type FROM campaign WHERE campaign.status = 'PAUSED'"
               ]
diff --git a/airbyte-integrations/connectors/source-google-ads/source_google_ads/streams.py b/airbyte-integrations/connectors/source-google-ads/source_google_ads/streams.py
diff --git a/airbyte-integrations/connectors/source-google-ads/unit_tests/test_google_ads.py b/airbyte-integrations/connectors/source-google-ads/unit_tests/test_google_ads.py

Original file line number	Diff line number	Diff line change
`@@ -90,7 +90,7 @@`
`90`	`90`	`"query": {`
`91`	`91`	`"type": "string",`
`92`	`92`	`"title": "Custom Query",`
`93`		`- "description": "A custom defined GAQL query for building the report. Should not contain segments.date expression because it is used by incremental streams. See Google's <a href=\"https://developers.google.com/google-ads/api/fields/v8/overview_query_builder\">query builder</a> for more information.",`
	`93`	`+ "description": "A custom defined GAQL query for building the report. Should not contain segments.date expression because it is used by incremental streams. See Google's <a href=\"https://developers.google.com/google-ads/api/fields/v9/overview_query_builder\">query builder</a> for more information.",`
`94`	`94`	`"examples": [`
`95`	`95`	`"SELECT segments.ad_destination_type, campaign.advertising_channel_sub_type FROM campaign WHERE campaign.status = 'PAUSED'"`
`96`	`96`	`]`