This is an automated email from the ASF dual-hosted git repository.
shahar1 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 24f33707582 Add query and report inspection operators for Google Bid
Manager (#71947)
24f33707582 is described below
commit 24f3370758227ce1bb13e5df251d26a9271ceeee
Author: Gbenga <[email protected]>
AuthorDate: Thu Oct 29 07:44:29 2026 +0100
Add query and report inspection operators for Google Bid Manager (#71947)
---
generated/known_airflow_exceptions.txt | 2 +-
.../operators/marketing_platform/bid_manager.rst | 76 ++++++-
.../google/marketing_platform/hooks/bid_manager.py | 21 +-
.../marketing_platform/operators/bid_manager.py | 237 ++++++++++++++++++++-
.../marketing_platform/sensors/bid_manager.py | 6 +-
.../marketing_platform/example_bid_manager.py | 37 +++-
.../marketing_platform/hooks/test_bid_manager.py | 37 +++-
.../operators/test_bid_manager.py | 135 +++++++++++-
.../marketing_platform/sensors/test_bid_manager.py | 23 +-
9 files changed, 544 insertions(+), 30 deletions(-)
diff --git a/generated/known_airflow_exceptions.txt
b/generated/known_airflow_exceptions.txt
index ecf8a7634ce..1484c68ce3f 100644
--- a/generated/known_airflow_exceptions.txt
+++ b/generated/known_airflow_exceptions.txt
@@ -310,7 +310,7 @@
providers/google/src/airflow/providers/google/firebase/hooks/firestore.py::1
providers/google/src/airflow/providers/google/leveldb/hooks/leveldb.py::4
providers/google/src/airflow/providers/google/marketing_platform/hooks/campaign_manager.py::2
providers/google/src/airflow/providers/google/marketing_platform/hooks/search_ads.py::1
-providers/google/src/airflow/providers/google/marketing_platform/operators/bid_manager.py::3
+providers/google/src/airflow/providers/google/marketing_platform/operators/bid_manager.py::2
providers/google/src/airflow/providers/google/marketing_platform/sensors/display_video.py::1
providers/google/src/airflow/providers/google/suite/hooks/calendar.py::1
providers/google/src/airflow/providers/google/suite/hooks/sheets.py::1
diff --git a/providers/google/docs/operators/marketing_platform/bid_manager.rst
b/providers/google/docs/operators/marketing_platform/bid_manager.rst
index c29e2687ed5..ce45177368f 100644
--- a/providers/google/docs/operators/marketing_platform/bid_manager.rst
+++ b/providers/google/docs/operators/marketing_platform/bid_manager.rst
@@ -45,6 +45,42 @@ parameters which allow you to dynamically determine values.
You can provide body
as this operator supports this template extension.
The result is saved to :ref:`XCom <concepts:xcom>`, which allows the result to
be used by other operators.
+.. _howto/operator:GoogleBidManagerGetQueryOperator:
+
+Getting a Query
+^^^^^^^^^^^^^^^
+
+To retrieve a query using Bid Manager, use
+:class:`~airflow.providers.google.marketing_platform.operators.bid_manager.GoogleBidManagerGetQueryOperator`.
+
+.. exampleinclude::
/../../google/tests/system/google/marketing_platform/example_bid_manager.py
+ :language: python
+ :dedent: 4
+ :start-after: [START howto_google_bid_manager_get_query_operator]
+ :end-before: [END howto_google_bid_manager_get_query_operator]
+
+Use :ref:`Jinja templating <concepts:jinja-templating>` with
+:template-fields:`airflow.providers.google.marketing_platform.operators.bid_manager.GoogleBidManagerGetQueryOperator`
+parameters which allow you to dynamically determine values.
+
+.. _howto/operator:GoogleBidManagerListQueriesOperator:
+
+Listing Queries
+^^^^^^^^^^^^^^^
+
+To list queries using Bid Manager, use
+:class:`~airflow.providers.google.marketing_platform.operators.bid_manager.GoogleBidManagerListQueriesOperator`.
+
+.. exampleinclude::
/../../google/tests/system/google/marketing_platform/example_bid_manager.py
+ :language: python
+ :dedent: 4
+ :start-after: [START howto_google_bid_manager_list_queries_operator]
+ :end-before: [END howto_google_bid_manager_list_queries_operator]
+
+Use :ref:`Jinja templating <concepts:jinja-templating>` with
+:template-fields:`airflow.providers.google.marketing_platform.operators.bid_manager.GoogleBidManagerListQueriesOperator`
+parameters which allow you to dynamically determine values.
+
.. _howto/operator:GoogleBidManagerRunQueryOperator:
Run Query
@@ -64,6 +100,42 @@ You can use :ref:`Jinja templating
<concepts:jinja-templating>` with
parameters which allow you to dynamically determine values.
The result is saved to :ref:`XCom <concepts:xcom>`, which allows the result to
be used by other operators.
+.. _howto/operator:GoogleBidManagerGetReportOperator:
+
+Getting a Report
+^^^^^^^^^^^^^^^^
+
+To retrieve a report using Bid Manager, use
+:class:`~airflow.providers.google.marketing_platform.operators.bid_manager.GoogleBidManagerGetReportOperator`.
+
+.. exampleinclude::
/../../google/tests/system/google/marketing_platform/example_bid_manager.py
+ :language: python
+ :dedent: 4
+ :start-after: [START howto_google_bid_manager_get_report_operator]
+ :end-before: [END howto_google_bid_manager_get_report_operator]
+
+Use :ref:`Jinja templating <concepts:jinja-templating>` with
+:template-fields:`airflow.providers.google.marketing_platform.operators.bid_manager.GoogleBidManagerGetReportOperator`
+parameters which allow you to dynamically determine values.
+
+.. _howto/operator:GoogleBidManagerListReportsOperator:
+
+Listing Reports
+^^^^^^^^^^^^^^^
+
+To list reports for a query using Bid Manager, use
+:class:`~airflow.providers.google.marketing_platform.operators.bid_manager.GoogleBidManagerListReportsOperator`.
+
+.. exampleinclude::
/../../google/tests/system/google/marketing_platform/example_bid_manager.py
+ :language: python
+ :dedent: 4
+ :start-after: [START howto_google_bid_manager_list_reports_operator]
+ :end-before: [END howto_google_bid_manager_list_reports_operator]
+
+Use :ref:`Jinja templating <concepts:jinja-templating>` with
+:template-fields:`airflow.providers.google.marketing_platform.operators.bid_manager.GoogleBidManagerListReportsOperator`
+parameters which allow you to dynamically determine values.
+
.. _howto/operator:GoogleBidManagerDeleteQueryOperator:
Deleting a Query
@@ -111,8 +183,8 @@ To download a report to GCS bucket use
.. exampleinclude::
/../../google/tests/system/google/marketing_platform/example_bid_manager.py
:language: python
:dedent: 4
- :start-after: [START howto_google_bid_manager_get_report_operator]
- :end-before: [END howto_google_bid_manager_get_report_operator]
+ :start-after: [START howto_google_bid_manager_download_report_operator]
+ :end-before: [END howto_google_bid_manager_download_report_operator]
Use :ref:`Jinja templating <concepts:jinja-templating>` with
:template-fields:`airflow.providers.google.marketing_platform.operators.bid_manager.GoogleBidManagerDownloadReportOperator`
diff --git
a/providers/google/src/airflow/providers/google/marketing_platform/hooks/bid_manager.py
b/providers/google/src/airflow/providers/google/marketing_platform/hooks/bid_manager.py
index a3d262ab7fb..4d86e67b561 100644
---
a/providers/google/src/airflow/providers/google/marketing_platform/hooks/bid_manager.py
+++
b/providers/google/src/airflow/providers/google/marketing_platform/hooks/bid_manager.py
@@ -87,8 +87,14 @@ class GoogleBidManagerHook(GoogleBaseHook):
def list_queries(self) -> list[dict]:
"""Retrieve stored queries."""
- response =
self.get_conn().queries().list().execute(num_retries=self.num_retries)
- return response.get("queries", [])
+ conn = self.get_conn()
+ request = conn.queries().list()
+ queries: list[dict] = []
+ while request is not None:
+ response = request.execute(num_retries=self.num_retries)
+ queries.extend(response.get("queries", []))
+ request = conn.queries().list_next(previous_request=request,
previous_response=response)
+ return queries
def run_query(self, query_id: str, params: dict[str, Any] | None) -> dict:
"""
@@ -122,6 +128,11 @@ class GoogleBidManagerHook(GoogleBaseHook):
:param query_id: Query ID for which report was generated.
"""
- return (
-
self.get_conn().queries().reports().list(queryId=query_id).execute(num_retries=self.num_retries)
- )
+ conn = self.get_conn()
+ request = conn.queries().reports().list(queryId=query_id)
+ reports: list[dict] = []
+ while request is not None:
+ response = request.execute(num_retries=self.num_retries)
+ reports.extend(response.get("reports", []))
+ request =
conn.queries().reports().list_next(previous_request=request,
previous_response=response)
+ return {"reports": reports}
diff --git
a/providers/google/src/airflow/providers/google/marketing_platform/operators/bid_manager.py
b/providers/google/src/airflow/providers/google/marketing_platform/operators/bid_manager.py
index ea00d7ffc7d..71096c65c33 100644
---
a/providers/google/src/airflow/providers/google/marketing_platform/operators/bid_manager.py
+++
b/providers/google/src/airflow/providers/google/marketing_platform/operators/bid_manager.py
@@ -29,6 +29,7 @@ from typing import TYPE_CHECKING, Any
from urllib.parse import urlparse, urlsplit
from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowFailException
from airflow.providers.google.cloud.hooks.gcs import GCSHook
from airflow.providers.google.marketing_platform.hooks.bid_manager import
GoogleBidManagerHook
from airflow.providers.google.version_compat import BaseOperator
@@ -105,6 +106,114 @@ class GoogleBidManagerCreateQueryOperator(BaseOperator):
return response
+class GoogleBidManagerGetQueryOperator(BaseOperator):
+ """
+ Retrieves a stored query.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the
guide:
+ :ref:`howto/operator:GoogleBidManagerGetQueryOperator`
+
+ .. seealso::
+ Check also the official API docs:
+ `https://developers.google.com/bid-manager/v2/queries/get`
+
+ :param query_id: Query ID to retrieve.
+ :param api_version: The version of the api that will be requested for
example 'v2'.
+ :param gcp_conn_id: The connection ID to use when fetching connection info.
+ :param impersonation_chain: Optional service account to impersonate using
short-term
+ credentials, or chained list of accounts required to get the
access_token
+ of the last account in the list, which will be impersonated in the
request.
+ If set as a string, the account must grant the originating account
+ the Service Account Token Creator IAM role.
+ If set as a sequence, the identities from the list must grant
+ 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).
+ """
+
+ template_fields: Sequence[str] = (
+ "query_id",
+ "impersonation_chain",
+ "gcp_conn_id",
+ )
+
+ def __init__(
+ self,
+ *,
+ query_id: str,
+ api_version: str = "v2",
+ gcp_conn_id: str = "google_cloud_default",
+ impersonation_chain: str | Sequence[str] | None = None,
+ **kwargs,
+ ) -> None:
+ super().__init__(**kwargs)
+ self.query_id = query_id
+ self.api_version = api_version
+ self.gcp_conn_id = gcp_conn_id
+ self.impersonation_chain = impersonation_chain
+
+ def execute(self, context: Context) -> dict:
+ hook = GoogleBidManagerHook(
+ gcp_conn_id=self.gcp_conn_id,
+ api_version=self.api_version,
+ impersonation_chain=self.impersonation_chain,
+ )
+ self.log.info("Retrieving Bid Manager API query with id: %s",
self.query_id)
+ return hook.get_query(query_id=self.query_id)
+
+
+class GoogleBidManagerListQueriesOperator(BaseOperator):
+ """
+ Retrieves stored queries.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the
guide:
+ :ref:`howto/operator:GoogleBidManagerListQueriesOperator`
+
+ .. seealso::
+ Check also the official API docs:
+ `https://developers.google.com/bid-manager/v2/queries/list`
+
+ :param api_version: The version of the api that will be requested for
example 'v2'.
+ :param gcp_conn_id: The connection ID to use when fetching connection info.
+ :param impersonation_chain: Optional service account to impersonate using
short-term
+ credentials, or chained list of accounts required to get the
access_token
+ of the last account in the list, which will be impersonated in the
request.
+ If set as a string, the account must grant the originating account
+ the Service Account Token Creator IAM role.
+ If set as a sequence, the identities from the list must grant
+ 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).
+ """
+
+ template_fields: Sequence[str] = (
+ "impersonation_chain",
+ "gcp_conn_id",
+ )
+
+ def __init__(
+ self,
+ *,
+ api_version: str = "v2",
+ gcp_conn_id: str = "google_cloud_default",
+ impersonation_chain: str | Sequence[str] | None = None,
+ **kwargs,
+ ) -> None:
+ super().__init__(**kwargs)
+ self.api_version = api_version
+ self.gcp_conn_id = gcp_conn_id
+ self.impersonation_chain = impersonation_chain
+
+ def execute(self, context: Context) -> list[dict]:
+ hook = GoogleBidManagerHook(
+ gcp_conn_id=self.gcp_conn_id,
+ api_version=self.api_version,
+ impersonation_chain=self.impersonation_chain,
+ )
+ self.log.info("Retrieving list of Bid Manager API queries.")
+ return hook.list_queries()
+
+
class GoogleBidManagerRunQueryOperator(BaseOperator):
"""
Runs a stored query to generate a report.
@@ -230,6 +339,126 @@ class GoogleBidManagerDeleteQueryOperator(BaseOperator):
self.log.info("Report deleted.")
+class GoogleBidManagerGetReportOperator(BaseOperator):
+ """
+ Retrieves details of a specific report.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the
guide:
+ :ref:`howto/operator:GoogleBidManagerGetReportOperator`
+
+ .. seealso::
+ Check also the official API docs:
+ `https://developers.google.com/bid-manager/v2/queries.reports/get`
+
+ :param query_id: Query ID for which report was generated.
+ :param report_id: Report ID to retrieve.
+ :param api_version: The version of the api that will be requested for
example 'v2'.
+ :param gcp_conn_id: The connection ID to use when fetching connection info.
+ :param impersonation_chain: Optional service account to impersonate using
short-term
+ credentials, or chained list of accounts required to get the
access_token
+ of the last account in the list, which will be impersonated in the
request.
+ If set as a string, the account must grant the originating account
+ the Service Account Token Creator IAM role.
+ If set as a sequence, the identities from the list must grant
+ 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).
+ """
+
+ template_fields: Sequence[str] = (
+ "query_id",
+ "report_id",
+ "impersonation_chain",
+ "gcp_conn_id",
+ )
+
+ def __init__(
+ self,
+ *,
+ query_id: str,
+ report_id: str,
+ api_version: str = "v2",
+ gcp_conn_id: str = "google_cloud_default",
+ impersonation_chain: str | Sequence[str] | None = None,
+ **kwargs,
+ ) -> None:
+ super().__init__(**kwargs)
+ self.query_id = query_id
+ self.report_id = report_id
+ self.api_version = api_version
+ self.gcp_conn_id = gcp_conn_id
+ self.impersonation_chain = impersonation_chain
+
+ def execute(self, context: Context) -> dict:
+ hook = GoogleBidManagerHook(
+ gcp_conn_id=self.gcp_conn_id,
+ api_version=self.api_version,
+ impersonation_chain=self.impersonation_chain,
+ )
+ self.log.info(
+ "Retrieving Bid Manager API report with query_id: %s and
report_id: %s",
+ self.query_id,
+ self.report_id,
+ )
+ return hook.get_report(query_id=self.query_id,
report_id=self.report_id)
+
+
+class GoogleBidManagerListReportsOperator(BaseOperator):
+ """
+ Retrieves a list of reports for a specific query.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the
guide:
+ :ref:`howto/operator:GoogleBidManagerListReportsOperator`
+
+ .. seealso::
+ Check also the official API docs:
+ `https://developers.google.com/bid-manager/v2/queries.reports/list`
+
+ :param query_id: Query ID for which reports were generated.
+ :param api_version: The version of the api that will be requested for
example 'v2'.
+ :param gcp_conn_id: The connection ID to use when fetching connection info.
+ :param impersonation_chain: Optional service account to impersonate using
short-term
+ credentials, or chained list of accounts required to get the
access_token
+ of the last account in the list, which will be impersonated in the
request.
+ If set as a string, the account must grant the originating account
+ the Service Account Token Creator IAM role.
+ If set as a sequence, the identities from the list must grant
+ 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).
+ """
+
+ template_fields: Sequence[str] = (
+ "query_id",
+ "impersonation_chain",
+ "gcp_conn_id",
+ )
+
+ def __init__(
+ self,
+ *,
+ query_id: str,
+ api_version: str = "v2",
+ gcp_conn_id: str = "google_cloud_default",
+ impersonation_chain: str | Sequence[str] | None = None,
+ **kwargs,
+ ) -> None:
+ super().__init__(**kwargs)
+ self.query_id = query_id
+ self.api_version = api_version
+ self.gcp_conn_id = gcp_conn_id
+ self.impersonation_chain = impersonation_chain
+
+ def execute(self, context: Context) -> dict:
+ hook = GoogleBidManagerHook(
+ gcp_conn_id=self.gcp_conn_id,
+ api_version=self.api_version,
+ impersonation_chain=self.impersonation_chain,
+ )
+ self.log.info("Retrieving list of Bid Manager API reports for
query_id: %s", self.query_id)
+ return hook.list_reports(query_id=self.query_id)
+
+
class GoogleBidManagerDownloadReportOperator(BaseOperator):
"""
Retrieves a stored query.
@@ -317,8 +546,12 @@ class GoogleBidManagerDownloadReportOperator(BaseOperator):
resource = hook.get_report(query_id=self.query_id,
report_id=self.report_id)
status = resource.get("metadata", {}).get("status", {}).get("state")
- if resource and status not in ["DONE", "FAILED"]:
- raise AirflowException(f"Report {self.report_id} for query
{self.query_id} is still running")
+ if status == "FAILED":
+ raise AirflowFailException(f"Report {self.report_id} for query
{self.query_id} failed execution.")
+ if status != "DONE":
+ raise ValueError(
+ f"Report {self.report_id} for query {self.query_id} is in
state '{status}'. Expected 'DONE'."
+ )
# If no custom report_name provided, use Bid Manager name
file_url = resource["metadata"]["googleCloudStoragePath"]
diff --git
a/providers/google/src/airflow/providers/google/marketing_platform/sensors/bid_manager.py
b/providers/google/src/airflow/providers/google/marketing_platform/sensors/bid_manager.py
index ca4c10ab2c4..d4d016cd550 100644
---
a/providers/google/src/airflow/providers/google/marketing_platform/sensors/bid_manager.py
+++
b/providers/google/src/airflow/providers/google/marketing_platform/sensors/bid_manager.py
@@ -21,7 +21,7 @@ from __future__ import annotations
from collections.abc import Sequence
from typing import TYPE_CHECKING
-from airflow.providers.common.compat.sdk import BaseSensorOperator
+from airflow.providers.common.compat.sdk import AirflowFailException,
BaseSensorOperator
from airflow.providers.google.marketing_platform.hooks.bid_manager import
GoogleBidManagerHook
if TYPE_CHECKING:
@@ -84,6 +84,8 @@ class GoogleBidManagerRunQuerySensor(BaseSensorOperator):
response = hook.get_report(query_id=self.query_id,
report_id=self.report_id)
status = response.get("metadata", {}).get("status", {}).get("state")
self.log.info("STATUS OF THE REPORT %s FOR QUERY %s: %s",
self.report_id, self.query_id, status)
- if response and status in ["DONE", "FAILED"]:
+ if status == "FAILED":
+ raise AirflowFailException(f"Report {self.report_id} for query
{self.query_id} failed execution.")
+ if status == "DONE":
return True
return False
diff --git
a/providers/google/tests/system/google/marketing_platform/example_bid_manager.py
b/providers/google/tests/system/google/marketing_platform/example_bid_manager.py
index 2c08a035581..07bc31f0516 100644
---
a/providers/google/tests/system/google/marketing_platform/example_bid_manager.py
+++
b/providers/google/tests/system/google/marketing_platform/example_bid_manager.py
@@ -33,6 +33,10 @@ from
airflow.providers.google.marketing_platform.operators.bid_manager import (
GoogleBidManagerCreateQueryOperator,
GoogleBidManagerDeleteQueryOperator,
GoogleBidManagerDownloadReportOperator,
+ GoogleBidManagerGetQueryOperator,
+ GoogleBidManagerGetReportOperator,
+ GoogleBidManagerListQueriesOperator,
+ GoogleBidManagerListReportsOperator,
GoogleBidManagerRunQueryOperator,
)
from airflow.providers.google.marketing_platform.sensors.bid_manager import (
@@ -147,6 +151,14 @@ with DAG(
query_id = cast("str", XComArg(create_query, key="query_id"))
# [END howto_google_bid_manager_create_query_operator]
+ # [START howto_google_bid_manager_get_query_operator]
+ get_query = GoogleBidManagerGetQueryOperator(query_id=query_id,
task_id="get_query", gcp_conn_id=CONN_ID)
+ # [END howto_google_bid_manager_get_query_operator]
+
+ # [START howto_google_bid_manager_list_queries_operator]
+ list_queries = GoogleBidManagerListQueriesOperator(task_id="list_queries",
gcp_conn_id=CONN_ID)
+ # [END howto_google_bid_manager_list_queries_operator]
+
# [START howto_google_bid_manager_run_query_report_operator]
run_query = GoogleBidManagerRunQueryOperator(
query_id=query_id, parameters=PARAMETERS, task_id="run_report",
gcp_conn_id=CONN_ID
@@ -166,14 +178,32 @@ with DAG(
# [END howto_google_bid_manager_wait_run_query_sensor]
# [START howto_google_bid_manager_get_report_operator]
- get_report = GoogleBidManagerDownloadReportOperator(
+ get_report = GoogleBidManagerGetReportOperator(
query_id=query_id,
report_id=report_id,
task_id="get_report",
+ gcp_conn_id=CONN_ID,
+ )
+ # [END howto_google_bid_manager_get_report_operator]
+
+ # [START howto_google_bid_manager_list_reports_operator]
+ list_reports = GoogleBidManagerListReportsOperator(
+ query_id=query_id,
+ task_id="list_reports",
+ gcp_conn_id=CONN_ID,
+ )
+ # [END howto_google_bid_manager_list_reports_operator]
+
+ # [START howto_google_bid_manager_download_report_operator]
+ download_report = GoogleBidManagerDownloadReportOperator(
+ query_id=query_id,
+ report_id=report_id,
+ task_id="download_report",
bucket_name=BUCKET_NAME,
report_name="test1.csv",
gcp_conn_id=CONN_ID,
)
+ # [END howto_google_bid_manager_download_report_operator]
delete_bucket = GCSDeleteBucketOperator(
task_id="delete_bucket",
@@ -181,7 +211,6 @@ with DAG(
gcp_conn_id=CONN_ID,
trigger_rule=TriggerRule.ALL_DONE,
)
- # [END howto_google_bid_manager_get_report_operator]
# [START howto_google_bid_manager_delete_query_operator]
delete_query = GoogleBidManagerDeleteQueryOperator(
@@ -194,9 +223,13 @@ with DAG(
>> create_connection_display_video_task # type: ignore
>> create_bucket
>> create_query
+ >> get_query
+ >> list_queries
>> run_query
>> wait_for_query
>> get_report
+ >> list_reports
+ >> download_report
>> delete_query
>> delete_bucket
>> delete_connection_task
diff --git
a/providers/google/tests/unit/google/marketing_platform/hooks/test_bid_manager.py
b/providers/google/tests/unit/google/marketing_platform/hooks/test_bid_manager.py
index 5f460e0ce7e..38fc3ae94e3 100644
---
a/providers/google/tests/unit/google/marketing_platform/hooks/test_bid_manager.py
+++
b/providers/google/tests/unit/google/marketing_platform/hooks/test_bid_manager.py
@@ -19,6 +19,8 @@ from __future__ import annotations
from unittest import mock
+from googleapiclient.http import HttpRequest
+
from airflow.providers.google.marketing_platform.hooks.bid_manager import
GoogleBidManagerHook
from unit.google.cloud.utils.base_gcp_mock import
mock_base_gcp_hook_default_project_id
@@ -93,14 +95,20 @@ class TestGoogleBidManagerHook:
@mock.patch("airflow.providers.google.marketing_platform.hooks.bid_manager.GoogleBidManagerHook.get_conn")
def test_list_queries(self, get_conn_mock):
- queries = ["test"]
- return_value = {"queries": queries}
-
get_conn_mock.return_value.queries.return_value.list.return_value.execute.return_value
= return_value
+ queries_page1 = [{"queryId": "q1"}]
+ queries_page2 = [{"queryId": "q2"}]
+ mock_request = mock.MagicMock(spec=HttpRequest)
+ mock_next_request = mock.MagicMock(spec=HttpRequest)
+ get_conn_mock.return_value.queries.return_value.list.return_value =
mock_request
+ mock_request.execute.return_value = {"queries": queries_page1}
+ mock_next_request.execute.return_value = {"queries": queries_page2}
+ get_conn_mock.return_value.queries.return_value.list_next.side_effect
= [mock_next_request, None]
+
result = self.hook.list_queries()
get_conn_mock.return_value.queries.return_value.list.assert_called_once_with()
-
- assert queries == result
+ assert
get_conn_mock.return_value.queries.return_value.list_next.call_count == 2
+ assert result == queries_page1 + queries_page2
@mock.patch("airflow.providers.google.marketing_platform.hooks.bid_manager.GoogleBidManagerHook.get_conn")
def test_run_query(self, get_conn_mock):
@@ -135,15 +143,22 @@ class TestGoogleBidManagerHook:
@mock.patch("airflow.providers.google.marketing_platform.hooks.bid_manager.GoogleBidManagerHook.get_conn")
def test_list_reports(self, get_conn_mock):
query_id = "QUERY_ID"
- reports = ["report1", "report2"]
- return_value = {"reports": reports}
- (
-
get_conn_mock.return_value.queries.return_value.reports.return_value.list.return_value.execute.return_value
- ) = return_value
+ reports_page1 = [{"key": {"reportId": "r1"}}]
+ reports_page2 = [{"key": {"reportId": "r2"}}]
+ mock_request = mock.MagicMock(spec=HttpRequest)
+ mock_next_request = mock.MagicMock(spec=HttpRequest)
+
get_conn_mock.return_value.queries.return_value.reports.return_value.list.return_value
= mock_request
+ mock_request.execute.return_value = {"reports": reports_page1}
+ mock_next_request.execute.return_value = {"reports": reports_page2}
+
get_conn_mock.return_value.queries.return_value.reports.return_value.list_next.side_effect
= [
+ mock_next_request,
+ None,
+ ]
result = self.hook.list_reports(query_id=query_id)
get_conn_mock.return_value.queries.return_value.reports.return_value.list.assert_called_once_with(
queryId=query_id
)
- assert reports == result["reports"]
+ assert
get_conn_mock.return_value.queries.return_value.reports.return_value.list_next.call_count
== 2
+ assert result == {"reports": reports_page1 + reports_page2}
diff --git
a/providers/google/tests/unit/google/marketing_platform/operators/test_bid_manager.py
b/providers/google/tests/unit/google/marketing_platform/operators/test_bid_manager.py
index 1f7b2928f7d..de9ac36c7ca 100644
---
a/providers/google/tests/unit/google/marketing_platform/operators/test_bid_manager.py
+++
b/providers/google/tests/unit/google/marketing_platform/operators/test_bid_manager.py
@@ -26,11 +26,15 @@ from sqlalchemy import delete
from airflow.exceptions import AirflowException
from airflow.models import TaskInstance as TI
-from airflow.providers.common.compat.sdk import timezone
+from airflow.providers.common.compat.sdk import AirflowFailException, timezone
from airflow.providers.google.marketing_platform.operators.bid_manager import (
GoogleBidManagerCreateQueryOperator,
GoogleBidManagerDeleteQueryOperator,
GoogleBidManagerDownloadReportOperator,
+ GoogleBidManagerGetQueryOperator,
+ GoogleBidManagerGetReportOperator,
+ GoogleBidManagerListQueriesOperator,
+ GoogleBidManagerListReportsOperator,
GoogleBidManagerRunQueryOperator,
)
from airflow.utils.session import create_session
@@ -426,6 +430,31 @@ class TestGoogleBidManagerDownloadReportOperator:
object_name=REPORT_NAME + ".gz",
)
+ @pytest.mark.parametrize(
+ ("state", "expected_exception", "match"),
+ [
+ ("FAILED", AirflowFailException, "failed execution"),
+ ("RUNNING", ValueError, "is in state 'RUNNING'. Expected 'DONE'"),
+ ],
+ )
+
@mock.patch("airflow.providers.google.marketing_platform.operators.bid_manager.GoogleBidManagerHook")
+ def test_execute_non_done_status_raises(self, mock_hook, state,
expected_exception, match):
+ mock_hook.return_value.get_report.return_value = {
+ "metadata": {
+ "status": {"state": state},
+ }
+ }
+ mock_context = {"task_instance": mock.Mock()}
+ op = GoogleBidManagerDownloadReportOperator(
+ query_id=QUERY_ID,
+ report_id=REPORT_ID,
+ bucket_name=BUCKET_NAME,
+ report_name=REPORT_NAME,
+ task_id="test_task",
+ )
+ with pytest.raises(expected_exception, match=match):
+ op.execute(context=mock_context)
+
class TestGoogleBidManagerRunQueryOperator:
@mock.patch("airflow.providers.google.marketing_platform.operators.bid_manager.GoogleBidManagerHook")
@@ -483,8 +512,110 @@ class TestGoogleBidManagerCreateQueryOperator:
with NamedTemporaryFile("w+", suffix=".json") as f:
f.write(json.dumps(body))
f.flush()
- op = GoogleBidManagerCreateQueryOperator(body=body,
task_id="test_task")
+ op = GoogleBidManagerCreateQueryOperator(
+ body=f.name,
+ api_version=API_VERSION,
+ gcp_conn_id=GCP_CONN_ID,
+ impersonation_chain=IMPERSONATION_CHAIN,
+ task_id="test_task",
+ )
op.prepare_template()
assert isinstance(op.body, dict)
assert op.body == body
+
+
+class TestGoogleBidManagerGetQueryOperator:
+
@mock.patch("airflow.providers.google.marketing_platform.operators.bid_manager.GoogleBidManagerHook")
+ def test_execute(self, hook_mock):
+ return_value = {"queryId": QUERY_ID}
+ hook_mock.return_value.get_query.return_value = return_value
+
+ op = GoogleBidManagerGetQueryOperator(
+ query_id=QUERY_ID,
+ api_version=API_VERSION,
+ gcp_conn_id=GCP_CONN_ID,
+ impersonation_chain=IMPERSONATION_CHAIN,
+ task_id="test_task",
+ )
+ result = op.execute(context=None)
+
+ hook_mock.assert_called_once_with(
+ gcp_conn_id=GCP_CONN_ID,
+ api_version=API_VERSION,
+ impersonation_chain=IMPERSONATION_CHAIN,
+ )
+
hook_mock.return_value.get_query.assert_called_once_with(query_id=QUERY_ID)
+ assert result == return_value
+
+
+class TestGoogleBidManagerListQueriesOperator:
+
@mock.patch("airflow.providers.google.marketing_platform.operators.bid_manager.GoogleBidManagerHook")
+ def test_execute(self, hook_mock):
+ return_value = [{"queryId": QUERY_ID}]
+ hook_mock.return_value.list_queries.return_value = return_value
+
+ op = GoogleBidManagerListQueriesOperator(
+ api_version=API_VERSION,
+ gcp_conn_id=GCP_CONN_ID,
+ impersonation_chain=IMPERSONATION_CHAIN,
+ task_id="test_task",
+ )
+ result = op.execute(context=None)
+
+ hook_mock.assert_called_once_with(
+ gcp_conn_id=GCP_CONN_ID,
+ api_version=API_VERSION,
+ impersonation_chain=IMPERSONATION_CHAIN,
+ )
+ hook_mock.return_value.list_queries.assert_called_once_with()
+ assert result == return_value
+
+
+class TestGoogleBidManagerGetReportOperator:
+
@mock.patch("airflow.providers.google.marketing_platform.operators.bid_manager.GoogleBidManagerHook")
+ def test_execute(self, hook_mock):
+ return_value = {"key": {"queryId": QUERY_ID, "reportId": REPORT_ID}}
+ hook_mock.return_value.get_report.return_value = return_value
+
+ op = GoogleBidManagerGetReportOperator(
+ query_id=QUERY_ID,
+ report_id=REPORT_ID,
+ api_version=API_VERSION,
+ gcp_conn_id=GCP_CONN_ID,
+ impersonation_chain=IMPERSONATION_CHAIN,
+ task_id="test_task",
+ )
+ result = op.execute(context=None)
+
+ hook_mock.assert_called_once_with(
+ gcp_conn_id=GCP_CONN_ID,
+ api_version=API_VERSION,
+ impersonation_chain=IMPERSONATION_CHAIN,
+ )
+
hook_mock.return_value.get_report.assert_called_once_with(query_id=QUERY_ID,
report_id=REPORT_ID)
+ assert result == return_value
+
+
+class TestGoogleBidManagerListReportsOperator:
+
@mock.patch("airflow.providers.google.marketing_platform.operators.bid_manager.GoogleBidManagerHook")
+ def test_execute(self, hook_mock):
+ return_value = {"reports": [{"key": {"queryId": QUERY_ID, "reportId":
REPORT_ID}}]}
+ hook_mock.return_value.list_reports.return_value = return_value
+
+ op = GoogleBidManagerListReportsOperator(
+ query_id=QUERY_ID,
+ api_version=API_VERSION,
+ gcp_conn_id=GCP_CONN_ID,
+ impersonation_chain=IMPERSONATION_CHAIN,
+ task_id="test_task",
+ )
+ result = op.execute(context=None)
+
+ hook_mock.assert_called_once_with(
+ gcp_conn_id=GCP_CONN_ID,
+ api_version=API_VERSION,
+ impersonation_chain=IMPERSONATION_CHAIN,
+ )
+
hook_mock.return_value.list_reports.assert_called_once_with(query_id=QUERY_ID)
+ assert result == return_value
diff --git
a/providers/google/tests/unit/google/marketing_platform/sensors/test_bid_manager.py
b/providers/google/tests/unit/google/marketing_platform/sensors/test_bid_manager.py
index f48810e60de..aac1c34901e 100644
---
a/providers/google/tests/unit/google/marketing_platform/sensors/test_bid_manager.py
+++
b/providers/google/tests/unit/google/marketing_platform/sensors/test_bid_manager.py
@@ -19,6 +19,9 @@ from __future__ import annotations
from unittest import mock
+import pytest
+
+from airflow.providers.common.compat.sdk import AirflowFailException
from airflow.providers.google.marketing_platform.sensors.bid_manager import (
GoogleBidManagerRunQuerySensor,
)
@@ -29,13 +32,27 @@ GCP_CONN_ID = "google_cloud_default"
class TestGoogleBidManagerRunQuerySensor:
+ @pytest.mark.parametrize(
+ ("state", "expected_result", "expected_exception"),
+ [
+ ("DONE", True, None),
+ ("RUNNING", False, None),
+ ("FAILED", None, AirflowFailException),
+ ],
+ )
@mock.patch(f"{MODULE_NAME}.GoogleBidManagerHook")
- @mock.patch(f"{MODULE_NAME}.BaseSensorOperator")
- def test_poke(self, mock_base_op, hook_mock):
+ def test_poke(self, hook_mock, state, expected_result, expected_exception):
query_id = "QUERY_ID"
report_id = "REPORT_ID"
+ hook_mock.return_value.get_report.return_value = {"metadata":
{"status": {"state": state}}}
op = GoogleBidManagerRunQuerySensor(query_id=query_id,
report_id=report_id, task_id="test_task")
- op.poke(context=None)
+ if expected_exception:
+ with pytest.raises(expected_exception, match="failed execution"):
+ op.poke(context=None)
+ else:
+ result = op.poke(context=None)
+ assert result is expected_result
+
hook_mock.assert_called_once_with(
gcp_conn_id=GCP_CONN_ID,
api_version="v2",
