def poke(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('Checking status of %s report.', self.report_id)
response = hook.get(report_id=self.report_id)
return response['isReportReady']
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
def execute(self, context: dict):
hook = GoogleSearchAdsHook(
gcp_conn_id=self.gcp_conn_id,
delegate_to=self.delegate_to,
api_version=self.api_version,
impersonation_chain=self.impersonation_chain,
)
gcs_hook = GCSHook(
gcp_conn_id=self.gcp_conn_id,
delegate_to=self.delegate_to,
impersonation_chain=self.impersonation_chain,
)
# 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)
def setUp(self):
with mock.patch(
"airflow.contrib.hooks.gcp_api_base_hook.GoogleCloudBaseHook.__init__",
new=mock_base_gcp_hook_default_project_id,
):
self.hook = GoogleSearchAdsHook(gcp_conn_id=GCP_CONN_ID)
class TestSearchAdsHook(TestCase):
def setUp(self):
with mock.patch(
"airflow.contrib.hooks.gcp_api_base_hook.GoogleCloudBaseHook.__init__",
new=mock_base_gcp_hook_default_project_id,
):
self.hook = GoogleSearchAdsHook(gcp_conn_id=GCP_CONN_ID)
@mock.patch("airflow.providers.google.marketing_platform.hooks."
"search_ads.GoogleSearchAdsHook._authorize")
@mock.patch(
"airflow.providers.google.marketing_platform.hooks.search_ads.build")
def test_gen_conn(self, mock_build, mock_authorize):
result = self.hook.get_conn()
mock_build.assert_called_once_with(
"doubleclicksearch",
API_VERSION,
http=mock_authorize.return_value,
cache_discovery=False,
)
self.assertEqual(mock_build.return_value, result)
@mock.patch("airflow.providers.google.marketing_platform.hooks."
"search_ads.GoogleSearchAdsHook.get_conn")
def test_insert(self, get_conn_mock):
report = {"report": "test"}
return_value = "TEST"
get_conn_mock.return_value.reports.return_value.request.return_value.execute.return_value = (
return_value)
result = self.hook.insert_report(report=report)
get_conn_mock.return_value.reports.return_value.request.assert_called_once_with(
body=report)
self.assertEqual(return_value, result)
@mock.patch("airflow.providers.google.marketing_platform.hooks."
"search_ads.GoogleSearchAdsHook.get_conn")
def test_get(self, get_conn_mock):
report_id = "REPORT_ID"
return_value = "TEST"
get_conn_mock.return_value.reports.return_value.get.return_value.execute.return_value = (
return_value)
result = self.hook.get(report_id=report_id)
get_conn_mock.return_value.reports.return_value.get.assert_called_once_with(
reportId=report_id)
self.assertEqual(return_value, result)
@mock.patch("airflow.providers.google.marketing_platform.hooks."
"search_ads.GoogleSearchAdsHook.get_conn")
def test_get_file(self, get_conn_mock):
report_fragment = 42
report_id = "REPORT_ID"
return_value = "TEST"
get_conn_mock.return_value.reports.return_value.getFile.return_value.execute.return_value = (
return_value)
result = self.hook.get_file(report_fragment=report_fragment,
report_id=report_id)
get_conn_mock.return_value.reports.return_value.getFile.assert_called_once_with(
reportFragment=report_fragment, reportId=report_id)
self.assertEqual(return_value, result)
def setUp(self):
with mock.patch(
"airflow.providers.google.marketing_platform.hooks.search_ads.GoogleBaseHook.__init__",
new=mock_base_gcp_hook_default_project_id,
):
self.hook = GoogleSearchAdsHook(gcp_conn_id=GCP_CONN_ID)