Use api version only in GoogleAdsHook not operators (#15266)

Jacob Kim · web-flow · commit aa4713e43f92 · 2021-05-21T10:38:37.000+02:00
Add parameter api_version to both GoogleAdsToGcsOperator &amp; GoogleAdsListAccountsOperator.
Remove repeated line self.gcp_conn_id = gcp_conn_id.
Change the default api_version to v5 since v3 is deprecated.
Add api_version to GoogleAdsHook's docstring.
diff --git a/airflow/providers/google/ads/hooks/ads.py b/airflow/providers/google/ads/hooks/ads.py
@@ -17,7 +17,7 @@
 # under the License.
 """This module contains Google Ad hook."""
 from tempfile import NamedTemporaryFile
-from typing import IO, Any, Dict, Generator, List
+from typing import IO, Any, Dict, Generator, List, Optional
 
 try:
     from functools import cached_property
@@ -69,22 +69,25 @@ class GoogleAdsHook(BaseHook):
     :type gcp_conn_id: str
     :param google_ads_conn_id: The connection ID with the details of Google Ads config.yaml file.
     :type google_ads_conn_id: str
+    :param api_version: The Google Ads API version to use.
+    :type api_version: str
 
     :return: list of Google Ads Row object(s)
     :rtype: list[GoogleAdsRow]
     """
 
+    default_api_version = "v5"
+
     def __init__(
         self,
+        api_version: Optional[str],
         gcp_conn_id: str = "google_cloud_default",
         google_ads_conn_id: str = "google_ads_default",
-        api_version: str = "v3",
     ) -> None:
         super().__init__()
+        self.api_version = api_version or self.default_api_version
         self.gcp_conn_id = gcp_conn_id
         self.google_ads_conn_id = google_ads_conn_id
-        self.gcp_conn_id = gcp_conn_id
-        self.api_version = api_version
         self.google_ads_config: Dict[str, Any] = {}
 
     @cached_property
diff --git a/airflow/providers/google/ads/operators/ads.py b/airflow/providers/google/ads/operators/ads.py
@@ -61,6 +61,8 @@ class GoogleAdsListAccountsOperator(BaseOperator):
         Service Account Token Creator IAM role to the directly preceding identity, with first
         account from the list granting this role to the originating account (templated).
     :type impersonation_chain: Union[str, Sequence[str]]
+    :param api_version: Optional Google Ads API version to use.
+    :type api_version: Optional[str]
     """
 
     template_fields = (
@@ -78,6 +80,7 @@ def __init__(
         google_ads_conn_id: str = "google_ads_default",
         gzip: bool = False,
         impersonation_chain: Optional[Union[str, Sequence[str]]] = None,
+        api_version: Optional[str] = None,
         **kwargs,
     ) -> None:
         super().__init__(**kwargs)
@@ -87,11 +90,16 @@ def __init__(
         self.google_ads_conn_id = google_ads_conn_id
         self.gzip = gzip
         self.impersonation_chain = impersonation_chain
+        self.api_version = api_version
 
     def execute(self, context: dict) -> str:
         uri = f"gs://{self.bucket}/{self.object_name}"
 
-        ads_hook = GoogleAdsHook(gcp_conn_id=self.gcp_conn_id, google_ads_conn_id=self.google_ads_conn_id)
+        ads_hook = GoogleAdsHook(
+            gcp_conn_id=self.gcp_conn_id,
+            google_ads_conn_id=self.google_ads_conn_id,
+            api_version=self.api_version,
+        )
 
         gcs_hook = GCSHook(gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain)
         with NamedTemporaryFile("w+") as temp_file:
diff --git a/airflow/providers/google/ads/transfers/ads_to_gcs.py b/airflow/providers/google/ads/transfers/ads_to_gcs.py
@@ -66,6 +66,8 @@ class GoogleAdsToGcsOperator(BaseOperator):
         Service Account Token Creator IAM role to the directly preceding identity, with first
         account from the list granting this role to the originating account (templated).
     :type impersonation_chain: Union[str, Sequence[str]]
+    :param api_version: Optional Google Ads API version to use.
+    :type api_version: Optional[str]
     """
 
     template_fields = (
@@ -90,6 +92,7 @@ def __init__(
         page_size: int = 10000,
         gzip: bool = False,
         impersonation_chain: Optional[Union[str, Sequence[str]]] = None,
+        api_version: Optional[str] = None,
         **kwargs,
     ) -> None:
         super().__init__(**kwargs)
@@ -103,9 +106,14 @@ def __init__(
         self.page_size = page_size
         self.gzip = gzip
         self.impersonation_chain = impersonation_chain
+        self.api_version = api_version
 
     def execute(self, context: dict) -> None:
-        service = GoogleAdsHook(gcp_conn_id=self.gcp_conn_id, google_ads_conn_id=self.google_ads_conn_id)
+        service = GoogleAdsHook(
+            gcp_conn_id=self.gcp_conn_id,
+            google_ads_conn_id=self.google_ads_conn_id,
+            api_version=self.api_version,
+        )
         rows = service.search(client_ids=self.client_ids, query=self.query, page_size=self.page_size)
 
         try:
diff --git a/tests/providers/google/ads/operators/test_ads.py b/tests/providers/google/ads/operators/test_ads.py
@@ -37,6 +37,7 @@
 
 gcp_conn_id = "gcp_conn_id"
 google_ads_conn_id = "google_ads_conn_id"
+api_version = "v5"
 
 
 class TestGoogleAdsListAccountsOperator:
@@ -58,10 +59,15 @@ def test_execute(self, mocks_csv_writer, mock_tempfile, mock_gcs_hook, mock_ads_
             bucket=BUCKET,
             task_id="run_operator",
             impersonation_chain=IMPERSONATION_CHAIN,
+            api_version=api_version,
         )
         op.execute({})
 
-        mock_ads_hook.assert_called_once_with(gcp_conn_id=gcp_conn_id, google_ads_conn_id=google_ads_conn_id)
+        mock_ads_hook.assert_called_once_with(
+            gcp_conn_id=gcp_conn_id,
+            google_ads_conn_id=google_ads_conn_id,
+            api_version=api_version,
+        )
         mock_gcs_hook.assert_called_once_with(
             gcp_conn_id=gcp_conn_id,
             impersonation_chain=IMPERSONATION_CHAIN,
diff --git a/tests/providers/google/ads/transfers/test_ads_to_gcs.py b/tests/providers/google/ads/transfers/test_ads_to_gcs.py
@@ -25,6 +25,7 @@
     GCS_OBJ_PATH,
     IMPERSONATION_CHAIN,
     QUERY,
+    api_version,
     gcp_conn_id,
     google_ads_conn_id,
 )
@@ -44,9 +45,14 @@ def test_execute(self, mock_gcs_hook, mock_ads_hook):
             bucket=BUCKET,
             task_id="run_operator",
             impersonation_chain=IMPERSONATION_CHAIN,
+            api_version=api_version,
         )
         op.execute({})
-        mock_ads_hook.assert_called_once_with(gcp_conn_id=gcp_conn_id, google_ads_conn_id=google_ads_conn_id)
+        mock_ads_hook.assert_called_once_with(
+            gcp_conn_id=gcp_conn_id,
+            google_ads_conn_id=google_ads_conn_id,
+            api_version=api_version,
+        )
         mock_ads_hook.return_value.search.assert_called_once_with(
             client_ids=CLIENT_IDS, query=QUERY, page_size=10000
         )
