apache-airflow-providers-google 8.12.0rc1__py3-none-any.whl → 9.0.0rc2__py3-none-any.whl

airflow/providers/google/cloud/transfers/gcs_to_gcs.py

@@ -150,7 +150,7 @@ class GCSToGCSOperator(BaseOperator):
         )
 
     The following Operator would move all the Avro files from ``sales/sales-2019``
-     and ``sales/sales-2020` folder in ``data`` bucket to the same folder in the
+     and ``sales/sales-2020`` folder in ``data`` bucket to the same folder in the
      ``data_backup`` bucket, deleting the original files in the process. ::
 
         move_files = GCSToGCSOperator(

airflow/providers/google/cloud/triggers/kubernetes_engine.py

@@ -18,14 +18,108 @@
 from __future__ import annotations
 
 import asyncio
+from datetime import datetime
 from typing import Any, AsyncIterator, Sequence
 
 from google.cloud.container_v1.types import Operation
 
-from airflow.providers.google.cloud.hooks.kubernetes_engine import AsyncGKEHook
+try:
+    from airflow.providers.cncf.kubernetes.triggers.pod import KubernetesPodTrigger
+except ImportError:
+    # preserve backward compatibility for older versions of cncf.kubernetes provider
+    from airflow.providers.cncf.kubernetes.triggers.kubernetes_pod import KubernetesPodTrigger
+from airflow.providers.google.cloud.hooks.kubernetes_engine import GKEAsyncHook, GKEPodAsyncHook
 from airflow.triggers.base import BaseTrigger, TriggerEvent
 
 
+class GKEStartPodTrigger(KubernetesPodTrigger):
+    """
+    Trigger for checking pod status until it finishes its job.
+
+    :param pod_name: The name of the pod.
+    :param pod_namespace: The namespace of the pod.
+    :param cluster_url: The URL pointed to the cluster.
+    :param ssl_ca_cert: SSL certificate that is used for authentication to the pod.
+    :param cluster_context: Context that points to kubernetes cluster.
+    :param poll_interval: Polling period in seconds to check for the status.
+    :param trigger_start_time: time in Datetime format when the trigger was started
+    :param in_cluster: run kubernetes client with in_cluster configuration.
+    :param should_delete_pod: What to do when the pod reaches its final
+        state, or the execution is interrupted. If True (default), delete the
+        pod; if False, leave the pod.
+    :param get_logs: get the stdout of the container as logs of the tasks.
+    :param startup_timeout: timeout in seconds to start up the pod.
+    :param base_container_name: The name of the base container in the pod. This container's logs
+        will appear as part of this task's logs if get_logs is True. Defaults to None. If None,
+        will consult the class variable BASE_CONTAINER_NAME (which defaults to "base") for the base
+        container name to use.
+    """
+
+    def __init__(
+        self,
+        pod_name: str,
+        pod_namespace: str,
+        cluster_url: str,
+        ssl_ca_cert: str,
+        base_container_name: str,
+        trigger_start_time: datetime,
+        cluster_context: str | None = None,
+        poll_interval: float = 2,
+        in_cluster: bool | None = None,
+        should_delete_pod: bool = True,
+        get_logs: bool = True,
+        startup_timeout: int = 120,
+        *args,
+        **kwargs,
+    ):
+        super().__init__(
+            pod_name,
+            pod_namespace,
+            trigger_start_time,
+            base_container_name,
+            *args,
+            **kwargs,
+        )
+        self.pod_name = pod_name
+        self.pod_namespace = pod_namespace
+        self.trigger_start_time = trigger_start_time
+        self.base_container_name = base_container_name
+        self.poll_interval = poll_interval
+        self.cluster_context = cluster_context
+        self.in_cluster = in_cluster
+        self.should_delete_pod = should_delete_pod
+        self.get_logs = get_logs
+        self.startup_timeout = startup_timeout
+
+        self._cluster_url = cluster_url
+        self._ssl_ca_cert = ssl_ca_cert
+
+    def serialize(self) -> tuple[str, dict[str, Any]]:
+        return (
+            "airflow.providers.google.cloud.triggers.kubernetes_engine.GKEStartPodTrigger",
+            {
+                "pod_name": self.pod_name,
+                "pod_namespace": self.pod_namespace,
+                "cluster_url": self._cluster_url,
+                "ssl_ca_cert": self._ssl_ca_cert,
+                "poll_interval": self.poll_interval,
+                "cluster_context": self.cluster_context,
+                "in_cluster": self.in_cluster,
+                "should_delete_pod": self.should_delete_pod,
+                "get_logs": self.get_logs,
+                "startup_timeout": self.startup_timeout,
+                "trigger_start_time": self.trigger_start_time,
+                "base_container_name": self.base_container_name,
+            },
+        )
+
+    def _get_async_hook(self) -> GKEPodAsyncHook:  # type: ignore[override]
+        return GKEPodAsyncHook(
+            cluster_url=self._cluster_url,
+            ssl_ca_cert=self._ssl_ca_cert,
+        )
+
+
 class GKEOperationTrigger(BaseTrigger):
     """Trigger which checks status of the operation."""
 
@@ -49,7 +143,7 @@ class GKEOperationTrigger(BaseTrigger):
         self.impersonation_chain = impersonation_chain
         self.poll_interval = poll_interval
 
-        self._hook: AsyncGKEHook | None = None
+        self._hook: GKEAsyncHook | None = None
 
     def serialize(self) -> tuple[str, dict[str, Any]]:
         """Serializes GKEOperationTrigger arguments and classpath."""
@@ -110,9 +204,9 @@ class GKEOperationTrigger(BaseTrigger):
                 )
                 return
 
-    def _get_hook(self) -> AsyncGKEHook:
+    def _get_hook(self) -> GKEAsyncHook:
         if self._hook is None:
-            self._hook = AsyncGKEHook(
+            self._hook = GKEAsyncHook(
                 gcp_conn_id=self.gcp_conn_id,
                 location=self.location,
                 delegate_to=self.delegate_to,

airflow/providers/google/get_provider_info.py

@@ -27,7 +27,9 @@ def get_provider_info():
         "package-name": "apache-airflow-providers-google",
         "name": "Google",
         "description": "Google services including:\n\n  - `Google Ads <https://ads.google.com/>`__\n  - `Google Cloud (GCP) <https://cloud.google.com/>`__\n  - `Google Firebase <https://firebase.google.com/>`__\n  - `Google LevelDB <https://github.com/google/leveldb/>`__\n  - `Google Marketing Platform <https://marketingplatform.google.com/>`__\n  - `Google Workspace <https://workspace.google.com/>`__ (formerly Google Suite)\n",
