kaxil commented on a change in pull request #4354: [AIRFLOW-3446] Add Google Cloud BigTable operators URL: https://github.com/apache/incubator-airflow/pull/4354#discussion_r245115365
########## File path: airflow/contrib/operators/gcp_bigtable_operator.py ########## @@ -0,0 +1,424 @@ +# -*- coding: utf-8 -*- +# +# 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. + +import google.api_core.exceptions + +from airflow import AirflowException +from airflow.models import BaseOperator +from airflow.sensors.base_sensor_operator import BaseSensorOperator +from airflow.contrib.hooks.gcp_bigtable_hook import BigtableHook +from airflow.utils.decorators import apply_defaults +from google.cloud.bigtable_admin_v2 import enums +from google.cloud.bigtable.table import ClusterState + + +class BigtableValidationMixin(object): + """ + Common class for Cloud Bigtable operators for validating required fields. + """ + + REQUIRED_ATTRIBUTES = [] + + def _validate_inputs(self): + for attr_name in self.REQUIRED_ATTRIBUTES: + if not getattr(self, attr_name): + raise AirflowException('Empty parameter: {}'.format(attr_name)) + + +class BigtableInstanceCreateOperator(BaseOperator, BigtableValidationMixin): + """ + Creates a new Cloud Bigtable instance. + If the Cloud Bigtable instance with the given ID exists, the operator does not compare its configuration + and immediately succeeds. No changes are made to the existing instance. + + For more details about instance creation have a look at the reference: + https://googleapis.github.io/google-cloud-python/latest/bigtable/instance.html#google.cloud.bigtable.instance.Instance.create + + :type project_id: str + :param project_id: The ID of the GCP project. + :type instance_id: str + :param instance_id: The ID of the Cloud Bigtable instance to create. + :type main_cluster_id: str + :param main_cluster_id: The ID for main cluster for the new instance. + :type main_cluster_zone: str + :param main_cluster_zone: The zone for main cluster + See https://cloud.google.com/bigtable/docs/locations for more details. + :type replica_cluster_id: str + :param replica_cluster_id: (optional) The ID for replica cluster for the new instance. + :type replica_cluster_zone: str + :param replica_cluster_zone: (optional) The zone for replica cluster. + :type instance_type: IntEnum + :param instance_type: (optional) The type of the instance. + :type instance_display_name: str + :param instance_display_name: (optional) Human-readable name of the instance. Defaults to ``instance_id``. + :type instance_labels: dict + :param instance_labels: (optional) Dictionary of labels to associate with the instance. + :type cluster_nodes: int + :param cluster_nodes: (optional) Number of nodes for cluster. + :type cluster_storage_type: IntEnum + :param cluster_storage_type: (optional) The type of storage. + :type timeout: int + :param timeout: (optional) timeout (in seconds) for instance creation. + If None is not specified, Operator will wait indefinitely. + """ + + REQUIRED_ATTRIBUTES = ('project_id', 'instance_id', 'main_cluster_id', 'main_cluster_zone') + template_fields = ['project_id', 'instance_id', 'main_cluster_id', 'main_cluster_zone'] + + @apply_defaults + def __init__(self, + project_id, + instance_id, + main_cluster_id, + main_cluster_zone, + replica_cluster_id=None, + replica_cluster_zone=None, + instance_display_name=None, + instance_type=None, + instance_labels=None, + cluster_nodes=None, + cluster_storage_type=None, + timeout=None, + *args, **kwargs): + self.project_id = project_id + self.instance_id = instance_id + self.main_cluster_id = main_cluster_id + self.main_cluster_zone = main_cluster_zone + self.replica_cluster_id = replica_cluster_id + self.replica_cluster_zone = replica_cluster_zone + self.instance_display_name = instance_display_name + self.instance_type = instance_type + self.instance_labels = instance_labels + self.cluster_nodes = cluster_nodes + self.cluster_storage_type = cluster_storage_type + self.timeout = timeout + self._validate_inputs() + self.hook = BigtableHook() + super(BigtableInstanceCreateOperator, self).__init__(*args, **kwargs) + + def execute(self, context): + instance = self.hook.get_instance(self.project_id, self.instance_id) + if instance: + # Based on Instance.__eq__ instance with the same ID and client is considered as equal. + self.log.info( + "The instance '%s' already exists in this project. Consider it as created", + self.instance_id + ) + return + try: + self.hook.create_instance( + self.project_id, + self.instance_id, + self.main_cluster_id, + self.main_cluster_zone, + self.replica_cluster_id, + self.replica_cluster_zone, + self.instance_display_name, + self.instance_type, + self.instance_labels, + self.cluster_nodes, + self.cluster_storage_type, + self.timeout, + ) + except google.api_core.exceptions.GoogleAPICallError as e: + self.log.error('An error occurred. Exiting.') + raise e + + +class BigtableInstanceDeleteOperator(BaseOperator, BigtableValidationMixin): + """ + Deletes the Cloud Bigtable instance, including its clusters and all related tables. + + For more details about deleting instance have a look at the reference: + https://googleapis.github.io/google-cloud-python/latest/bigtable/instance.html#google.cloud.bigtable.instance.Instance.delete + + :type project_id: str + :param project_id: The ID of the GCP project. + :type instance_id: str + :param instance_id: The ID of the Cloud Bigtable instance to delete. + """ + REQUIRED_ATTRIBUTES = ('project_id', 'instance_id') Review comment: This is a new operator, hence I would still feel as it is just a matter of removing an attribute. However, I don't have a strong feeling about it - so it can wait for the next PR ---------------------------------------------------------------- This is an automated message from the Apache Git Service. To respond to the message, please log on GitHub and use the URL above to go to the specific comment. For queries about this service, please contact Infrastructure at: us...@infra.apache.org With regards, Apache Git Services
