turbaszek commented on a change in pull request #21251: URL: https://github.com/apache/airflow/pull/21251#discussion_r800086573
########## File path: airflow/providers/google/cloud/operators/cloud_composer.py ########## @@ -0,0 +1,696 @@ +# +# 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 datetime import datetime +from typing import TYPE_CHECKING, Dict, Optional, Sequence, Tuple, Union + +from google.api_core.exceptions import AlreadyExists +from google.api_core.retry import Retry +from google.cloud.orchestration.airflow.service_v1 import ImageVersion +from google.cloud.orchestration.airflow.service_v1.types import Environment +from google.protobuf.field_mask_pb2 import FieldMask + +from airflow import AirflowException +from airflow.models import BaseOperator, BaseOperatorLink, XCom +from airflow.providers.google.cloud.hooks.cloud_composer import CloudComposerHook +from airflow.providers.google.cloud.triggers.cloud_composer import CloudComposerExecutionTrigger +from airflow.providers.google.common.consts import GOOGLE_DEFAULT_DEFERRABLE_METHOD_NAME + +if TYPE_CHECKING: + from airflow.utils.context import Context + +CLOUD_COMPOSER_BASE_LINK = "https://console.cloud.google.com/composer/environments"; +CLOUD_COMPOSER_DETAILS_LINK = ( + CLOUD_COMPOSER_BASE_LINK + "/detail/{region}/{environment_id}/monitoring?project={project_id}" +) +CLOUD_COMPOSER_ENVIRONMENTS_LINK = CLOUD_COMPOSER_BASE_LINK + '?project={project_id}' + + +class CloudComposerEnvironmentLink(BaseOperatorLink): + """Helper class for constructing Cloud Composer Environment Link""" + + name = "Cloud Composer Environment" + + def get_link(self, operator: BaseOperator, dttm: datetime) -> str: + conf = XCom.get_one( + key="composer_conf", dag_id=operator.dag.dag_id, task_id=operator.task_id, execution_date=dttm + ) + return ( + CLOUD_COMPOSER_DETAILS_LINK.format( + project_id=conf["project_id"], region=conf['region'], environment_id=conf['environment_id'] + ) + if conf + else "" + ) + + +class CloudComposerEnvironmentsLink(BaseOperatorLink): + """Helper class for constructing Cloud Composer Environment Link""" + + name = "Cloud Composer Environment List" + + def get_link(self, operator: BaseOperator, dttm: datetime) -> str: + conf = XCom.get_one( + key="composer_conf", dag_id=operator.dag.dag_id, task_id=operator.task_id, execution_date=dttm + ) + return CLOUD_COMPOSER_ENVIRONMENTS_LINK.format(project_id=conf["project_id"]) if conf else "" + + +class CloudComposerCreateEnvironmentOperator(BaseOperator): + """ + Create a new 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 environment: The environment to create. + :param gcp_conn_id: + :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). + :param delegate_to: The account to impersonate using domain-wide delegation of authority, + if any. For this to work, the service account making the request must have + domain-wide delegation enabled. + :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. + :param deferrable: Run operator in the deferrable mode + :param pooling_period_seconds: Optional: Control the rate of the poll for the result of deferrable run. + By default the trigger will poll every 30 seconds. + """ + + template_fields = ( + 'project_id', + 'region', + 'environment_id', + 'environment', + 'impersonation_chain', + ) + + operator_extra_links = (CloudComposerEnvironmentLink(),) + + def __init__( + self, + *, + project_id: str, + region: str, + environment_id: str, + environment: Union[Environment, Dict], + gcp_conn_id: str = "google_cloud_default", + impersonation_chain: Optional[Union[str, Sequence[str]]] = None, + delegate_to: Optional[str] = None, + retry: Optional[Retry] = None, + timeout: Optional[float] = None, + metadata: Sequence[Tuple[str, str]] = (), + deferrable: bool = False, + pooling_period_seconds: int = 30, + **kwargs, + ) -> None: + super().__init__(**kwargs) + self.project_id = project_id + self.region = region + self.environment_id = environment_id + self.environment = environment + self.retry = retry + self.timeout = timeout + self.metadata = metadata + self.gcp_conn_id = gcp_conn_id + self.impersonation_chain = impersonation_chain + self.delegate_to = delegate_to + self.deferrable = deferrable + self.pooling_period_seconds = pooling_period_seconds + + def execute(self, context: 'Context'): + hook = CloudComposerHook( + gcp_conn_id=self.gcp_conn_id, + impersonation_chain=self.impersonation_chain, + delegate_to=self.delegate_to, + ) + + name = hook.get_environment_name(self.project_id, self.region, self.environment_id) + if isinstance(self.environment, Environment): + self.environment.name = name + else: + self.environment["name"] = name + + self.xcom_push( + context, + key='composer_conf', + value={ + "project_id": self.project_id, + "region": self.region, + "environment_id": self.environment_id, + }, + ) Review comment: Same comment as here https://github.com/apache/airflow/pull/21267#discussion_r800086099, the minimal change I would expect is to extract the key value as class link attribute so we cane reference it instead of hardcoding it in two places. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