+        "suspended": False,
         "versions": [
+            "9.0.0",
             "8.12.0",
             "8.11.0",
             "8.10.0",

airflow/providers/google/marketing_platform/example_dags/example_display_video.py

@@ -29,11 +29,14 @@ from airflow.models.xcom_arg import XComArg
 from airflow.providers.google.cloud.transfers.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 (
+    GoogleDisplayVideo360CreateQueryOperator,
     GoogleDisplayVideo360CreateReportOperator,
     GoogleDisplayVideo360CreateSDFDownloadTaskOperator,
     GoogleDisplayVideo360DeleteReportOperator,
     GoogleDisplayVideo360DownloadLineItemsOperator,
     GoogleDisplayVideo360DownloadReportOperator,
+    GoogleDisplayVideo360DownloadReportV2Operator,
+    GoogleDisplayVideo360RunQueryOperator,
     GoogleDisplayVideo360RunReportOperator,
     GoogleDisplayVideo360SDFtoGCSOperator,
     GoogleDisplayVideo360UploadLineItemsOperator,
@@ -41,6 +44,7 @@ from airflow.providers.google.marketing_platform.operators.display_video import
 from airflow.providers.google.marketing_platform.sensors.display_video import (
     GoogleDisplayVideo360GetSDFDownloadOperationSensor,
     GoogleDisplayVideo360ReportSensor,
+    GoogleDisplayVideo360RunQuerySensor,
 )
 
 # [START howto_display_video_env_variables]
@@ -50,7 +54,7 @@ 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")
 BUCKET_FILE_LOCATION = PATH_TO_UPLOAD_FILE.rpartition("/")[-1]
-SDF_VERSION = os.environ.get("GMP_SDF_VERSION", "SDF_VERSION_5_1")
+SDF_VERSION = os.environ.get("GMP_SDF_VERSION", "SDF_VERSION_5_5")
 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")
@@ -74,7 +78,25 @@ REPORT = {
     "schedule": {"frequency": "ONE_TIME"},
 }
 
-PARAMETERS = {"dataRange": "LAST_14_DAYS", "timezoneCode": "America/New_York"}
+REPORT_V2 = {
+    "metadata": {
+        "title": "Airflow Test Report",
+        "dataRange": {"range": "LAST_7_DAYS"},
+        "format": "CSV",
+        "sendNotification": False,
+    },
+    "params": {
+        "type": "STANDARD",
+        "groupBys": ["FILTER_DATE", "FILTER_PARTNER"],
+        "filters": [{"type": "FILTER_PARTNER", "value": ADVERTISER_ID}],
+        "metrics": ["METRIC_IMPRESSIONS", "METRIC_CLICKS"],
+    },
+    "schedule": {"frequency": "ONE_TIME"},
+}
+
+PARAMETERS = {
+    "dataRange": {"range": "LAST_7_DAYS"},
+}
 
 CREATE_SDF_DOWNLOAD_TASK_BODY_REQUEST: dict = {
     "version": SDF_VERSION,
@@ -209,3 +231,46 @@ with models.DAG(
 
     # Task dependency created via `XComArgs`:
     #   save_sdf_in_gcs >> upload_sdf_to_big_query
+
+with models.DAG(
+    "example_display_video_v2",
+    start_date=START_DATE,
+    catchup=False,
+) as dag:
+    # [START howto_google_display_video_create_query_operator]
+    create_query_v2 = GoogleDisplayVideo360CreateQueryOperator(body=REPORT_V2, task_id="create_query")
+
+    query_id = cast(str, XComArg(create_query_v2, key="query_id"))
+    # [END howto_google_display_video_create_query_operator]
+
+    # [START howto_google_display_video_run_query_report_operator]
+    run_query_v2 = GoogleDisplayVideo360RunQueryOperator(
+        query_id=query_id, parameters=PARAMETERS, task_id="run_report"
+    )
+
+    query_id = cast(str, XComArg(run_query_v2, key="query_id"))
+    report_id = cast(str, XComArg(run_query_v2, key="report_id"))
+    # [END howto_google_display_video_run_query_report_operator]
+
+    # [START howto_google_display_video_wait_run_query_sensor]
+    wait_for_query = GoogleDisplayVideo360RunQuerySensor(
+        task_id="wait_for_query",
+        query_id=query_id,
+        report_id=report_id,
+    )
+    # [END howto_google_display_video_wait_run_query_sensor]
+
+    # [START howto_google_display_video_get_report_operator]
+    get_report_v2 = GoogleDisplayVideo360DownloadReportV2Operator(
+        query_id=query_id,
+        report_id=report_id,
+        task_id="get_report",
+        bucket_name=BUCKET,
+        report_name="test1.csv",
+    )
+    # # [END howto_google_display_video_get_report_operator]
+    # # [START howto_google_display_video_delete_query_report_operator]
+    delete_report_v2 = GoogleDisplayVideo360DeleteReportOperator(report_id=report_id, task_id="delete_report")
+    # # [END howto_google_display_video_delete_query_report_operator]
+
+    create_query_v2 >> run_query_v2 >> wait_for_query >> get_report_v2 >> delete_report_v2

airflow/providers/google/marketing_platform/hooks/display_video.py

@@ -18,6 +18,7 @@
 """This module contains Google DisplayVideo hook."""
 from __future__ import annotations
 
+import warnings
 from typing import Any, Sequence
 
 from googleapiclient.discovery import Resource, build
@@ -32,7 +33,7 @@ class GoogleDisplayVideo360Hook(GoogleBaseHook):
 
     def __init__(
         self,
-        api_version: str = "v1",
+        api_version: str = "v2",
         gcp_conn_id: str = "google_cloud_default",
         delegate_to: str | None = None,
         impersonation_chain: str | Sequence[str] | None = None,
@@ -42,6 +43,11 @@ class GoogleDisplayVideo360Hook(GoogleBaseHook):
             delegate_to=delegate_to,
             impersonation_chain=impersonation_chain,
         )
+        if api_version in ["v1", "v1.1"]:
+            warnings.warn(
+                f"API {api_version} is deprecated and shortly will be removed please use v2",
+                DeprecationWarning,
+            )
         self.api_version = api_version
 
     def get_conn(self) -> Resource:
@@ -93,7 +99,10 @@ class GoogleDisplayVideo360Hook(GoogleBaseHook):
 
         :param query: Query object to be passed to request body.
         """
-        response = self.get_conn().queries().createquery(body=query).execute(num_retries=self.num_retries)
+        if self.api_version in ["v1", "v1.1"]:
+            response = self.get_conn().queries().createquery(body=query).execute(num_retries=self.num_retries)
+        else:
+            response = self.get_conn().queries().create(body=query).execute(num_retries=self.num_retries)
         return response
 
     def delete_query(self, query_id: str) -> None:
@@ -102,7 +111,10 @@ class GoogleDisplayVideo360Hook(GoogleBaseHook):
 
         :param query_id: Query ID to delete.
         """
-        (self.get_conn().queries().deletequery(queryId=query_id).execute(num_retries=self.num_retries))
+        if self.api_version in ["v1", "v1.1"]:
+            self.get_conn().queries().deletequery(queryId=query_id).execute(num_retries=self.num_retries)
+        else:
+            self.get_conn().queries().delete(queryId=query_id).execute(num_retries=self.num_retries)
 
     def get_query(self, query_id: str) -> dict:
         """
@@ -110,25 +122,56 @@ class GoogleDisplayVideo360Hook(GoogleBaseHook):
 
         :param query_id: Query ID to retrieve.
         """
-        response = self.get_conn().queries().getquery(queryId=query_id).execute(num_retries=self.num_retries)
+        if self.api_version in ["v1", "v1.1"]:
+            response = (
+                self.get_conn().queries().getquery(queryId=query_id).execute(num_retries=self.num_retries)
+            )
+        else:
+            response = self.get_conn().queries().get(queryId=query_id).execute(num_retries=self.num_retries)
         return response
 
     def list_queries(self) -> list[dict]:
         """Retrieves stored queries."""
-        response = self.get_conn().queries().listqueries().execute(num_retries=self.num_retries)
+        if self.api_version in ["v1", "v1.1"]:
+            response = self.get_conn().queries().listqueries().execute(num_retries=self.num_retries)
+        else:
+            response = self.get_conn().queries().list().execute(num_retries=self.num_retries)
         return response.get("queries", [])
 
-    def run_query(self, query_id: str, params: dict[str, Any] | None) -> None:
+    def run_query(self, query_id: str, params: dict[str, Any] | None) -> dict:
         """
         Runs a stored query to generate a report.
 
         :param query_id: Query ID to run.
         :param params: Parameters for the report.
         """
-        (
+        if self.api_version in ["v1", "v1.1"]:
+            return (
+                self.get_conn()
+                .queries()
+                .runquery(queryId=query_id, body=params)
+                .execute(num_retries=self.num_retries)
+            )
+        else:
+            return (
+                self.get_conn()
+                .queries()
+                .run(queryId=query_id, body=params)
+                .execute(num_retries=self.num_retries)
+            )
+
+    def get_report(self, query_id: str, report_id: str) -> dict:
+        """
+        Retrieves a report.
+
+        :param query_id: Query ID for which report was generated.
+        :param report_id: Report ID to retrieve.
+        """
+        return (
             self.get_conn()
             .queries()
-            .runquery(queryId=query_id, body=params)
+            .reports()
+            .get(queryId=query_id, reportId=report_id)
             .execute(num_retries=self.num_retries)
         )