apache-airflow-providers-google 14.0.0__py3-none-any.whl → 19.1.0rc1__py3-none-any.whl

airflow/providers/google/cloud/hooks/bigquery_dts.py

@@ -23,12 +23,6 @@ from collections.abc import Sequence
 from copy import copy
 from typing import TYPE_CHECKING
 
-from airflow.providers.google.common.consts import CLIENT_INFO
-from airflow.providers.google.common.hooks.base_google import (
-    PROVIDE_PROJECT_ID,
-    GoogleBaseAsyncHook,
-    GoogleBaseHook,
-)
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.bigquery_datatransfer_v1 import DataTransferServiceAsyncClient, DataTransferServiceClient
 from google.cloud.bigquery_datatransfer_v1.types import (
@@ -37,11 +31,17 @@ from google.cloud.bigquery_datatransfer_v1.types import (
     TransferRun,
 )
 
-if TYPE_CHECKING:
-    from googleapiclient.discovery import Resource
+from airflow.providers.google.common.consts import CLIENT_INFO
+from airflow.providers.google.common.hooks.base_google import (
+    PROVIDE_PROJECT_ID,
+    GoogleBaseAsyncHook,
+    GoogleBaseHook,
+)
 
+if TYPE_CHECKING:
     from google.api_core.retry import Retry
     from google.api_core.retry_async import AsyncRetry
+    from googleapiclient.discovery import Resource
 
 
 def get_object_id(obj: dict) -> str:

airflow/providers/google/cloud/hooks/bigtable.py

@@ -22,13 +22,14 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
-from airflow.providers.google.common.consts import CLIENT_INFO
-from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
 from google.cloud.bigtable import Client, enums
 from google.cloud.bigtable.cluster import Cluster
 from google.cloud.bigtable.instance import Instance
 from google.cloud.bigtable.table import ClusterState, Table
 
+from airflow.providers.google.common.consts import CLIENT_INFO
+from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
+
 if TYPE_CHECKING:
     import enum
 

airflow/providers/google/cloud/hooks/cloud_batch.py

@@ -23,9 +23,6 @@ import time
 from collections.abc import Iterable, Sequence
 from typing import TYPE_CHECKING
 
-from airflow.exceptions import AirflowException
-from airflow.providers.google.common.consts import CLIENT_INFO
-from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
 from google.cloud.batch import ListJobsRequest, ListTasksRequest
 from google.cloud.batch_v1 import (
     BatchServiceAsyncClient,
@@ -36,6 +33,10 @@ from google.cloud.batch_v1 import (
     Task,
 )
 
+from airflow.exceptions import AirflowException
+from airflow.providers.google.common.consts import CLIENT_INFO
+from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
+
 if TYPE_CHECKING:
     from google.api_core import operation
     from google.cloud.batch_v1.services.batch_service import pagers
@@ -146,20 +147,18 @@ class CloudBatchHook(GoogleBaseHook):
                 status: JobStatus.State = job.status.state
                 if status == JobStatus.State.SUCCEEDED:
                     return job
-                elif status == JobStatus.State.FAILED:
+                if status == JobStatus.State.FAILED:
                     message = (
                         "Unexpected error in the operation: "
                         "Batch job with name {job_name} has failed its execution."
                     )
                     raise AirflowException(message)
-                elif status == JobStatus.State.DELETION_IN_PROGRESS:
+                if status == JobStatus.State.DELETION_IN_PROGRESS:
                     message = (
-                        "Unexpected error in the operation: "
-                        "Batch job with name {job_name} is being deleted."
+                        "Unexpected error in the operation: Batch job with name {job_name} is being deleted."
                     )
                     raise AirflowException(message)
-                else:
-                    time.sleep(polling_period_seconds)
+                time.sleep(polling_period_seconds)
             except Exception as e:
                 self.log.exception("Exception occurred while checking for job completion.")
                 raise e

airflow/providers/google/cloud/hooks/cloud_build.py

@@ -22,15 +22,16 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
-from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
-from airflow.providers.google.common.consts import CLIENT_INFO
-from airflow.providers.google.common.deprecated import deprecated
-from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
 from google.api_core.client_options import ClientOptions
 from google.api_core.exceptions import AlreadyExists
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.devtools.cloudbuild_v1 import CloudBuildAsyncClient, CloudBuildClient, GetBuildRequest
 
+from airflow.exceptions import AirflowException
+from airflow.providers.google.common.consts import CLIENT_INFO
+from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
+from airflow.providers.google.common.hooks.operation_helpers import OperationHelper
+
 if TYPE_CHECKING:
     from google.api_core.operation import Operation
     from google.api_core.retry import Retry
@@ -41,7 +42,7 @@ if TYPE_CHECKING:
 TIME_TO_SLEEP_IN_SECONDS = 5
 
 
-class CloudBuildHook(GoogleBaseHook):
+class CloudBuildHook(GoogleBaseHook, OperationHelper):
     """
     Hook for the Google Cloud Build Service.
 
@@ -79,14 +80,6 @@ class CloudBuildHook(GoogleBaseHook):
         except Exception:
             raise AirflowException("Could not retrieve Build ID from Operation.")
 
-    def wait_for_operation(self, operation: Operation, timeout: float | None = None):
-        """Wait for long-lasting operation to complete."""
-        try:
-            return operation.result(timeout=timeout)
-        except Exception:
-            error = operation.exception(timeout=timeout)
-            raise AirflowException(error)
-
     def get_conn(self, location: str = "global") -> CloudBuildClient:
         """
         Retrieve the connection to Google Cloud Build.
@@ -184,58 +177,6 @@ class CloudBuildHook(GoogleBaseHook):
 
         return operation, id_
 
-    @GoogleBaseHook.fallback_to_default_project_id
-    @deprecated(
-        planned_removal_date="March 01, 2025",
-        use_instead="create_build_without_waiting_for_result",
-        category=AirflowProviderDeprecationWarning,
-    )
-    def create_build(
-        self,
-        build: dict | Build,
-        project_id: str = PROVIDE_PROJECT_ID,
-        wait: bool = True,
-        retry: Retry | _MethodDefault = DEFAULT,
-        timeout: float | None = None,
-        metadata: Sequence[tuple[str, str]] = (),
-    ) -> Build:
-        """
-        Start a build with the specified configuration.
-
-        :param build: The build resource to create. If a dict is provided, it must be of the same form
-            as the protobuf message `google.cloud.devtools.cloudbuild_v1.types.Build`
-        :param project_id: Optional, Google Cloud Project project_id where the function belongs.
-            If set to None or missing, the default project_id from the GCP connection is used.
-        :param wait: Optional, wait for operation to finish.
-        :param retry: Optional, a retry object used  to retry requests. If `None` is specified, requests
-            will not be retried.
-        :param timeout: Optional, the amount of time, in seconds, to wait for the request to complete.
-            Note that if `retry` is specified, the timeout applies to each individual attempt.
-        :param metadata: Optional, additional metadata that is provided to the method.
-
-        """
-        client = self.get_conn()
-
-        self.log.info("Start creating build...")
-
-        operation = client.create_build(
-            request={"project_id": project_id, "build": build},
-            retry=retry,
-            timeout=timeout,
-            metadata=metadata,
-        )
-
-        id_ = self._get_build_id_from_operation(operation)
-
-        if not wait:
-            return self.get_build(id_=id_, project_id=project_id)
-
-        operation.result()
-
-        self.log.info("Build has been created: %s.", id_)
-
-        return self.get_build(id_=id_, project_id=project_id)
-
     @GoogleBaseHook.fallback_to_default_project_id
     def create_build_trigger(
         self,

airflow/providers/google/cloud/hooks/cloud_composer.py

@@ -18,15 +18,16 @@
 from __future__ import annotations
 
 import asyncio
+import json
 import time
 from collections.abc import MutableSequence, Sequence
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
+from urllib.parse import urlencode, urljoin
 
-from airflow.exceptions import AirflowException
-from airflow.providers.google.common.consts import CLIENT_INFO
-from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
+from aiohttp import ClientSession
 from google.api_core.client_options import ClientOptions
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
+from google.auth.transport.requests import AuthorizedSession, Request
 from google.cloud.orchestration.airflow.service_v1 import (
     EnvironmentsAsyncClient,
     EnvironmentsClient,
@@ -34,6 +35,10 @@ from google.cloud.orchestration.airflow.service_v1 import (
     PollAirflowCommandResponse,
 )
 
+from airflow.exceptions import AirflowException
+from airflow.providers.google.common.consts import CLIENT_INFO
+from airflow.providers.google.common.hooks.base_google import GoogleBaseAsyncHook, GoogleBaseHook
+
 if TYPE_CHECKING:
     from google.api_core.operation import Operation
     from google.api_core.operation_async import AsyncOperation
@@ -51,8 +56,10 @@ if TYPE_CHECKING:
     )
     from google.protobuf.field_mask_pb2 import FieldMask
 
+from airflow.providers.google.common.hooks.operation_helpers import OperationHelper
+
 
-class CloudComposerHook(GoogleBaseHook):
+class CloudComposerHook(GoogleBaseHook, OperationHelper):
     """Hook for Google Cloud Composer APIs."""
 
     client_options = ClientOptions(api_endpoint="composer.googleapis.com:443")
@@ -73,13 +80,33 @@ class CloudComposerHook(GoogleBaseHook):
             client_options=self.client_options,
         )
 
-    def wait_for_operation(self, operation: Operation, timeout: float | None = None):
-        """Wait for long-lasting operation to complete."""
-        try:
-            return operation.result(timeout=timeout)
-        except Exception:
-            error = operation.exception(timeout=timeout)
-            raise AirflowException(error)
+    def make_composer_airflow_api_request(
+        self,
+        method: str,
+        airflow_uri: str,
+        path: str,
+        data: Any | None = None,
+        timeout: float | None = None,
+    ):
+        """
+        Make a request to Cloud Composer environment's web server.
+
+        :param method: The request method to use ('GET', 'OPTIONS', 'HEAD', 'POST', 'PUT', 'PATCH', 'DELETE').
+        :param airflow_uri: The URI of the Apache Airflow Web UI hosted within this environment.
+        :param path: The path to send the request.
+        :param data: Dictionary, list of tuples, bytes, or file-like object to send in the body of the request.
+        :param timeout: The timeout for this request.
+        """
+        authed_session = AuthorizedSession(self.get_credentials())
+
+        resp = authed_session.request(
+            method=method,
+            url=urljoin(airflow_uri, path),
+            data=data,
+            headers={"Content-Type": "application/json"},
+            timeout=timeout,
+        )
+        return resp
 
     def get_operation(self, operation_name):
         return self.get_environment_client().transport.operations_client.get_operation(name=operation_name)
@@ -413,20 +440,160 @@ class CloudComposerHook(GoogleBaseHook):
             self.log.info("Waiting for result...")
             time.sleep(poll_interval)
 
+    def trigger_dag_run(
+        self,
+        composer_airflow_uri: str,
+        composer_dag_id: str,
+        composer_dag_conf: dict | None = None,
+        timeout: float | None = None,
+    ) -> dict:
+        """
+        Trigger DAG run for provided Apache Airflow Web UI hosted within Composer environment.
+
+        :param composer_airflow_uri: The URI of the Apache Airflow Web UI hosted within Composer environment.
+        :param composer_dag_id: The ID of DAG which will be triggered.
+        :param composer_dag_conf: Configuration parameters for the DAG run.
+        :param timeout: The timeout for this request.
+        """
+        response = self.make_composer_airflow_api_request(
+            method="POST",
+            airflow_uri=composer_airflow_uri,
+            path=f"/api/v1/dags/{composer_dag_id}/dagRuns",
+            data=json.dumps(
+                {
+                    "conf": composer_dag_conf or {},
+                }
+            ),
+            timeout=timeout,
+        )
+
+        if response.status_code != 200:
+            self.log.error(response.text)
+            response.raise_for_status()
+
+        return response.json()
+
+    def get_dag_runs(
+        self,
+        composer_airflow_uri: str,
+        composer_dag_id: str,
+        timeout: float | None = None,
+    ) -> dict:
+        """
+        Get the list of dag runs for provided DAG.
 
-class CloudComposerAsyncHook(GoogleBaseHook):
+        :param composer_airflow_uri: The URI of the Apache Airflow Web UI hosted within Composer environment.
+        :param composer_dag_id: The ID of DAG.
+        :param timeout: The timeout for this request.
+        """
+        response = self.make_composer_airflow_api_request(
+            method="GET",
+            airflow_uri=composer_airflow_uri,
+            path=f"/api/v1/dags/{composer_dag_id}/dagRuns",
+            timeout=timeout,
+        )
+
+        if response.status_code != 200:
+            self.log.error(
+                "Failed to get DAG runs for dag_id=%s from %s (status=%s): %s",
+                composer_dag_id,
+                composer_airflow_uri,
+                response.status_code,
+                response.text,
+            )
+            response.raise_for_status()
+
+        return response.json()
+
+    def get_task_instances(
+        self,
+        composer_airflow_uri: str,
+        composer_dag_id: str,
+        query_parameters: dict | None = None,
+        timeout: float | None = None,
+    ) -> dict:
+        """
+        Get the list of task instances for provided DAG.
+
+        :param composer_airflow_uri: The URI of the Apache Airflow Web UI hosted within Composer environment.
+        :param composer_dag_id: The ID of DAG.
+        :query_parameters: Query parameters for this request.
+        :param timeout: The timeout for this request.
+        """
+        query_string = f"?{urlencode(query_parameters)}" if query_parameters else ""
+
+        response = self.make_composer_airflow_api_request(
+            method="GET",
+            airflow_uri=composer_airflow_uri,
+            path=f"/api/v1/dags/{composer_dag_id}/dagRuns/~/taskInstances{query_string}",
+            timeout=timeout,
+        )
+
+        if response.status_code != 200:
+            self.log.error(
+                "Failed to get task instances for dag_id=%s from %s (status=%s): %s",
+                composer_dag_id,
+                composer_airflow_uri,
+                response.status_code,
+                response.text,
+            )
+            response.raise_for_status()
+
+        return response.json()
+
+
+class CloudComposerAsyncHook(GoogleBaseAsyncHook):
     """Hook for Google Cloud Composer async APIs."""
 
+    sync_hook_class = CloudComposerHook
+
     client_options = ClientOptions(api_endpoint="composer.googleapis.com:443")
 
-    def get_environment_client(self) -> EnvironmentsAsyncClient:
+    async def get_environment_client(self) -> EnvironmentsAsyncClient:
         """Retrieve client library object that allow access Environments service."""
+        sync_hook = await self.get_sync_hook()
         return EnvironmentsAsyncClient(
-            credentials=self.get_credentials(),
+            credentials=sync_hook.get_credentials(),
             client_info=CLIENT_INFO,
             client_options=self.client_options,
         )
 
+    async def make_composer_airflow_api_request(
+        self,
+        method: str,
+        airflow_uri: str,
+        path: str,
+        data: Any | None = None,
+        timeout: float | None = None,
+    ):
+        """
+        Make a request to Cloud Composer environment's web server.
+
+        :param method: The request method to use ('GET', 'OPTIONS', 'HEAD', 'POST', 'PUT', 'PATCH', 'DELETE').
+        :param airflow_uri: The URI of the Apache Airflow Web UI hosted within this environment.
+        :param path: The path to send the request.
+        :param data: Dictionary, list of tuples, bytes, or file-like object to send in the body of the request.
+        :param timeout: The timeout for this request.
+        """
+        sync_hook = await self.get_sync_hook()
+        credentials = sync_hook.get_credentials()
+
+        if not credentials.valid:
+            credentials.refresh(Request())
+
+        async with ClientSession() as session:
+            async with session.request(
+                method=method,
+                url=urljoin(airflow_uri, path),
+                data=data,
+                headers={
+                    "Content-Type": "application/json",
+                    "Authorization": f"Bearer {credentials.token}",
+                },
+                timeout=timeout,
+            ) as response:
+                return await response.json(), response.status
+
     def get_environment_name(self, project_id, region, environment_id):
         return f"projects/{project_id}/locations/{region}/environments/{environment_id}"
 
@@ -434,9 +601,8 @@ class CloudComposerAsyncHook(GoogleBaseHook):
         return f"projects/{project_id}/locations/{region}"
 
     async def get_operation(self, operation_name):
-        return await self.get_environment_client().transport.operations_client.get_operation(
-            name=operation_name
-        )
+        client = await self.get_environment_client()
+        return await client.transport.operations_client.get_operation(name=operation_name)
 
     @GoogleBaseHook.fallback_to_default_project_id
     async def create_environment(
@@ -459,7 +625,7 @@ class CloudComposerAsyncHook(GoogleBaseHook):
         :param timeout: The timeout for this request.
         :param metadata: Strings which should be sent along with the request as metadata.
         """
-        client = self.get_environment_client()
+        client = await self.get_environment_client()
         return await client.create_environment(
             request={"parent": self.get_parent(project_id, region), "environment": environment},
             retry=retry,
@@ -487,7 +653,7 @@ class CloudComposerAsyncHook(GoogleBaseHook):
         :param timeout: The timeout for this request.
         :param metadata: Strings which should be sent along with the request as metadata.
         """
-        client = self.get_environment_client()
+        client = await self.get_environment_client()
         name = self.get_environment_name(project_id, region, environment_id)
         return await client.delete_environment(
             request={"name": name}, retry=retry, timeout=timeout, metadata=metadata
@@ -523,7 +689,7 @@ class CloudComposerAsyncHook(GoogleBaseHook):
         :param timeout: The timeout for this request.
         :param metadata: Strings which should be sent along with the request as metadata.
         """
-        client = self.get_environment_client()
+        client = await self.get_environment_client()
         name = self.get_environment_name(project_id, region, environment_id)
 
         return await client.update_environment(
@@ -533,6 +699,35 @@ class CloudComposerAsyncHook(GoogleBaseHook):
             metadata=metadata,
         )
 
+    @GoogleBaseHook.fallback_to_default_project_id
+    async def get_environment(
+        self,
+        project_id: str,
+        region: str,
+        environment_id: str,
+        retry: AsyncRetry | _MethodDefault = DEFAULT,
+        timeout: float | None = None,
+        metadata: Sequence[tuple[str, str]] = (),
+    ) -> Environment:
+        """
+        Get an existing environment.
+
+        :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
+        :param region: Required. The ID of the Google Cloud region that the service belongs to.
+        :param environment_id: Required. The ID of the Google Cloud environment that the service belongs to.
+        :param retry: Designation of what errors, if any, should be retried.
+        :param timeout: The timeout for this request.
+        :param metadata: Strings which should be sent along with the request as metadata.
+        """
+        client = await self.get_environment_client()
+
+        return await client.get_environment(
+            request={"name": self.get_environment_name(project_id, region, environment_id)},
+            retry=retry,
+            timeout=timeout,
+            metadata=metadata,
+        )
+
     @GoogleBaseHook.fallback_to_default_project_id
     async def execute_airflow_command(
         self,
@@ -561,7 +756,7 @@ class CloudComposerAsyncHook(GoogleBaseHook):
         :param timeout: The timeout for this request.
         :param metadata: Strings which should be sent along with the request as metadata.
         """
-        client = self.get_environment_client()
+        client = await self.get_environment_client()
 
         return await client.execute_airflow_command(
             request={
@@ -603,7 +798,7 @@ class CloudComposerAsyncHook(GoogleBaseHook):
         :param timeout: The timeout for this request.
         :param metadata: Strings which should be sent along with the request as metadata.
         """
-        client = self.get_environment_client()
+        client = await self.get_environment_client()
 
         return await client.poll_airflow_command(
             request={
@@ -647,9 +842,82 @@ class CloudComposerAsyncHook(GoogleBaseHook):
                 self.log.exception("Exception occurred while polling CMD result")
                 raise AirflowException(ex)
 
-            result_dict = PollAirflowCommandResponse.to_dict(result)
+            try:
+                result_dict = PollAirflowCommandResponse.to_dict(result)
+            except Exception as ex:
+                self.log.exception("Exception occurred while transforming PollAirflowCommandResponse")
+                raise AirflowException(ex)
+
             if result_dict["output_end"]:
                 return result_dict
 
             self.log.info("Sleeping for %s seconds.", poll_interval)
             await asyncio.sleep(poll_interval)
+
+    async def get_dag_runs(
+        self,
+        composer_airflow_uri: str,
+        composer_dag_id: str,
+        timeout: float | None = None,
+    ) -> dict:
+        """
+        Get the list of dag runs for provided DAG.
+
+        :param composer_airflow_uri: The URI of the Apache Airflow Web UI hosted within Composer environment.
+        :param composer_dag_id: The ID of DAG.
+        :param timeout: The timeout for this request.
+        """
+        response_body, response_status_code = await self.make_composer_airflow_api_request(
+            method="GET",
+            airflow_uri=composer_airflow_uri,
+            path=f"/api/v1/dags/{composer_dag_id}/dagRuns",
+            timeout=timeout,
+        )
+
+        if response_status_code != 200:
+            self.log.error(
+                "Failed to get DAG runs for dag_id=%s from %s (status=%s): %s",
+                composer_dag_id,
+                composer_airflow_uri,
+                response_status_code,
+                response_body["title"],
+            )
+            raise AirflowException(response_body["title"])
+
+        return response_body
+
+    async def get_task_instances(
+        self,
+        composer_airflow_uri: str,
+        composer_dag_id: str,
+        query_parameters: dict | None = None,
+        timeout: float | None = None,
+    ) -> dict:
+        """
+        Get the list of task instances for provided DAG.
+
+        :param composer_airflow_uri: The URI of the Apache Airflow Web UI hosted within Composer environment.
+        :param composer_dag_id: The ID of DAG.
+        :query_parameters: Query parameters for this request.
+        :param timeout: The timeout for this request.
+        """
+        query_string = f"?{urlencode(query_parameters)}" if query_parameters else ""
+
+        response_body, response_status_code = await self.make_composer_airflow_api_request(
+            method="GET",
+            airflow_uri=composer_airflow_uri,
+            path=f"/api/v1/dags/{composer_dag_id}/dagRuns/~/taskInstances{query_string}",
+            timeout=timeout,
+        )
+
+        if response_status_code != 200:
+            self.log.error(
+                "Failed to get task instances for dag_id=%s from %s (status=%s): %s",
+                composer_dag_id,
+                composer_airflow_uri,
+                response_status_code,
+                response_body["title"],
+            )
+            raise AirflowException(response_body["title"])
+
+        return response_body