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

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

@@ -19,15 +19,15 @@ from __future__ import annotations
 
 import shlex
 from collections.abc import Sequence
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
-from google.api_core.exceptions import AlreadyExists
+from google.api_core.exceptions import AlreadyExists, NotFound
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.orchestration.airflow.service_v1 import ImageVersion
 from google.cloud.orchestration.airflow.service_v1.types import Environment, ExecuteAirflowCommandResponse
 
 from airflow.configuration import conf
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.cloud.hooks.cloud_composer import CloudComposerHook
 from airflow.providers.google.cloud.links.base import BaseGoogleLink
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
@@ -41,7 +41,7 @@ if TYPE_CHECKING:
     from google.api_core.retry import Retry
     from google.protobuf.field_mask_pb2 import FieldMask
 
-    from airflow.utils.context import Context
+    from airflow.providers.common.compat.sdk import Context
 
 CLOUD_COMPOSER_BASE_LINK = "https://console.cloud.google.com/composer/environments"
 CLOUD_COMPOSER_DETAILS_LINK = (
@@ -57,25 +57,6 @@ class CloudComposerEnvironmentLink(BaseGoogleLink):
     key = "composer_conf"
     format_str = CLOUD_COMPOSER_DETAILS_LINK
 
-    @staticmethod
-    def persist(
-        operator_instance: (
-            CloudComposerCreateEnvironmentOperator
-            | CloudComposerUpdateEnvironmentOperator
-            | CloudComposerGetEnvironmentOperator
-        ),
-        context: Context,
-    ) -> None:
-        operator_instance.xcom_push(
-            context,
-            key=CloudComposerEnvironmentLink.key,
-            value={
-                "project_id": operator_instance.project_id,
-                "region": operator_instance.region,
-                "environment_id": operator_instance.environment_id,
-            },
-        )
-
 
 class CloudComposerEnvironmentsLink(BaseGoogleLink):
     """Helper class for constructing Cloud Composer Environment Link."""
@@ -84,16 +65,6 @@ class CloudComposerEnvironmentsLink(BaseGoogleLink):
     key = "composer_conf"
     format_str = CLOUD_COMPOSER_ENVIRONMENTS_LINK
 
-    @staticmethod
-    def persist(operator_instance: CloudComposerListEnvironmentsOperator, context: Context) -> None:
-        operator_instance.xcom_push(
-            context,
-            key=CloudComposerEnvironmentsLink.key,
-            value={
-                "project_id": operator_instance.project_id,
-            },
-        )
-
 
 class CloudComposerCreateEnvironmentOperator(GoogleCloudBaseOperator):
     """
@@ -159,6 +130,14 @@ class CloudComposerCreateEnvironmentOperator(GoogleCloudBaseOperator):
         self.deferrable = deferrable
         self.pooling_period_seconds = pooling_period_seconds
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "project_id": self.project_id,
+            "region": self.region,
+            "environment_id": self.environment_id,
+        }
+
     def execute(self, context: Context):
         hook = CloudComposerHook(
             gcp_conn_id=self.gcp_conn_id,
@@ -171,7 +150,7 @@ class CloudComposerCreateEnvironmentOperator(GoogleCloudBaseOperator):
         else:
             self.environment["name"] = name
 
-        CloudComposerEnvironmentLink.persist(operator_instance=self, context=context)
+        CloudComposerEnvironmentLink.persist(context=context)
         try:
             result = hook.create_environment(
                 project_id=self.project_id,
@@ -370,6 +349,14 @@ class CloudComposerGetEnvironmentOperator(GoogleCloudBaseOperator):
         self.gcp_conn_id = gcp_conn_id
         self.impersonation_chain = impersonation_chain
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "project_id": self.project_id,
+            "region": self.region,
+            "environment_id": self.environment_id,
+        }
+
     def execute(self, context: Context):
         hook = CloudComposerHook(
             gcp_conn_id=self.gcp_conn_id,
@@ -384,8 +371,7 @@ class CloudComposerGetEnvironmentOperator(GoogleCloudBaseOperator):
             timeout=self.timeout,
             metadata=self.metadata,
         )
-
-        CloudComposerEnvironmentLink.persist(operator_instance=self, context=context)
+        CloudComposerEnvironmentLink.persist(context=context)
         return Environment.to_dict(result)
 
 
@@ -445,12 +431,17 @@ class CloudComposerListEnvironmentsOperator(GoogleCloudBaseOperator):
         self.gcp_conn_id = gcp_conn_id
         self.impersonation_chain = impersonation_chain
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context):
         hook = CloudComposerHook(
             gcp_conn_id=self.gcp_conn_id,
             impersonation_chain=self.impersonation_chain,
         )
-        CloudComposerEnvironmentsLink.persist(operator_instance=self, context=context)
         result = hook.list_environments(
             project_id=self.project_id,
             region=self.region,
@@ -532,6 +523,14 @@ class CloudComposerUpdateEnvironmentOperator(GoogleCloudBaseOperator):
         self.deferrable = deferrable
         self.pooling_period_seconds = pooling_period_seconds
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "project_id": self.project_id,
+            "region": self.region,
+            "environment_id": self.environment_id,
+        }
+
     def execute(self, context: Context):
         hook = CloudComposerHook(
             gcp_conn_id=self.gcp_conn_id,
@@ -549,7 +548,7 @@ class CloudComposerUpdateEnvironmentOperator(GoogleCloudBaseOperator):
             metadata=self.metadata,
         )
 
-        CloudComposerEnvironmentLink.persist(operator_instance=self, context=context)
+        CloudComposerEnvironmentLink.persist(context=context)
         if not self.deferrable:
             environment = hook.wait_for_operation(timeout=self.timeout, operation=result)
             return Environment.to_dict(environment)
@@ -765,9 +764,15 @@ class CloudComposerRunAirflowCLICommandOperator(GoogleCloudBaseOperator):
             metadata=self.metadata,
             poll_interval=self.poll_interval,
         )
-        result_str = self._merge_cmd_output_result(result)
-        self.log.info("Command execution result:\n%s", result_str)
-        return result
+        exit_code = result.get("exit_info", {}).get("exit_code")
+        if exit_code == 0:
+            result_str = self._merge_cmd_output_result(result)
+            self.log.info("Command execution result:\n%s", result_str)
+            return result
+
+        error_output = "".join(line["content"] for line in result.get("error", []))
+        message = f"Airflow CLI command failed with exit code {exit_code}.\nError output:\n{error_output}"
+        raise AirflowException(message)
 
     def execute_complete(self, context: Context, event: dict) -> dict:
         if event and event["status"] == "error":
@@ -793,3 +798,86 @@ class CloudComposerRunAirflowCLICommandOperator(GoogleCloudBaseOperator):
         """Merge output to one string."""
         result_str = "\n".join(line_dict["content"] for line_dict in result["output"])
         return result_str
+
+
+class CloudComposerTriggerDAGRunOperator(GoogleCloudBaseOperator):
+    """
+    Trigger DAG run for provided Composer environment.
+
+    :param project_id: The ID of the Google Cloud project that the service belongs to.
+    :param region: The ID of the Google Cloud region that the service belongs to.
+    :param environment_id: The ID of the Google Cloud environment that the service belongs to.
+    :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.
+    :param gcp_conn_id: The connection ID used to connect to Google Cloud Platform.
+    :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 = (
+        "project_id",
+        "region",
+        "environment_id",
+        "composer_dag_id",
+        "impersonation_chain",
+    )
+
+    def __init__(
+        self,
+        *,
+        project_id: str,
+        region: str,
+        environment_id: str,
+        composer_dag_id: str,
+        composer_dag_conf: dict | None = None,
+        timeout: float | None = None,
+        gcp_conn_id: str = "google_cloud_default",
+        impersonation_chain: str | Sequence[str] | None = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.region = region
+        self.environment_id = environment_id
+        self.composer_dag_id = composer_dag_id
+        self.composer_dag_conf = composer_dag_conf or {}
+        self.timeout = timeout
+        self.gcp_conn_id = gcp_conn_id
+        self.impersonation_chain = impersonation_chain
+
+    def execute(self, context: Context):
+        hook = CloudComposerHook(
+            gcp_conn_id=self.gcp_conn_id,
+            impersonation_chain=self.impersonation_chain,
+        )
+        try:
+            environment = hook.get_environment(
+                project_id=self.project_id,
+                region=self.region,
+                environment_id=self.environment_id,
+                timeout=self.timeout,
+            )
+        except NotFound as not_found_err:
+            self.log.info("The Composer environment %s does not exist.", self.environment_id)
+            raise AirflowException(not_found_err)
+        composer_airflow_uri = environment.config.airflow_uri
+
+        self.log.info(
+            "Triggering the DAG %s on the %s environment...", self.composer_dag_id, self.environment_id
+        )
+        dag_run = hook.trigger_dag_run(
+            composer_airflow_uri=composer_airflow_uri,
+            composer_dag_id=self.composer_dag_id,
+            composer_dag_conf=self.composer_dag_conf,
+            timeout=self.timeout,
+        )
+        self.log.info("The DAG %s was triggered with Run ID: %s", self.composer_dag_id, dag_run["dag_run_id"])
+
+        return dag_run

airflow/providers/google/cloud/operators/cloud_logging_sink.py

@@ -0,0 +1,341 @@
+#
+# 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, Any
+
+import google.cloud.exceptions
+from google.api_core.exceptions import AlreadyExists
+from google.cloud.logging_v2.types import LogSink
+
+from airflow.providers.common.compat.sdk import AirflowException
+from airflow.providers.google.cloud.hooks.cloud_logging import CloudLoggingHook
+from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
+
+if TYPE_CHECKING:
+    from google.protobuf.field_mask_pb2 import FieldMask
+
+    from airflow.providers.common.compat.sdk import Context
+
+
+def _validate_inputs(obj, required_fields: list[str]) -> None:
+    """Validate that all required fields are present on self."""
+    missing = [field for field in required_fields if not getattr(obj, field, None)]
+    if missing:
+        raise AirflowException(
+            f"Required parameters are missing: {missing}. These must be passed as keyword parameters."
+        )
+
+
+def _get_field(obj, field_name):
+    """Supports both dict and protobuf-like objects."""
+    if isinstance(obj, dict):
+        return obj.get(field_name)
+    return getattr(obj, field_name, None)
+
+
+class CloudLoggingCreateSinkOperator(GoogleCloudBaseOperator):
+    """
+    Creates a Cloud Logging export sink in a GCP project.
+
+    This operator creates a sink that exports log entries from Cloud Logging
+    to destinations like Cloud Storage, BigQuery, or Pub/Sub.
+
+    :param project_id: Required. ID of the Google Cloud project where the sink will be created.
+    :param sink_config: Required. The full sink configuration as a dictionary or a LogSink object.
+        See: https://cloud.google.com/logging/docs/reference/v2/rest/v2/projects.sinks
+    :param unique_writer_identity: If True, creates a unique service account for the sink.
+        If False, uses the default Google-managed service account.
+    :param gcp_conn_id: Optional. The connection ID used to connect to Google Cloud. Defaults to "google_cloud_default".
+    :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] = (
+        "project_id",
+        "sink_config",
+        "gcp_conn_id",
+        "impersonation_chain",
+        "unique_writer_identity",
+    )
+
+    def __init__(
+        self,
+        project_id: str,
+        sink_config: dict | LogSink,
+        unique_writer_identity: bool = False,
+        gcp_conn_id: str = "google_cloud_default",
+        impersonation_chain: str | Sequence[str] | None = None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.sink_config = sink_config
+        self.unique_writer_identity = unique_writer_identity
+        self.gcp_conn_id = gcp_conn_id
+        self.impersonation_chain = impersonation_chain
+
+    def execute(self, context: Context) -> dict[str, Any]:
+        """Execute the operator."""
+        _validate_inputs(self, required_fields=["project_id", "sink_config"])
+        hook = CloudLoggingHook(gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain)
+
+        try:
+            self.log.info(
+                "Creating log sink '%s' in project '%s'",
+                _get_field(self.sink_config, "name"),
+                self.project_id,
+            )
+            self.log.info("Destination: %s", _get_field(self.sink_config, "destination"))
+
+            response = hook.create_sink(
+                sink=self.sink_config,
+                unique_writer_identity=self.unique_writer_identity,
+                project_id=self.project_id,
+            )
+
+            self.log.info("Log sink created successfully: %s", response.name)
+
+            if self.unique_writer_identity and hasattr(response, "writer_identity"):
+                self.log.info("Writer identity: %s", response.writer_identity)
+                self.log.info("Remember to grant appropriate permissions to the writer identity")
+
+            return LogSink.to_dict(response)
+
+        except AlreadyExists:
+            self.log.info(
+                "Already existed log sink, sink_name=%s, project_id=%s",
+                _get_field(self.sink_config, "name"),
+                self.project_id,
+            )
+            existing_sink = hook.get_sink(
+                sink_name=_get_field(self.sink_config, "name"), project_id=self.project_id
+            )
+            return LogSink.to_dict(existing_sink)
+
+        except google.cloud.exceptions.GoogleCloudError as e:
+            self.log.error("An error occurred. Exiting.")
+            raise e
+
+
+class CloudLoggingDeleteSinkOperator(GoogleCloudBaseOperator):
+    """
+    Deletes a Cloud Logging export sink from a GCP project.
+
+    :param sink_name: Required. Name of the sink to delete.
+    :param project_id: Required. The ID of the Google Cloud project.
+    :param gcp_conn_id: Optional. The connection ID to use for connecting to Google Cloud.
+        Defaults to "google_cloud_default".
+    :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] = ("sink_name", "project_id", "gcp_conn_id", "impersonation_chain")
+
+    def __init__(
+        self,
+        sink_name: str,
+        project_id: str,
+        gcp_conn_id: str = "google_cloud_default",
+        impersonation_chain: str | Sequence[str] | None = None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.sink_name = sink_name
+        self.project_id = project_id
+        self.gcp_conn_id = gcp_conn_id
+        self.impersonation_chain = impersonation_chain
+
+    def execute(self, context: Context) -> None:
+        """Execute the operator."""
+        _validate_inputs(self, ["sink_name", "project_id"])
+        hook = CloudLoggingHook(gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain)
+
+        try:
+            self.log.info("Deleting log sink '%s' from project '%s'", self.sink_name, self.project_id)
+            hook.delete_sink(sink_name=self.sink_name, project_id=self.project_id)
+            self.log.info("Log sink '%s' deleted successfully", self.sink_name)
+
+        except google.cloud.exceptions.NotFound as e:
+            self.log.error("An error occurred. Not Found.")
+            raise e
+        except google.cloud.exceptions.GoogleCloudError as e:
+            self.log.error("An error occurred. Exiting.")
+            raise e
+
+
+class CloudLoggingUpdateSinkOperator(GoogleCloudBaseOperator):
+    """
+    Updates an existing Cloud Logging export sink.
+
+    :param project_id: Required. The ID of the Google Cloud project that contains the sink.
+    :param sink_name: Required. The name of the sink to update.
+    :param sink_config: Required. The updated sink configuration. Can be a dictionary or a
+        `google.cloud.logging_v2.types.LogSink` object. Refer to:
+        https://cloud.google.com/logging/docs/reference/v2/rest/v2/projects.sinks
+    :param update_mask: Required. A FieldMask or dictionary specifying which fields of the sink
+        should be updated. For example, to update the destination and filter, use:
+        `{"paths": ["destination", "filter"]}`.
+    :param unique_writer_identity: Optional. When set to True, a new unique service account
+        will be created for the sink. Defaults to False.
+    :param gcp_conn_id: Optional. The connection ID used to connect to Google Cloud.
+        Defaults to "google_cloud_default".
+    :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] = (
+        "sink_name",
+        "project_id",
+        "update_mask",
+        "sink_config",
+        "unique_writer_identity",
+        "gcp_conn_id",
+        "impersonation_chain",
+    )
+
+    def __init__(
+        self,
+        project_id: str,
+        sink_name: str,
+        sink_config: dict | LogSink,
+        update_mask: FieldMask | dict,
+        unique_writer_identity: bool = False,
+        gcp_conn_id: str = "google_cloud_default",
+        impersonation_chain: str | Sequence[str] | None = None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.sink_name = sink_name
+        self.sink_config = sink_config
+        self.update_mask = update_mask
+        self.unique_writer_identity = unique_writer_identity
+        self.gcp_conn_id = gcp_conn_id
+        self.impersonation_chain = impersonation_chain
+
+    def execute(self, context: Context) -> dict[str, Any]:
+        """Execute the operator."""
+        _validate_inputs(self, ["sink_name", "project_id", "sink_config", "update_mask"])
+        hook = CloudLoggingHook(gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain)
+
+        try:
+            current_sink = hook.get_sink(sink_name=self.sink_name, project_id=self.project_id)
+            self.log.info("Current log sink configuration: '%s'.", LogSink.to_dict(current_sink))
+
+            self.log.info("Updating log sink '%s' in project '%s'", self.sink_name, self.project_id)
+            if isinstance(self.update_mask, dict) and "paths" in self.update_mask:
+                paths = self.update_mask["paths"]
+            elif hasattr(self.update_mask, "paths"):
+                paths = self.update_mask.paths
+
+            self.log.info("Updating fields: %s", ", ".join(paths))
+
+            response = hook.update_sink(
+                sink_name=self.sink_name,
+                sink=self.sink_config,
+                unique_writer_identity=self.unique_writer_identity,
+                project_id=self.project_id,
+                update_mask=self.update_mask,
+            )
+            self.log.info("Log sink updated successfully: %s", response.name)
+            return LogSink.to_dict(response)
+
+        except google.cloud.exceptions.NotFound as e:
+            self.log.error("An error occurred. Not Found.")
+            raise e
+        except google.cloud.exceptions.GoogleCloudError as e:
+            self.log.error("An error occurred. Exiting.")
+            raise e
+
+
+class CloudLoggingListSinksOperator(GoogleCloudBaseOperator):
+    """
+    Lists Cloud Logging export sinks in a Google Cloud project.
+
+    :param project_id: Required. The ID of the Google Cloud project to list sinks from.
+    :param page_size: Optional. The maximum number of sinks to return per page. Must be greater than 0.
+        If None, the server will use a default value.
+    :param gcp_conn_id: Optional. The connection ID used to connect to Google Cloud.
+        Defaults to "google_cloud_default".
+    :param impersonation_chain: Optional. Service account or chained list of accounts to impersonate.
+        If a string, the service account must grant the originating account the
+        'Service Account Token Creator' IAM role.
+
+        If a sequence, each account in the chain must grant this role to the next.
+        The first account must grant it to the originating account (templated).
+    """
+
+    template_fields: Sequence[str] = ("project_id", "gcp_conn_id", "impersonation_chain", "page_size")
+
+    def __init__(
+        self,
+        project_id: str,
+        page_size: int | None = None,
+        gcp_conn_id: str = "google_cloud_default",
+        impersonation_chain: str | Sequence[str] | None = None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.page_size = page_size
+        self.gcp_conn_id = gcp_conn_id
+        self.impersonation_chain = impersonation_chain
+
+    def execute(self, context: Context) -> list[dict[str, Any]]:
+        """Execute the operator."""
+        _validate_inputs(self, ["project_id"])
+
+        if self.page_size is not None and self.page_size < 1:
+            raise AirflowException("The page_size for the list sinks request must be greater than zero")
+
+        hook = CloudLoggingHook(gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain)
+
+        try:
+            self.log.info("Listing log sinks in project '%s'", self.project_id)
+
+            sinks = hook.list_sinks(project_id=self.project_id, page_size=self.page_size)
+
+            result = [LogSink.to_dict(sink) for sink in sinks]
+            self.log.info("Found %d log sinks", len(result))
+
+            return result
+
+        except google.cloud.exceptions.GoogleCloudError as e:
+            self.log.error("An error occurred. Exiting.")
+            raise e