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

Add parameter api_version to both GoogleAdsToGcsOperator & 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.
apache · May 21, 2021 · aa4713e · aa4713e
1 parent 85b2ccb
commit aa4713e
Show file tree

Hide file tree

Showing 5 changed files with 39 additions and 8 deletions.
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
  )