[GitHub] [airflow] turbaszek commented on a change in pull request #10340: Add GCP Bigtable Update Instance Hook/Operator

[GitBox <gi...@apache.org>]

turbaszek commented on a change in pull request #10340:
URL: https://github.com/apache/airflow/pull/10340#discussion_r471015015



##########
File path: airflow/providers/google/cloud/operators/bigtable.py
##########
@@ -158,6 +159,81 @@ def execute(self, context):
             raise e
 
 
+class BigtableUpdateInstanceOperator(BaseOperator, BigtableValidationMixin):
+    """
+    Updates an existing Cloud Bigtable instance.
+
+    For more details about instance creation have a look at the reference:
+    https://googleapis.dev/python/bigtable/latest/instance.html#google.cloud.bigtable.instance.Instance.update
+
+    .. seealso::
+        For more information on how to use this operator, take a look at the guide:
+        :ref:`howto/operator:BigtableUpdateInstanceOperator`
+
+    :type instance_id: str
+    :param instance_id: The ID of the Cloud Bigtable instance to update.
+    :type project_id: str
+    :param project_id: Optional, the ID of the GCP project. If set to None or missing,
+            the default project_id from the GCP connection is used.
+    :type instance_display_name: str
+    :param instance_display_name: (optional) Human-readable name of the instance.
+    :type instance_type: enums.Instance.Type or enum.IntEnum
+    :param instance_type: (optional) The type of the instance.
+    :type instance_labels: dict
+    :param instance_labels: (optional) Dictionary of labels to associate
+        with the instance.
+    :type timeout: int
+    :param timeout: (optional) timeout (in seconds) for instance update.
+                    If None is not specified, Operator will wait indefinitely.
+    :param gcp_conn_id: The connection ID to use to connect to Google Cloud Platform.
+    :type gcp_conn_id: str
+    """
+
+    REQUIRED_ATTRIBUTES: Iterable[str] = ['instance_id']
+    template_fields: Iterable[str] = ['project_id', 'instance_id']
+
+    @apply_defaults
+    def __init__(self, *,
+                 instance_id: str,
+                 project_id: Optional[str] = None,
+                 instance_display_name: Optional[str] = None,
+                 instance_type: Optional[Union[enums.Instance.Type, enum.IntEnum]] = None,
+                 instance_labels: Optional[Dict] = None,
+                 timeout: Optional[float] = None,
+                 gcp_conn_id: str = 'google_cloud_default',
+                 **kwargs) -> None:
+        self.project_id = project_id
+        self.instance_id = instance_id
+        self.instance_display_name = instance_display_name
+        self.instance_type = instance_type
+        self.instance_labels = instance_labels
+        self.timeout = timeout
+        self._validate_inputs()
+        self.gcp_conn_id = gcp_conn_id
+        super().__init__(**kwargs)
+
+    def execute(self, context):
+        hook = BigtableHook(gcp_conn_id=self.gcp_conn_id)
+        instance = hook.get_instance(project_id=self.project_id,
+                                     instance_id=self.instance_id)
+        if not instance:
+            raise AirflowException("Dependency: instance '{}' does not exist.".format(
+                self.instance_id))

Review comment:
       ```suggestion
               raise AirflowException(f"Dependency: instance '{self.instance_id}' does not exist.")
   ```

##########
File path: airflow/providers/google/cloud/hooks/bigtable.py
##########
@@ -184,6 +185,51 @@ def create_instance(
         operation.result(timeout)
         return instance
 
+    @GoogleBaseHook.fallback_to_default_project_id
+    def update_instance(
+        self,
+        instance_id: str,
+        project_id: str,
+        instance_display_name: Optional[str] = None,
+        instance_type: Optional[Union[enums.Instance.Type, enum.IntEnum]] = None,
+        instance_labels: Optional[Dict] = None,
+        timeout: Optional[float] = None
+    ) -> Instance:
+        """
+        Update an existing instance.
+
+        :type instance_id: str
+        :param instance_id: The ID for the existing instance.
+        :type project_id: str
+        :param project_id: Optional, Google Cloud Platform project ID where the
+            BigTable exists. If set to None or missing,
+            the default project_id from the GCP connection is used.
+        :type instance_display_name: str
+        :param instance_display_name: (optional) Human-readable name of the instance.
+        :type instance_type: enums.Instance.Type or enum.IntEnum
+        :param instance_type: (optional) The type of the instance.
+        :type instance_labels: dict
+        :param instance_labels: (optional) Dictionary of labels to associate with the
+            instance.
+        :type timeout: int
+        :param timeout: (optional) timeout (in seconds) for instance update.
+            If None is not specified, Operator will wait indefinitely.
+        """
+        instance_type = enums.Instance.Type(instance_type)
+
+        instance = Instance(
+            instance_id=instance_id,
+            client=self._get_client(project_id=project_id),
+            display_name=instance_display_name,
+            instance_type=instance_type,
+            labels=instance_labels,
+        )
+
+        operation = instance.update()
+        operation.result(timeout)

Review comment:
       Nice! I didn't know that timeout can be passed like that 👍 




----------------------------------------------------------------
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.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org