Add query and report inspection operators for Google Bid Manager (#71947)

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

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`

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}

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"]

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

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

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}

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

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",