Repository: apache/airflow
File: providers/google/src/airflow/providers/google/cloud/hooks/bigtable.py (Python)

Code before:

        table = instance.table(table_id=table_id)
        try:
            table.delete()
        except google.api_core.exceptions.NotFound:
            self.log.info("The table '%s' no longer exists. Consider it as deleted", table_id)

    @staticmethod
    def update_cluster(instance: Instance, instance_id: str, cluster_id: str, nodes: int) -> None:
        """
        Update number of nodes in the specified Cloud Bigtable cluster.

        Raises google.api_core.exceptions.NotFound if the cluster does not exist.

        :param instance: The Cloud Bigtable instance that owns the cluster.
        :param cluster_id: The ID of the cluster.
        :param nodes: The desired number of nodes.
        """
        cluster = Cluster(cluster_id, instance)
        # "reload" is required to set location_id attribute on cluster.
        try:
            cluster.reload()
        except google.api_core.exceptions.NotFound:
            raise AirflowException(
                f"Dependency: cluster '{cluster_id}' does not exist for instance '{instance_id}'."
            )
        cluster.serve_nodes = nodes
        cluster.update()

    @staticmethod
    def get_column_families_for_table(instance: Instance, table_id: str) -> dict[str, ColumnFamily]:
        """
        Fetch Column Families for the specified table in Cloud Bigtable.

        :param instance: The Cloud Bigtable instance that owns the table.
        :param table_id: The ID of the table in Cloud Bigtable to fetch Column Families
            from.
        """
        table = Table(table_id, instance)
        return table.list_column_families()

    @staticmethod
    def get_cluster_states_for_table(instance: Instance, table_id: str) -> dict[str, ClusterState]:
        """
        Fetch Cluster States for the specified table in Cloud Bigtable.

        Raises google.api_core.exceptions.NotFound if the table does not exist.

        :param instance: The Cloud Bigtable instance that owns the table.
        :param table_id: The ID of the table in Cloud Bigtable to fetch Cluster States
            from.

Review comment (suggestion):
1. I'll waive the usage of `AirflowException` here because it was relocated from the operator, but eventually we'll need to deprecate it.
2. If you relocate the logic - it should be tested (you could use the deleted tests from the operator as a reference)

Code after:

        table = instance.table(table_id=table_id)
        try:
            table.delete()
        except google.api_core.exceptions.NotFound:
            self.log.info("The table '%s' no longer exists. Consider it as deleted", table_id)

    @staticmethod
    def update_cluster(instance: Instance, cluster_id: str, nodes: int) -> None:
        """
        Update number of nodes in the specified Cloud Bigtable cluster.

        If the cluster does not exist, raises AirflowException.

        :param instance: The Cloud Bigtable instance that owns the cluster.
        :param cluster_id: The ID of the cluster.
        :param nodes: The desired number of nodes.
        """
        cluster = Cluster(cluster_id, instance)
        # "reload" is required to set location_id attribute on cluster.
        try:
            cluster.reload()
        except google.api_core.exceptions.NotFound:
            raise AirflowException(
                f"Dependency: cluster '{cluster_id}' does not exist for instance '{instance.instance_id}'."
            )
        cluster.serve_nodes = nodes
        cluster.update()

    @staticmethod
    def get_column_families_for_table(instance: Instance, table_id: str) -> dict[str, ColumnFamily]:
        """
        Fetch Column Families for the specified table in Cloud Bigtable.

        :param instance: The Cloud Bigtable instance that owns the table.
        :param table_id: The ID of the table in Cloud Bigtable to fetch Column Families
            from.
        """
        table = Table(table_id, instance)
        return table.list_column_families()

    @staticmethod
    def get_cluster_states_for_table(instance: Instance, table_id: str) -> dict[str, ClusterState]:
        """
        Fetch Cluster States for the specified table in Cloud Bigtable.

        Raises google.api_core.exceptions.NotFound if the table does not exist.

        :param instance: The Cloud Bigtable instance that owns the table.
        :param table_id: The ID of the table in Cloud Bigtable to fetch Cluster States
            from.