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

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

@@ -18,12 +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 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,
@@ -33,7 +37,7 @@ from google.cloud.orchestration.airflow.service_v1 import (
 
 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 airflow.providers.google.common.hooks.base_google import GoogleBaseAsyncHook, GoogleBaseHook
 
 if TYPE_CHECKING:
     from google.api_core.operation import Operation
@@ -76,6 +80,34 @@ class CloudComposerHook(GoogleBaseHook, OperationHelper):
             client_options=self.client_options,
         )
 
+    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)
 
@@ -408,20 +440,160 @@ class CloudComposerHook(GoogleBaseHook, OperationHelper):
             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,
+        )
 
-class CloudComposerAsyncHook(GoogleBaseHook):
+        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.
+
+        :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}"
 
@@ -429,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(
@@ -454,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,
@@ -482,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
@@ -518,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(
@@ -528,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,
@@ -556,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={
@@ -598,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={
@@ -642,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

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

@@ -0,0 +1,109 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import TYPE_CHECKING
+
+from google.cloud.logging_v2.services.config_service_v2 import ConfigServiceV2Client
+from google.cloud.logging_v2.types import (
+    CreateSinkRequest,
+    DeleteSinkRequest,
+    GetSinkRequest,
+    ListSinksRequest,
+    LogSink,
+    UpdateSinkRequest,
+)
+
+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.protobuf.field_mask_pb2 import FieldMask
+
+
+class CloudLoggingHook(GoogleBaseHook):
+    """
+    Hook for Google Cloud Logging Log Sinks API.
+
+    :param gcp_conn_id: The connection ID to use when fetching connection info.
+    :param impersonation_chain: Optional service account to impersonate.
+    """
+
+    def __init__(
+        self,
+        gcp_conn_id: str = "google_cloud_default",
+        impersonation_chain: str | Sequence[str] | None = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(gcp_conn_id=gcp_conn_id, impersonation_chain=impersonation_chain, **kwargs)
+        self._client: ConfigServiceV2Client | None = None
+
+    def get_conn(self) -> ConfigServiceV2Client:
+        """Return the Google Cloud Logging Config client."""
+        if not self._client:
+            self._client = ConfigServiceV2Client(credentials=self.get_credentials(), client_info=CLIENT_INFO)
+        return self._client
+
+    def get_parent(self, project_id):
+        return f"projects/{project_id}"
+
+    @GoogleBaseHook.fallback_to_default_project_id
+    def create_sink(
+        self, sink: LogSink | dict, unique_writer_identity: bool = True, project_id: str = PROVIDE_PROJECT_ID
+    ) -> LogSink:
+        if isinstance(sink, dict):
+            sink = LogSink(**sink)
+        request = CreateSinkRequest(
+            parent=self.get_parent(project_id), sink=sink, unique_writer_identity=unique_writer_identity
+        )
+        return self.get_conn().create_sink(request=request)
+
+    @GoogleBaseHook.fallback_to_default_project_id
+    def get_sink(self, sink_name: str, project_id: str = PROVIDE_PROJECT_ID) -> LogSink:
+        request = GetSinkRequest(sink_name=f"projects/{project_id}/sinks/{sink_name}")
+        return self.get_conn().get_sink(request=request)
+
+    @GoogleBaseHook.fallback_to_default_project_id
+    def list_sinks(self, page_size: int | None = None, project_id: str = PROVIDE_PROJECT_ID) -> list[LogSink]:
+        request = ListSinksRequest(parent=self.get_parent(project_id), page_size=page_size)
+        return list(self.get_conn().list_sinks(request=request))
+
+    @GoogleBaseHook.fallback_to_default_project_id
+    def delete_sink(self, sink_name: str, project_id: str = PROVIDE_PROJECT_ID) -> None:
+        request = DeleteSinkRequest(sink_name=f"projects/{project_id}/sinks/{sink_name}")
+        self.get_conn().delete_sink(request=request)
+
+    @GoogleBaseHook.fallback_to_default_project_id
+    def update_sink(
+        self,
+        sink_name: str,
+        sink: LogSink | dict,
+        unique_writer_identity: bool,
+        update_mask: FieldMask | dict,
+        project_id: str = PROVIDE_PROJECT_ID,
+    ) -> LogSink:
+        if isinstance(sink, dict):
+            sink = LogSink(**sink)
+        request = UpdateSinkRequest(
+            sink_name=f"projects/{project_id}/sinks/{sink_name}",
+            sink=sink,
+            unique_writer_identity=unique_writer_identity,
+            update_mask=update_mask,
+        )
+        return self.get_conn().update_sink(request=request)

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

@@ -38,11 +38,15 @@ from google.cloud.run_v2 import (
     ServicesClient,
     UpdateJobRequest,
 )
-from google.longrunning import operations_pb2  # type: ignore[attr-defined]
+from google.longrunning import operations_pb2
 
 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.base_google import (
+    PROVIDE_PROJECT_ID,
+    GoogleBaseAsyncHook,
+    GoogleBaseHook,
+)
 
 if TYPE_CHECKING:
     from google.api_core import operation
@@ -159,7 +163,7 @@ class CloudRunHook(GoogleBaseHook):
         return list(itertools.islice(jobs, limit))
 
 
-class CloudRunAsyncHook(GoogleBaseHook):
+class CloudRunAsyncHook(GoogleBaseAsyncHook):
     """
     Async hook for the Google Cloud Run service.
 
@@ -174,6 +178,8 @@ class CloudRunAsyncHook(GoogleBaseHook):
         account from the list granting this role to the originating account.
     """
 
+    sync_hook_class = CloudRunHook
+
     def __init__(
         self,
         gcp_conn_id: str = "google_cloud_default",
@@ -183,16 +189,16 @@ class CloudRunAsyncHook(GoogleBaseHook):
         self._client: JobsAsyncClient | None = None
         super().__init__(gcp_conn_id=gcp_conn_id, impersonation_chain=impersonation_chain, **kwargs)
 
-    def get_conn(self):
+    async def get_conn(self):
         if self._client is None:
-            self._client = JobsAsyncClient(credentials=self.get_credentials(), client_info=CLIENT_INFO)
+            sync_hook = await self.get_sync_hook()
+            self._client = JobsAsyncClient(credentials=sync_hook.get_credentials(), client_info=CLIENT_INFO)
 
         return self._client
 
     async def get_operation(self, operation_name: str) -> operations_pb2.Operation:
-        return await self.get_conn().get_operation(
-            operations_pb2.GetOperationRequest(name=operation_name), timeout=120
-        )
+        conn = await self.get_conn()
+        return await conn.get_operation(operations_pb2.GetOperationRequest(name=operation_name), timeout=120)
 
 
 class CloudRunServiceHook(GoogleBaseHook):
@@ -258,7 +264,7 @@ class CloudRunServiceHook(GoogleBaseHook):
         return operation.result()
 
 
-class CloudRunServiceAsyncHook(GoogleBaseHook):
+class CloudRunServiceAsyncHook(GoogleBaseAsyncHook):
     """
     Async hook for the Google Cloud Run services.
 
@@ -273,6 +279,8 @@ class CloudRunServiceAsyncHook(GoogleBaseHook):
         account from the list granting this role to the originating account.
     """
 
+    sync_hook_class = CloudRunServiceHook
+
     def __init__(
         self,
         gcp_conn_id: str = "google_cloud_default",