Python BigtableHook.update_cluster方法代码示例

# 需要导入模块: from airflow.contrib.hooks.gcp_bigtable_hook import BigtableHook [as 别名]
# 或者: from airflow.contrib.hooks.gcp_bigtable_hook.BigtableHook import update_cluster [as 别名]
class BigtableClusterUpdateOperator(BaseOperator, BigtableValidationMixin):
    """
    Updates a Cloud Bigtable cluster.

    For more details about updating a Cloud Bigtable cluster,
    have a look at the reference:
    https://googleapis.github.io/google-cloud-python/latest/bigtable/cluster.html#google.cloud.bigtable.cluster.Cluster.update

    .. seealso::
        For more information on how to use this operator, take a look at the guide:
        :ref:`howto/operator:BigtableClusterUpdateOperator`

    :type instance_id: str
    :param instance_id: The ID of the Cloud Bigtable instance.
    :type cluster_id: str
    :param cluster_id: The ID of the Cloud Bigtable cluster to update.
    :type nodes: int
    :param nodes: The desired number of nodes for the Cloud Bigtable cluster.
    :type project_id: str
    :param project_id: Optional, the ID of the GCP project.
    """
    REQUIRED_ATTRIBUTES = ('instance_id', 'cluster_id', 'nodes')
    template_fields = ['project_id', 'instance_id', 'cluster_id', 'nodes']

    @apply_defaults
    def __init__(self,
                 instance_id,
                 cluster_id,
                 nodes,
                 project_id=None,
                 *args, **kwargs):
        self.project_id = project_id
        self.instance_id = instance_id
        self.cluster_id = cluster_id
        self.nodes = nodes
        self._validate_inputs()
        self.hook = BigtableHook()
        super(BigtableClusterUpdateOperator, self).__init__(*args, **kwargs)

    def execute(self, context):
        instance = self.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))

        try:
            self.hook.update_cluster(
                instance=instance,
                cluster_id=self.cluster_id,
                nodes=self.nodes
            )
        except google.api_core.exceptions.NotFound:
            raise AirflowException(
                "Dependency: cluster '{}' does not exist for instance '{}'.".
                format(self.cluster_id, self.instance_id))
        except google.api_core.exceptions.GoogleAPICallError as e:
            self.log.error('An error occurred. Exiting.')
            raise e

# 需要导入模块: from airflow.contrib.hooks.gcp_bigtable_hook import BigtableHook [as 别名]
# 或者: from airflow.contrib.hooks.gcp_bigtable_hook.BigtableHook import update_cluster [as 别名]

#.........这里部分代码省略.........
        instance_create.return_value = operation
        res = self.bigtable_hook_default_project_id.create_instance(
            project_id='new-project',
            instance_id=CBT_INSTANCE,
            main_cluster_id=CBT_CLUSTER,
            main_cluster_zone=CBT_ZONE)
        get_client.assert_called_once_with(project_id='new-project')
        instance_create.assert_called_once_with(clusters=mock.ANY)
        self.assertEqual(res.instance_id, 'instance')

    @mock.patch('airflow.contrib.hooks.gcp_bigtable_hook.BigtableHook._get_client')
    def test_delete_table(self, get_client):
        instance_method = get_client.return_value.instance
        instance_exists_method = instance_method.return_value.exists
        table_delete_method = instance_method.return_value.table.return_value.delete
        instance_exists_method.return_value = True
        self.bigtable_hook_default_project_id.delete_table(
            instance_id=CBT_INSTANCE,
            table_id=CBT_TABLE)
        get_client.assert_called_once_with(project_id='example-project')
        instance_exists_method.assert_called_once_with()
        table_delete_method.assert_called_once_with()

    @mock.patch('airflow.contrib.hooks.gcp_bigtable_hook.BigtableHook._get_client')
    def test_delete_table_overridden_project_id(self, get_client):
        instance_method = get_client.return_value.instance
        instance_exists_method = instance_method.return_value.exists
        table_delete_method = instance_method.return_value.table.return_value.delete
        instance_exists_method.return_value = True
        self.bigtable_hook_default_project_id.delete_table(
            project_id='new-project',
            instance_id=CBT_INSTANCE,
            table_id=CBT_TABLE)
        get_client.assert_called_once_with(project_id='new-project')
        instance_exists_method.assert_called_once_with()
        table_delete_method.assert_called_once_with()

    @mock.patch('google.cloud.bigtable.table.Table.create')
    @mock.patch('airflow.contrib.hooks.gcp_bigtable_hook.BigtableHook._get_client')
    def test_create_table(self, get_client, create):
        instance_method = get_client.return_value.instance
        instance_exists_method = instance_method.return_value.exists
        instance_exists_method.return_value = True
        client = mock.Mock(Client)
        instance = google.cloud.bigtable.instance.Instance(
            instance_id=CBT_INSTANCE,
            client=client)
        self.bigtable_hook_default_project_id.create_table(
            instance=instance,
            table_id=CBT_TABLE)
        get_client.assert_not_called()
        create.assert_called_once_with([], {})

    @mock.patch('google.cloud.bigtable.cluster.Cluster.update')
    @mock.patch('airflow.contrib.hooks.gcp_bigtable_hook.BigtableHook._get_client')
    def test_update_cluster(self, get_client, update):
        instance_method = get_client.return_value.instance
        instance_exists_method = instance_method.return_value.exists
        instance_exists_method.return_value = True
        client = mock.Mock(Client)
        instance = google.cloud.bigtable.instance.Instance(
            instance_id=CBT_INSTANCE,
            client=client)
        self.bigtable_hook_default_project_id.update_cluster(
            instance=instance,
            cluster_id=CBT_CLUSTER,
            nodes=4)
        get_client.assert_not_called()
        update.assert_called_once_with()

    @mock.patch('google.cloud.bigtable.table.Table.list_column_families')
    @mock.patch('airflow.contrib.hooks.gcp_bigtable_hook.BigtableHook._get_client')
    def test_list_column_families(self, get_client, list_column_families):
        instance_method = get_client.return_value.instance
        instance_exists_method = instance_method.return_value.exists
        instance_exists_method.return_value = True
        client = mock.Mock(Client)
        get_client.return_value = client
        instance = google.cloud.bigtable.instance.Instance(
            instance_id=CBT_INSTANCE,
            client=client)
        self.bigtable_hook_default_project_id.get_column_families_for_table(
            instance=instance, table_id=CBT_TABLE)
        get_client.assert_not_called()
        list_column_families.assert_called_once_with()

    @mock.patch('google.cloud.bigtable.table.Table.get_cluster_states')
    @mock.patch('airflow.contrib.hooks.gcp_bigtable_hook.BigtableHook._get_client')
    def test_get_cluster_states(self, get_client, get_cluster_states):
        instance_method = get_client.return_value.instance
        instance_exists_method = instance_method.return_value.exists
        instance_exists_method.return_value = True
        client = mock.Mock(Client)
        instance = google.cloud.bigtable.instance.Instance(
            instance_id=CBT_INSTANCE,
            client=client)
        self.bigtable_hook_default_project_id.get_cluster_states_for_table(
            instance=instance, table_id=CBT_TABLE)
        get_client.assert_not_called()
        get_cluster_states.assert_called_once_with()