[AIRFLOW-5379] Add Google Search Ads 360 operators (#6228)

apache · Oct 22, 2019 · 4e661f5 · 4e661f5
1 parent b439221
commit 4e661f5
Show file tree

Hide file tree

Showing 12 changed files with 924 additions and 0 deletions.
diff --git a/airflow/providers/google/marketing_platform/example_dags/example_search_ads.py b/airflow/providers/google/marketing_platform/example_dags/example_search_ads.py
@@ -0,0 +1,76 @@
+# -*- coding: utf-8 -*-
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Example Airflow DAG that shows how to use SearchAds.
+"""
+import os
+
+from airflow import models
+from airflow.providers.google.marketing_platform.operators.search_ads import (
+ GoogleSearchAdsDownloadReportOperator, GoogleSearchAdsInsertReportOperator,
+)
+from airflow.providers.google.marketing_platform.sensors.search_ads import GoogleSearchAdsReportSensor
+from airflow.utils import dates
+
+# [START howto_search_ads_env_variables]
+AGENCY_ID = os.environ.get("GMP_AGENCY_ID")
+ADVERTISER_ID = os.environ.get("GMP_ADVERTISER_ID")
+GCS_BUCKET = os.environ.get("GMP_GCS_BUCKET", "test-cm-bucket")
+
+REPORT = {
+ "reportScope": {"agencyId": AGENCY_ID, "advertiserId": ADVERTISER_ID},
+ "reportType": "account",
+ "columns": [{"columnName": "agency"}, {"columnName": "lastModifiedTimestamp"}],
+ "includeRemovedEntities": False,
+ "statisticsCurrency": "usd",
+ "maxRowsPerFile": 1000000,
+ "downloadFormat": "csv",
+}
+# [END howto_search_ads_env_variables]
+
+default_args = {"start_date": dates.days_ago(1)}
+
+with models.DAG(
+ "example_search_ads",
+ default_args=default_args,
+ schedule_interval=None, # Override to match your needs
+) as dag:
+ # [START howto_search_ads_generate_report_operator]
+ generate_report = GoogleSearchAdsInsertReportOperator(
+ report=REPORT, task_id="generate_report"
+ )
+ # [END howto_search_ads_generate_report_operator]
+
+ # [START howto_search_ads_get_report_id]
+ report_id = "{{ task_instance.xcom_pull('generate_report', key='report_id') }}"
+ # [END howto_search_ads_get_report_id]
+
+ # [START howto_search_ads_get_report_operator]
+ wait_for_report = GoogleSearchAdsReportSensor(
+ report_id=report_id, task_id="wait_for_report"
+ )
+ # [END howto_search_ads_get_report_operator]
+
+ # [START howto_search_ads_getfile_report_operator]
+ download_report = GoogleSearchAdsDownloadReportOperator(
+ report_id=report_id, bucket_name=GCS_BUCKET, task_id="download_report"
+ )
+ # [END howto_search_ads_getfile_report_operator]
+
+ generate_report >> wait_for_report >> download_report
diff --git a/airflow/providers/google/marketing_platform/hooks/search_ads.py b/airflow/providers/google/marketing_platform/hooks/search_ads.py
@@ -0,0 +1,104 @@
+# -*- coding: utf-8 -*-
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+This module contains Google Search Ads 360 hook.
+"""
+from typing import Any, Dict, Optional
+
+from googleapiclient.discovery import build
+
+from airflow.gcp.hooks.base import GoogleCloudBaseHook
+
+
+class GoogleSearchAdsHook(GoogleCloudBaseHook):
+ """
+ Hook for Google Search Ads 360.
+ """
+
+ _conn = None # type: Optional[Any]
+
+ def __init__(
+ self,
+ api_version: str = "v2",
+ gcp_conn_id: str = "google_cloud_default",
+ delegate_to: Optional[str] = None,
+ ) -> None:
+ super().__init__(gcp_conn_id, delegate_to)
+ self.api_version = api_version
+
+ def get_conn(self):
+ """
+ Retrieves connection to Google SearchAds.
+ """
+ if not self._conn:
+ http_authorized = self._authorize()
+ self._conn = build(
+ "doubleclicksearch",
+ self.api_version,
+ http=http_authorized,
+ cache_discovery=False,
+ )
+ return self._conn
+
+ def insert_report(self, report: Dict[str, Any]) -> Any:
+ """
+ Inserts a report request into the reporting system.
+
+ :param report: Report to be generated.
+ :type report: Dict[str, Any]
+ """
+ response = (
+ self.get_conn() # pylint: disable=no-member
+ .reports()
+ .request(body=report)
+ .execute(num_retries=self.num_retries)
+ )
+ return response
+
+ def get(self, report_id: str) -> Any:
+ """
+ Polls for the status of a report request.
+
+ :param report_id: ID of the report request being polled.
+ :type report_id: str
+ """
+ response = (
+ self.get_conn() # pylint: disable=no-member
+ .reports()
+ .get(reportId=report_id)
+ .execute(num_retries=self.num_retries)
+ )
+ return response
+
+ def get_file(self, report_fragment: int, report_id: str) -> Any:
+ """
+ Downloads a report file encoded in UTF-8.
+
+ :param report_fragment: The index of the report fragment to download.
+ :type report_fragment: int
+ :param report_id: ID of the report.
+ :type report_id: str
+ """
+ response = (
+ self.get_conn() # pylint: disable=no-member
+ .reports()
+ .getFile(reportFragment=report_fragment, reportId=report_id)
+ .execute(num_retries=self.num_retries)
+ )
+ return response
diff --git a/airflow/providers/google/marketing_platform/operators/search_ads.py b/airflow/providers/google/marketing_platform/operators/search_ads.py
@@ -0,0 +1,210 @@
+# -*- coding: utf-8 -*-
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+This module contains Google Search Ads operators.
+"""
+from tempfile import NamedTemporaryFile
+from typing import Any, Dict, Optional
+
+from airflow import AirflowException
+from airflow.gcp.hooks.gcs import GoogleCloudStorageHook
+from airflow.models.baseoperator import BaseOperator
+from airflow.providers.google.marketing_platform.hooks.search_ads import GoogleSearchAdsHook
+from airflow.utils.decorators import apply_defaults
+
+
+class GoogleSearchAdsInsertReportOperator(BaseOperator):
+ """
+ Inserts a report request into the reporting system.
+
+ .. seealso:
+ For API documentation check:
+ https://developers.google.com/search-ads/v2/reference/reports/request
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:GoogleSearchAdsInsertReportOperator`
+
+ :param report: Report to be generated
+ :type report: Dict[str, Any]
+ :param api_version: The version of the api that will be requested for example 'v3'.
+ :type api_version: str
+ :param gcp_conn_id: The connection ID to use when fetching connection info.
+ :type gcp_conn_id: str
+ :param delegate_to: The account to impersonate, if any. For this to work, the service accountmaking the
+ request must have domain-wide delegation enabled.
+ :type delegate_to: str
+ """
+
+ template_fields = ("report",)
+ template_ext = (".json",)
+
+ @apply_defaults
+ def __init__(
+ self,
+ report: Dict[str, Any],
+ api_version: str = "v2",
+ gcp_conn_id: str = "google_cloud_default",
+ delegate_to: Optional[str] = None,
+ *args,
+ **kwargs
+ ):
+ super().__init__(*args, **kwargs)
+ self.report = report
+ self.api_version = api_version
+ self.gcp_conn_id = gcp_conn_id
+ self.delegate_to = delegate_to
+
+ def execute(self, context: Dict):
+ hook = GoogleSearchAdsHook(
+ gcp_conn_id=self.gcp_conn_id,
+ delegate_to=self.delegate_to,
+ api_version=self.api_version,
+ )
+ self.log.info("Generating Search Ads report")
+ response = hook.insert_report(report=self.report)
+ report_id = response.get("id")
+ self.xcom_push(context, key="report_id", value=report_id)
+ self.log.info("Report generated, id: %s", report_id)
+ return response
+
+
+class GoogleSearchAdsDownloadReportOperator(BaseOperator):
+ """
+ Downloads a report to GCS bucket.
+
+ .. seealso:
+ For API documentation check:
+ https://developers.google.com/search-ads/v2/reference/reports/getFile
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:GoogleSearchAdsGetfileReportOperator`
+
+ :param report_id: ID of the report.
+ :type report_id: str
+ :param bucket_name: The bucket to upload to.
+ :type bucket_name: str
+ :param report_name: The report name to set when uploading the local file. If not provided then
+ report_id is used.
+ :type report_name: str
+ :param gzip: Option to compress local file or file data for upload
+ :type gzip: bool
+ :param api_version: The version of the api that will be requested for example 'v3'.
+ :type api_version: str
+ :param gcp_conn_id: The connection ID to use when fetching connection info.
+ :type gcp_conn_id: str
+ :param delegate_to: The account to impersonate, if any. For this to work, the service accountmaking the
+ request must have domain-wide delegation enabled.
+ :type delegate_to: str
+ """
+
+ template_fields = ("report_name", "report_id", "bucket_name")
+
+ @apply_defaults
+ def __init__(
+ self,
+ report_id: str,
+ bucket_name: str,
+ report_name: Optional[str] = None,
+ gzip: bool = True,
+ chunk_size: int = 10 * 1024 * 1024,
+ api_version: str = "v2",
+ gcp_conn_id: str = "google_cloud_default",
+ delegate_to: Optional[str] = None,
+ *args,
+ **kwargs
+ ) -> None:
+ super().__init__(*args, **kwargs)
+ self.report_id = report_id
+ self.api_version = api_version
+ self.gcp_conn_id = gcp_conn_id
+ self.delegate_to = delegate_to
+ self.report_id = report_id
+ self.chunk_size = chunk_size
+ self.gzip = gzip
+ self.bucket_name = self._set_bucket_name(bucket_name)
+ self.report_name = report_name
+
+ def _resolve_file_name(self, name: str) -> str:
+ csv = ".csv"
+ gzip = ".gz"
+ if not name.endswith(csv):
+ name += csv
+ if self.gzip:
+ name += gzip
+ return name
+
+ @staticmethod
+ def _set_bucket_name(name: str) -> str:
+ bucket = name if not name.startswith("gs://") else name[5:]
+ return bucket.strip("/")
+
+ @staticmethod
+ def _handle_report_fragment(fragment: bytes) -> bytes:
+ fragment_records = fragment.split(b"\n", 1)
+ if len(fragment_records) > 1:
+ return fragment_records[1]
+ return b""
+
+ def execute(self, context: Dict):
+ hook = GoogleSearchAdsHook(
+ gcp_conn_id=self.gcp_conn_id,
+ delegate_to=self.delegate_to,
+ api_version=self.api_version,
+ )
+
+ gcs_hook = GoogleCloudStorageHook(
+ gcp_conn_id=self.gcp_conn_id, delegate_to=self.delegate_to
+ )
+
+ # Resolve file name of the report
+ report_name = self.report_name or self.report_id
+ report_name = self._resolve_file_name(report_name)
+
+ response = hook.get(report_id=self.report_id)
+ if not response['isReportReady']:
+ raise AirflowException('Report {} is not ready yet'.format(self.report_id))
+
+ # Resolve report fragments
+ fragments_count = len(response["files"])
+
+ # Download chunks of report's data
+ self.log.info("Downloading Search Ads report %s", self.report_id)
+ with NamedTemporaryFile() as temp_file:
+ for i in range(fragments_count):
+ byte_content = hook.get_file(
+ report_fragment=i, report_id=self.report_id
+ )
+ fragment = (
+ byte_content
+ if i == 0
+ else self._handle_report_fragment(byte_content)
+ )
+ temp_file.write(fragment)
+
+ temp_file.flush()
+
+ gcs_hook.upload(
+ bucket_name=self.bucket_name,
+ object_name=report_name,
+ gzip=self.gzip,
+ filename=temp_file.name,
+ )
+ self.xcom_push(context, key="file_name", value=report_name)