Added Upload Multiple Entity Read Files to specified big query dataset (#8610)

michalslowikowski00 · michalslowikowski00 · web-flow · commit 79ef8bed891c · 2020-05-10T13:27:14.000+02:00
Co-authored-by: michalslowikowski00 &lt;michal.slowikowski@polidea.com&gt;
diff --git a/airflow/providers/google/marketing_platform/example_dags/example_display_video.py b/airflow/providers/google/marketing_platform/example_dags/example_display_video.py
@@ -23,6 +23,7 @@
 
 from airflow import models
 from airflow.providers.google.cloud.operators.gcs_to_bigquery import GCSToBigQueryOperator
+from airflow.providers.google.marketing_platform.hooks.display_video import GoogleDisplayVideo360Hook
 from airflow.providers.google.marketing_platform.operators.display_video import (
     GoogleDisplayVideo360CreateReportOperator, GoogleDisplayVideo360CreateSDFDownloadTaskOperator,
     GoogleDisplayVideo360DeleteReportOperator, GoogleDisplayVideo360DownloadLineItemsOperator,
@@ -38,15 +39,14 @@
 BUCKET = os.environ.get("GMP_DISPLAY_VIDEO_BUCKET", "gs://test-display-video-bucket")
 ADVERTISER_ID = os.environ.get("GMP_ADVERTISER_ID", 1234567)
 OBJECT_NAME = os.environ.get("GMP_OBJECT_NAME", "files/report.csv")
-PATH_TO_UPLOAD_FILE = os.environ.get(
-    "GCP_GCS_PATH_TO_UPLOAD_FILE", "test-gcs-example.txt"
-)
-PATH_TO_SAVED_FILE = os.environ.get(
-    "GCP_GCS_PATH_TO_SAVED_FILE", "test-gcs-example-download.txt"
-)
+PATH_TO_UPLOAD_FILE = os.environ.get("GCP_GCS_PATH_TO_UPLOAD_FILE", "test-gcs-example.txt")
+PATH_TO_SAVED_FILE = os.environ.get("GCP_GCS_PATH_TO_SAVED_FILE", "test-gcs-example-download.txt")
 BUCKET_FILE_LOCATION = PATH_TO_UPLOAD_FILE.rpartition("/")[-1]
 SDF_VERSION = os.environ.get("GMP_SDF_VERSION", "SDF_VERSION_5_1")
 BQ_DATA_SET = os.environ.get("GMP_BQ_DATA_SET", "airflow_test")
+GMP_PARTNER_ID = os.environ.get("GMP_PARTNER_ID", 123)
+ENTITY_TYPE = os.environ.get("GMP_ENTITY_TYPE", "LineItem")
+ERF_SOURCE_OBJECT = GoogleDisplayVideo360Hook.erf_uri(GMP_PARTNER_ID, ENTITY_TYPE)
 
 REPORT = {
     "kind": "doubleclickbidmanager#query",
@@ -68,15 +68,17 @@
 
 PARAMS = {"dataRange": "LAST_14_DAYS", "timezoneCode": "America/New_York"}
 
-BODY_REQUEST: Dict = {
+CREATE_SDF_DOWNLOAD_TASK_BODY_REQUEST: Dict = {
     "version": SDF_VERSION,
     "advertiserId": ADVERTISER_ID,
     "inventorySourceFilter": {"inventorySourceIds": []},
 }
-# [END howto_display_video_env_variables]
 
-# download_line_items variables
-REQUEST_BODY = {"filterType": ADVERTISER_ID, "format": "CSV", "fileSpec": "EWF"}
+DOWNLOAD_LINE_ITEMS_REQUEST: Dict = {
+    "filterType": ADVERTISER_ID,
+    "format": "CSV",
+    "fileSpec": "EWF"}
+# [END howto_display_video_env_variables]
 
 default_args = {"start_date": dates.days_ago(1)}
 
@@ -119,10 +121,20 @@
     )
     # [END howto_google_display_video_deletequery_report_operator]
 
+    # [START howto_google_display_video_upload_multiple_entity_read_files_to_big_query]
+    upload_erf_to_bq = GCSToBigQueryOperator(
+        task_id='upload_erf_to_bq',
+        bucket=BUCKET,
+        source_objects=ERF_SOURCE_OBJECT,
+        destination_project_dataset_table=f"{BQ_DATA_SET}.gcs_to_bq_table",
+        write_disposition='WRITE_TRUNCATE',
+        dag=dag)
+    # [END howto_google_display_video_upload_multiple_entity_read_files_to_big_query]
+
     # [START howto_google_display_video_download_line_items_operator]
     download_line_items = GoogleDisplayVideo360DownloadLineItemsOperator(
         task_id="download_line_items",
-        request_body=REQUEST_BODY,
+        request_body=DOWNLOAD_LINE_ITEMS_REQUEST,
         bucket_name=BUCKET,
         object_name=OBJECT_NAME,
         gzip=False,
@@ -139,7 +151,7 @@
 
     # [START howto_google_display_video_create_sdf_download_task_operator]
     create_sdf_download_task = GoogleDisplayVideo360CreateSDFDownloadTaskOperator(
-        task_id="create_sdf_download_task", body_request=BODY_REQUEST
+        task_id="create_sdf_download_task", body_request=CREATE_SDF_DOWNLOAD_TASK_BODY_REQUEST
     )
     operation_name = '{{ task_instance.xcom_pull("create_sdf_download_task")["name"] }}'
     # [END howto_google_display_video_create_sdf_download_task_operator]
diff --git a/airflow/providers/google/marketing_platform/hooks/display_video.py b/airflow/providers/google/marketing_platform/hooks/display_video.py
@@ -70,6 +70,28 @@ def get_conn_to_display_video(self) -> Resource:
             )
         return self._conn
 
+    @staticmethod
+    def erf_uri(partner_id, entity_type) -> List[str]:
+        """
+        Return URI for all Entity Read Files in bucket.
+
+        For example, if you were generating a file name to retrieve the entity read file
+        for partner 123 accessing the line_item table from April 2, 2013, your filename
+        would look something like this:
+        gdbm-123/entity/20130402.0.LineItem.json
+
+        More information:
+        https://developers.google.com/bid-manager/guides/entity-read/overview
+
+        :param partner_id The numeric ID of your Partner.
+        :type partner_id: int
+        :param entity_type: The type of file Partner, Advertiser, InsertionOrder,
+        LineItem, Creative, Pixel, InventorySource, UserList, UniversalChannel, and summary.
+        :type entity_type: str
+        """
+
+        return [f"gdbm-{partner_id}/entity/{{{{ ds_nodash }}}}.*.{entity_type}.json"]
+
     def create_query(self, query: Dict[str, Any]) -> Dict:
         """
         Creates a query.
@@ -125,7 +147,7 @@ def list_queries(self, ) -> List[Dict]:
             .listqueries()
             .execute(num_retries=self.num_retries)
         )
-        return response.get("queries", [])
+        return response.get('queries', [])
 
     def run_query(self, query_id: str, params: Dict[str, Any]) -> None:
         """
