From cd3d9d93402f06a08f35e3586802f11a18c4f1f3 Mon Sep 17 00:00:00 2001 From: Tomek Urbaszek Date: Thu, 2 Jul 2020 10:43:44 +0200 Subject: [PATCH] Fix using .json template extension in GMP operators (#9566) --- .../operators/campaign_manager.py | 8 +++++++- .../operators/display_video.py | 7 +++++++ .../operators/search_ads.py | 7 +++++++ docs/howto/operator/gcp/campaign_manager.rst | 3 ++- docs/howto/operator/gcp/display_video.rst | 3 ++- docs/howto/operator/gcp/search_ads.rst | 3 ++- .../operators/test_campaign_manager.py | 19 +++++++++++++++++++ .../operators/test_display_video.py | 15 +++++++++++++++ .../operators/test_search_ads.py | 19 +++++++++++++++++-- 9 files changed, 78 insertions(+), 6 deletions(-) diff --git a/airflow/providers/google/marketing_platform/operators/campaign_manager.py b/airflow/providers/google/marketing_platform/operators/campaign_manager.py index 4814c9a851..b3cd0c81fc 100644 --- a/airflow/providers/google/marketing_platform/operators/campaign_manager.py +++ b/airflow/providers/google/marketing_platform/operators/campaign_manager.py @@ -18,6 +18,7 @@ """ This module contains Google CampaignManager operators. """ +import json import tempfile import uuid from typing import Any, Dict, List, Optional @@ -298,6 +299,12 @@ class GoogleCampaignManagerInsertReportOperator(BaseOperator): self.gcp_conn_id = gcp_conn_id self.delegate_to = delegate_to + def prepare_template(self) -> None: + # If .json is passed then we have to read the file + if isinstance(self.report, str) and self.report.endswith('.json'): + with open(self.report, 'r') as file: + self.report = json.load(file) + def execute(self, context: Dict): hook = GoogleCampaignManagerHook( gcp_conn_id=self.gcp_conn_id, @@ -349,7 +356,6 @@ class GoogleCampaignManagerRunReportOperator(BaseOperator): "gcp_conn_id", "delegate_to", ) - template_ext = (".json",) @apply_defaults def __init__( diff --git a/airflow/providers/google/marketing_platform/operators/display_video.py b/airflow/providers/google/marketing_platform/operators/display_video.py index 0b0af20d07..8f5f5b4a85 100644 --- a/airflow/providers/google/marketing_platform/operators/display_video.py +++ b/airflow/providers/google/marketing_platform/operators/display_video.py @@ -19,6 +19,7 @@ This module contains Google DisplayVideo operators. """ import csv +import json import shutil import tempfile import urllib.request @@ -75,6 +76,12 @@ class GoogleDisplayVideo360CreateReportOperator(BaseOperator): self.gcp_conn_id = gcp_conn_id self.delegate_to = delegate_to + def prepare_template(self) -> None: + # If .json is passed then we have to read the file + if isinstance(self.body, str) and self.body.endswith('.json'): + with open(self.body, 'r') as file: + self.body = json.load(file) + def execute(self, context: Dict): hook = GoogleDisplayVideo360Hook( gcp_conn_id=self.gcp_conn_id, diff --git a/airflow/providers/google/marketing_platform/operators/search_ads.py b/airflow/providers/google/marketing_platform/operators/search_ads.py index ef88315750..4f2200f15a 100644 --- a/airflow/providers/google/marketing_platform/operators/search_ads.py +++ b/airflow/providers/google/marketing_platform/operators/search_ads.py @@ -18,6 +18,7 @@ """ This module contains Google Search Ads operators. """ +import json from tempfile import NamedTemporaryFile from typing import Any, Dict, Optional @@ -70,6 +71,12 @@ class GoogleSearchAdsInsertReportOperator(BaseOperator): self.gcp_conn_id = gcp_conn_id self.delegate_to = delegate_to + def prepare_template(self) -> None: + # If .json is passed then we have to read the file + if isinstance(self.report, str) and self.report.endswith('.json'): + with open(self.report, 'r') as file: + self.report = json.load(file) + def execute(self, context: Dict): hook = GoogleSearchAdsHook( gcp_conn_id=self.gcp_conn_id, diff --git a/docs/howto/operator/gcp/campaign_manager.rst b/docs/howto/operator/gcp/campaign_manager.rst index a82de6942e..80b3b90cb2 100644 --- a/docs/howto/operator/gcp/campaign_manager.rst +++ b/docs/howto/operator/gcp/campaign_manager.rst @@ -104,7 +104,8 @@ Running this operator creates a new report. You can use :ref:`Jinja templating ` with :template-fields:`airflow.providers.google.marketing_platform.operators.campaign_manager.GoogleCampaignManagerInsertReportOperator` -parameters which allows you to dynamically determine values. +parameters which allows you to dynamically determine values. You can provide report definition using +``.json`` file as this operator supports this template extension. The result is saved to :ref:`XCom `, which allows it to be used by other operators. .. _howto/operator:GoogleCampaignManagerRunReportOperator: diff --git a/docs/howto/operator/gcp/display_video.rst b/docs/howto/operator/gcp/display_video.rst index 04f4814769..26832eebc9 100644 --- a/docs/howto/operator/gcp/display_video.rst +++ b/docs/howto/operator/gcp/display_video.rst @@ -45,7 +45,8 @@ To create Display&Video 360 report use Use :ref:`Jinja templating ` with :template-fields:`airflow.providers.google.marketing_platform.operators.display_video.GoogleDisplayVideo360CreateReportOperator` -parameters which allow you to dynamically determine values. +parameters which allow you to dynamically determine values. You can provide body definition using `` +.json`` file as this operator supports this template extension. The result is saved to :ref:`XCom `, which allows the result to be used by other operators. .. _howto/operator:GoogleDisplayVideo360DeleteReportOperator: diff --git a/docs/howto/operator/gcp/search_ads.rst b/docs/howto/operator/gcp/search_ads.rst index e157caeee4..38a3e1386a 100644 --- a/docs/howto/operator/gcp/search_ads.rst +++ b/docs/howto/operator/gcp/search_ads.rst @@ -46,7 +46,8 @@ To insert a Search Ads report use the You can use :ref:`Jinja templating ` with :template-fields:`airflow.providers.google.marketing_platform.operators.search_ads.GoogleSearchAdsInsertReportOperator` -parameters which allows you to dynamically determine values. +parameters which allows you to dynamically determine values. You can provide report definition using `` +.json`` file as this operator supports this template extension. The result is saved to :ref:`XCom `, which allows it to be used by other operators: .. exampleinclude:: ../../../../airflow/providers/google/marketing_platform/example_dags/example_search_ads.py diff --git a/tests/providers/google/marketing_platform/operators/test_campaign_manager.py b/tests/providers/google/marketing_platform/operators/test_campaign_manager.py index bc2bd899bc..07bf9e8b50 100644 --- a/tests/providers/google/marketing_platform/operators/test_campaign_manager.py +++ b/tests/providers/google/marketing_platform/operators/test_campaign_manager.py @@ -15,6 +15,8 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +import json +from tempfile import NamedTemporaryFile from unittest import TestCase, mock from airflow.providers.google.marketing_platform.operators.campaign_manager import ( @@ -183,6 +185,23 @@ class TestGoogleCampaignManagerInsertReportOperator(TestCase): ) xcom_mock.assert_called_once_with(None, key="report_id", value=report_id) + def test_prepare_template(self): + profile_id = "PROFILE_ID" + report = {"key": "value"} + with NamedTemporaryFile("w+", suffix=".json") as f: + f.write(json.dumps(report)) + f.flush() + op = GoogleCampaignManagerInsertReportOperator( + profile_id=profile_id, + report=f.name, + api_version=API_VERSION, + task_id="test_task", + ) + op.prepare_template() + + assert isinstance(op.report, dict) + assert op.report == report + class TestGoogleCampaignManagerRunReportOperator(TestCase): @mock.patch( diff --git a/tests/providers/google/marketing_platform/operators/test_display_video.py b/tests/providers/google/marketing_platform/operators/test_display_video.py index febcf3d120..fa8ffc65e4 100644 --- a/tests/providers/google/marketing_platform/operators/test_display_video.py +++ b/tests/providers/google/marketing_platform/operators/test_display_video.py @@ -15,6 +15,8 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +import json +from tempfile import NamedTemporaryFile from typing import Optional from unittest import TestCase, mock @@ -57,6 +59,19 @@ class TestGoogleDisplayVideo360CreateReportOperator(TestCase): hook_mock.return_value.create_query.assert_called_once_with(query=body) xcom_mock.assert_called_once_with(None, key="report_id", value=query_id) + def test_prepare_template(self): + body = {"key": "value"} + with NamedTemporaryFile("w+", suffix=".json") as f: + f.write(json.dumps(body)) + f.flush() + op = GoogleDisplayVideo360CreateReportOperator( + body=body, api_version=API_VERSION, task_id="test_task" + ) + op.prepare_template() + + assert isinstance(op.body, dict) + assert op.body == body + class TestGoogleDisplayVideo360DeleteReportOperator(TestCase): @mock.patch( diff --git a/tests/providers/google/marketing_platform/operators/test_search_ads.py b/tests/providers/google/marketing_platform/operators/test_search_ads.py index 9e69cab96e..68cb552401 100644 --- a/tests/providers/google/marketing_platform/operators/test_search_ads.py +++ b/tests/providers/google/marketing_platform/operators/test_search_ads.py @@ -15,6 +15,8 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +import json +from tempfile import NamedTemporaryFile from unittest import TestCase, mock from airflow.providers.google.marketing_platform.operators.search_ads import ( @@ -25,7 +27,7 @@ API_VERSION = "api_version" GCP_CONN_ID = "google_cloud_default" -class TestSearchAdsGenerateReportOperator(TestCase): +class TestGoogleSearchAdsInsertReportOperator(TestCase): @mock.patch( "airflow.providers.google.marketing_platform." "operators.search_ads.GoogleSearchAdsHook" @@ -52,8 +54,21 @@ class TestSearchAdsGenerateReportOperator(TestCase): hook_mock.return_value.insert_report.assert_called_once_with(report=report) xcom_mock.assert_called_once_with(None, key="report_id", value=report_id) + def test_prepare_template(self): + report = {"key": "value"} + with NamedTemporaryFile("w+", suffix=".json") as f: + f.write(json.dumps(report)) + f.flush() + op = GoogleSearchAdsInsertReportOperator( + report=report, api_version=API_VERSION, task_id="test_task" + ) + op.prepare_template() -class TestSearchAdsGetfileReportOperator(TestCase): + assert isinstance(op.report, dict) + assert op.report == report + + +class TestGoogleSearchAdsDownloadReportOperator(TestCase): @mock.patch( "airflow.providers.google.marketing_platform." "operators.search_ads.NamedTemporaryFile"