From 93df82998cc0218cbc4a1bc2ab41a48b7478758d Mon Sep 17 00:00:00 2001 From: Chris Rossi Date: Mon, 19 Apr 2021 16:03:49 -0400 Subject: [PATCH] feat: customer managed keys (CMEK) (#249) * feat: customer managed keys (CMEK) Implement customer managed keys (CMEK) feature. WIP. DO NOT MERGE. * Wrap Status. * Wrapper for Status, reorganize to avoid circular imports. * Blacken. * Make system tests in charge of their own key. * Consolidate system tests. Get KMS_KEY_NAME from user's environment. * Fix test. * Lint. * Put system tests where nox is expecting to find them. * Test backup with CMEK. * Differentiate instance and cluster names for cmek test, so tests aren't stepping on each other's toes. Remove bogus backup with cmek test. * rename `encryption.py` to `encryption_info.py` * make sure `kms_key_name` is set to `None` if `encryption_info` is not PB. * Fix typo. Use more realistic looking test strings. --- google/cloud/bigtable/backup.py | 21 +- google/cloud/bigtable/cluster.py | 28 +++ google/cloud/bigtable/encryption_info.py | 64 ++++++ google/cloud/bigtable/enums.py | 30 +++ google/cloud/bigtable/error.py | 64 ++++++ google/cloud/bigtable/instance.py | 24 ++- google/cloud/bigtable/table.py | 28 +++ tests/system.py | 236 ++++++++++++++++++++++- tests/unit/test_backup.py | 32 +++ tests/unit/test_cluster.py | 99 ++++++++++ tests/unit/test_table.py | 107 ++++++++++ 11 files changed, 724 insertions(+), 9 deletions(-) create mode 100644 google/cloud/bigtable/encryption_info.py create mode 100644 google/cloud/bigtable/error.py diff --git a/google/cloud/bigtable/backup.py b/google/cloud/bigtable/backup.py index 6dead1f74..3666b7132 100644 --- a/google/cloud/bigtable/backup.py +++ b/google/cloud/bigtable/backup.py @@ -19,6 +19,7 @@ from google.cloud._helpers import _datetime_to_pb_timestamp from google.cloud.bigtable_admin_v2 import BigtableTableAdminClient from google.cloud.bigtable_admin_v2.types import table +from google.cloud.bigtable.encryption_info import EncryptionInfo from google.cloud.bigtable.policy import Policy from google.cloud.exceptions import NotFound from google.protobuf import field_mask_pb2 @@ -67,13 +68,20 @@ class Backup(object): """ def __init__( - self, backup_id, instance, cluster_id=None, table_id=None, expire_time=None + self, + backup_id, + instance, + cluster_id=None, + table_id=None, + expire_time=None, + encryption_info=None, ): self.backup_id = backup_id self._instance = instance self._cluster = cluster_id self.table_id = table_id self._expire_time = expire_time + self._encryption_info = encryption_info self._parent = None self._source_table = None @@ -176,6 +184,15 @@ def expire_time(self): def expire_time(self, new_expire_time): self._expire_time = new_expire_time + @property + def encryption_info(self): + """Encryption info for this Backup. + + :rtype: :class:`google.cloud.bigtable.encryption.EncryptionInfo` + :returns: The encryption information for this backup. + """ + return self._encryption_info + @property def start_time(self): """The time this Backup was started. @@ -255,6 +272,7 @@ def from_pb(cls, backup_pb, instance): table_id = match.group("table_id") if match else None expire_time = backup_pb._pb.expire_time + encryption_info = EncryptionInfo._from_pb(backup_pb.encryption_info) backup = cls( backup_id, @@ -262,6 +280,7 @@ def from_pb(cls, backup_pb, instance): cluster_id=cluster_id, table_id=table_id, expire_time=expire_time, + encryption_info=encryption_info, ) backup._start_time = backup_pb._pb.start_time backup._end_time = backup_pb._pb.end_time diff --git a/google/cloud/bigtable/cluster.py b/google/cloud/bigtable/cluster.py index 5c4c355ff..f3e79c6c2 100644 --- a/google/cloud/bigtable/cluster.py +++ b/google/cloud/bigtable/cluster.py @@ -63,6 +63,19 @@ class Cluster(object): Defaults to :data:`google.cloud.bigtable.enums.StorageType.UNSPECIFIED`. + :type kms_key_name: str + :param kms_key_name: (Optional, Creation Only) The name of the KMS customer managed + encryption key (CMEK) to use for at-rest encryption of data in + this cluster. If omitted, Google's default encryption will be + used. If specified, the requirements for this key are: + + 1) The Cloud Bigtable service account associated with the + project that contains the cluster must be granted the + ``cloudkms.cryptoKeyEncrypterDecrypter`` role on the CMEK. + 2) Only regional keys can be used and the region of the CMEK + key must match the region of the cluster. + 3) All clusters within an instance must use the same CMEK key. + :type _state: int :param _state: (`OutputOnly`) The current state of the cluster. @@ -81,6 +94,7 @@ def __init__( location_id=None, serve_nodes=None, default_storage_type=None, + kms_key_name=None, _state=None, ): self.cluster_id = cluster_id @@ -88,6 +102,7 @@ def __init__( self.location_id = location_id self.serve_nodes = serve_nodes self.default_storage_type = default_storage_type + self._kms_key_name = kms_key_name self._state = _state @classmethod @@ -145,6 +160,10 @@ def _update_from_pb(self, cluster_pb): self.location_id = cluster_pb.location.split("/")[-1] self.serve_nodes = cluster_pb.serve_nodes self.default_storage_type = cluster_pb.default_storage_type + if cluster_pb.encryption_config: + self._kms_key_name = cluster_pb.encryption_config.kms_key_name + else: + self._kms_key_name = None self._state = cluster_pb.state @property @@ -187,6 +206,11 @@ def state(self): """ return self._state + @property + def kms_key_name(self): + """str: Customer managed encryption key for the cluster.""" + return self._kms_key_name + def __eq__(self, other): if not isinstance(other, self.__class__): return NotImplemented @@ -356,4 +380,8 @@ def _to_pb(self): serve_nodes=self.serve_nodes, default_storage_type=self.default_storage_type, ) + if self._kms_key_name: + cluster_pb.encryption_config = instance.Cluster.EncryptionConfig( + kms_key_name=self._kms_key_name, + ) return cluster_pb diff --git a/google/cloud/bigtable/encryption_info.py b/google/cloud/bigtable/encryption_info.py new file mode 100644 index 000000000..1757297bc --- /dev/null +++ b/google/cloud/bigtable/encryption_info.py @@ -0,0 +1,64 @@ +# Copyright 2021 Google LLC +# +# Licensed 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. + +"""Class for encryption info for tables and backups.""" + +from google.cloud.bigtable.error import Status + + +class EncryptionInfo: + """Encryption information for a given resource. + + If this resource is protected with customer managed encryption, the in-use Google + Cloud Key Management Service (KMS) key versions will be specified along with their + status. + + :type encryption_type: int + :param encryption_type: See :class:`enums.EncryptionInfo.EncryptionType` + + :type encryption_status: google.cloud.bigtable.encryption.Status + :param encryption_status: The encryption status. + + :type kms_key_version: str + :param kms_key_version: The key version used for encryption. + """ + + @classmethod + def _from_pb(cls, info_pb): + return cls( + info_pb.encryption_type, + Status(info_pb.encryption_status), + info_pb.kms_key_version, + ) + + def __init__(self, encryption_type, encryption_status, kms_key_version): + self.encryption_type = encryption_type + self.encryption_status = encryption_status + self.kms_key_version = kms_key_version + + def __eq__(self, other): + if self is other: + return True + + if not isinstance(other, type(self)): + return NotImplemented + + return ( + self.encryption_type == other.encryption_type + and self.encryption_status == other.encryption_status + and self.kms_key_version == other.kms_key_version + ) + + def __ne__(self, other): + return not self == other diff --git a/google/cloud/bigtable/enums.py b/google/cloud/bigtable/enums.py index 50c7f2e60..327b2f828 100644 --- a/google/cloud/bigtable/enums.py +++ b/google/cloud/bigtable/enums.py @@ -156,6 +156,7 @@ class View(object): NAME_ONLY = table.Table.View.NAME_ONLY SCHEMA_VIEW = table.Table.View.SCHEMA_VIEW REPLICATION_VIEW = table.Table.View.REPLICATION_VIEW + ENCRYPTION_VIEW = table.Table.View.ENCRYPTION_VIEW FULL = table.Table.View.FULL class ReplicationState(object): @@ -191,3 +192,32 @@ class ReplicationState(object): table.Table.ClusterState.ReplicationState.UNPLANNED_MAINTENANCE ) READY = table.Table.ClusterState.ReplicationState.READY + + +class EncryptionInfo: + class EncryptionType: + """Possible encryption types for a resource. + + Attributes: + ENCRYPTION_TYPE_UNSPECIFIED (int): Encryption type was not specified, though + data at rest remains encrypted. + GOOGLE_DEFAULT_ENCRYPTION (int): The data backing this resource is encrypted + at rest with a key that is fully managed by Google. No key version or + status will be populated. This is the default state. + CUSTOMER_MANAGED_ENCRYPTION (int): The data backing this resource is + encrypted at rest with a key that is managed by the customer. The in-use + version of the key and its status are populated for CMEK-protected + tables. CMEK-protected backups are pinned to the key version that was in + use at the time the backup was taken. This key version is populated but + its status is not tracked and is reported as `UNKNOWN`. + """ + + ENCRYPTION_TYPE_UNSPECIFIED = ( + table.EncryptionInfo.EncryptionType.ENCRYPTION_TYPE_UNSPECIFIED + ) + GOOGLE_DEFAULT_ENCRYPTION = ( + table.EncryptionInfo.EncryptionType.GOOGLE_DEFAULT_ENCRYPTION + ) + CUSTOMER_MANAGED_ENCRYPTION = ( + table.EncryptionInfo.EncryptionType.CUSTOMER_MANAGED_ENCRYPTION + ) diff --git a/google/cloud/bigtable/error.py b/google/cloud/bigtable/error.py new file mode 100644 index 000000000..261cfc2c3 --- /dev/null +++ b/google/cloud/bigtable/error.py @@ -0,0 +1,64 @@ +# Copyright 2021 Google LLC +# +# Licensed 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. + +"""Class for error status.""" + + +class Status: + """A status, comprising a code and a message. + + See: `Cloud APIs Errors `_ + + This is a thin wrapper for ``google.rpc.status_pb2.Status``. + + :type status_pb: google.rpc.status_pb2.Status + :param status_pb: The status protocol buffer. + """ + + def __init__(self, status_pb): + self.status_pb = status_pb + + @property + def code(self): + """The status code. + + Values are defined in ``google.rpc.code_pb2.Code``. + + See: `google.rpc.Code + `_ + + :rtype: int + :returns: The status code. + """ + return self.status_pb.code + + @property + def message(self): + """A human readable status message. + + :rypte: str + :returns: The status message. + """ + return self.status_pb.message + + def __repr__(self): + return repr(self.status_pb) + + def __eq__(self, other): + if isinstance(other, type(self)): + return self.status_pb == other.status_pb + return NotImplemented + + def __ne__(self, other): + return not self == other diff --git a/google/cloud/bigtable/instance.py b/google/cloud/bigtable/instance.py index d2fb5db07..138d3bfc1 100644 --- a/google/cloud/bigtable/instance.py +++ b/google/cloud/bigtable/instance.py @@ -540,7 +540,12 @@ def test_iam_permissions(self, permissions): return list(resp.permissions) def cluster( - self, cluster_id, location_id=None, serve_nodes=None, default_storage_type=None + self, + cluster_id, + location_id=None, + serve_nodes=None, + default_storage_type=None, + kms_key_name=None, ): """Factory to create a cluster associated with this instance. @@ -576,6 +581,22 @@ def cluster( :rtype: :class:`~google.cloud.bigtable.instance.Cluster` :returns: a cluster owned by this instance. + + :type kms_key_name: str + :param kms_key_name: (Optional, Creation Only) The name of the KMS customer + managed encryption key (CMEK) to use for at-rest encryption + of data in this cluster. If omitted, Google's default + encryption will be used. If specified, the requirements for + this key are: + + 1) The Cloud Bigtable service account associated with the + project that contains the cluster must be granted the + ``cloudkms.cryptoKeyEncrypterDecrypter`` role on the + CMEK. + 2) Only regional keys can be used and the region of the + CMEK key must match the region of the cluster. + 3) All clusters within an instance must use the same CMEK + key. """ return Cluster( cluster_id, @@ -583,6 +604,7 @@ def cluster( location_id=location_id, serve_nodes=serve_nodes, default_storage_type=default_storage_type, + kms_key_name=kms_key_name, ) def list_clusters(self): diff --git a/google/cloud/bigtable/table.py b/google/cloud/bigtable/table.py index c2d114362..95fb55c50 100644 --- a/google/cloud/bigtable/table.py +++ b/google/cloud/bigtable/table.py @@ -28,6 +28,7 @@ from google.cloud.bigtable.column_family import ColumnFamily from google.cloud.bigtable.batcher import MutationsBatcher from google.cloud.bigtable.batcher import FLUSH_COUNT, MAX_ROW_BYTES +from google.cloud.bigtable.encryption_info import EncryptionInfo from google.cloud.bigtable.policy import Policy from google.cloud.bigtable.row import AppendRow from google.cloud.bigtable.row import ConditionalRow @@ -484,6 +485,33 @@ def get_cluster_states(self): for cluster_id, value_pb in table_pb.cluster_states.items() } + def get_encryption_info(self): + """List the encryption info for each cluster owned by this table. + + Gets the current encryption info for the table across all of the clusters. The + returned dict will be keyed by cluster id and contain a status for all of the + keys in use. + + :rtype: dict + :returns: Dictionary of encryption info for this table. Keys are cluster ids and + values are tuples of :class:`google.cloud.bigtable.encryption.EncryptionInfo` instances. + """ + ENCRYPTION_VIEW = enums.Table.View.ENCRYPTION_VIEW + table_client = self._instance._client.table_admin_client + table_pb = table_client.get_table( + request={"name": self.name, "view": ENCRYPTION_VIEW} + ) + + return { + cluster_id: tuple( + ( + EncryptionInfo._from_pb(info_pb) + for info_pb in value_pb.encryption_info + ) + ) + for cluster_id, value_pb in table_pb.cluster_states.items() + } + def read_row(self, row_key, filter_=None): """Read a single row from this table. diff --git a/tests/system.py b/tests/system.py index 21a39eb29..48f7e3bdf 100644 --- a/tests/system.py +++ b/tests/system.py @@ -18,6 +18,8 @@ import time import unittest +import pytest + from google.api_core.datetime_helpers import DatetimeWithNanoseconds from google.api_core.exceptions import DeadlineExceeded from google.api_core.exceptions import TooManyRequests @@ -56,8 +58,8 @@ CLUSTER_ID = INSTANCE_ID + "-cluster" CLUSTER_ID_DATA = INSTANCE_ID_DATA + "-cluster" SERVE_NODES = 3 -COLUMN_FAMILY_ID1 = u"col-fam-id1" -COLUMN_FAMILY_ID2 = u"col-fam-id2" +COLUMN_FAMILY_ID1 = "col-fam-id1" +COLUMN_FAMILY_ID2 = "col-fam-id2" COL_NAME1 = b"col-name1" COL_NAME2 = b"col-name2" COL_NAME3 = b"col-name3-but-other-fam" @@ -68,13 +70,14 @@ ROW_KEY = b"row-key" ROW_KEY_ALT = b"row-key-alt" EXISTING_INSTANCES = [] -LABEL_KEY = u"python-system" +LABEL_KEY = "python-system" label_stamp = ( datetime.datetime.utcnow() .replace(microsecond=0, tzinfo=UTC) .strftime("%Y-%m-%dt%H-%M-%S") ) LABELS = {LABEL_KEY: str(label_stamp)} +KMS_KEY_NAME = os.environ.get("KMS_KEY_NAME", None) class Config(object): @@ -121,13 +124,13 @@ def setUpModule(): Config.INSTANCE = Config.CLIENT.instance(INSTANCE_ID, labels=LABELS) Config.CLUSTER = Config.INSTANCE.cluster( - CLUSTER_ID, location_id=LOCATION_ID, serve_nodes=SERVE_NODES + CLUSTER_ID, location_id=LOCATION_ID, serve_nodes=SERVE_NODES, ) Config.INSTANCE_DATA = Config.CLIENT.instance( INSTANCE_ID_DATA, instance_type=Instance.Type.DEVELOPMENT, labels=LABELS ) Config.CLUSTER_DATA = Config.INSTANCE_DATA.cluster( - CLUSTER_ID_DATA, location_id=LOCATION_ID + CLUSTER_ID_DATA, location_id=LOCATION_ID, ) if not Config.IN_EMULATOR: @@ -331,6 +334,220 @@ def test_create_instance_w_two_clusters(self): temp_table_id = "test-get-cluster-states" temp_table = instance.table(temp_table_id) temp_table.create() + + encryption_info = temp_table.get_encryption_info() + self.assertEqual( + encryption_info[ALT_CLUSTER_ID_1][0].encryption_type, + enums.EncryptionInfo.EncryptionType.GOOGLE_DEFAULT_ENCRYPTION, + ) + self.assertEqual( + encryption_info[ALT_CLUSTER_ID_2][0].encryption_type, + enums.EncryptionInfo.EncryptionType.GOOGLE_DEFAULT_ENCRYPTION, + ) + + result = temp_table.get_cluster_states() + ReplicationState = enums.Table.ReplicationState + expected_results = [ + ClusterState(ReplicationState.STATE_NOT_KNOWN), + ClusterState(ReplicationState.INITIALIZING), + ClusterState(ReplicationState.PLANNED_MAINTENANCE), + ClusterState(ReplicationState.UNPLANNED_MAINTENANCE), + ClusterState(ReplicationState.READY), + ] + cluster_id_list = result.keys() + self.assertEqual(len(cluster_id_list), 2) + self.assertIn(ALT_CLUSTER_ID_1, cluster_id_list) + self.assertIn(ALT_CLUSTER_ID_2, cluster_id_list) + for clusterstate in result.values(): + self.assertIn(clusterstate, expected_results) + + # Test create app profile with multi_cluster_routing policy + app_profiles_to_delete = [] + description = "routing policy-multy" + app_profile_id_1 = "app_profile_id_1" + routing = enums.RoutingPolicyType.ANY + self._test_create_app_profile_helper( + app_profile_id_1, + instance, + routing_policy_type=routing, + description=description, + ignore_warnings=True, + ) + app_profiles_to_delete.append(app_profile_id_1) + + # Test list app profiles + self._test_list_app_profiles_helper(instance, [app_profile_id_1]) + + # Test modify app profile app_profile_id_1 + # routing policy to single cluster policy, + # cluster -> ALT_CLUSTER_ID_1, + # allow_transactional_writes -> disallowed + # modify description + description = "to routing policy-single" + routing = enums.RoutingPolicyType.SINGLE + self._test_modify_app_profile_helper( + app_profile_id_1, + instance, + routing_policy_type=routing, + description=description, + cluster_id=ALT_CLUSTER_ID_1, + allow_transactional_writes=False, + ) + + # Test modify app profile app_profile_id_1 + # cluster -> ALT_CLUSTER_ID_2, + # allow_transactional_writes -> allowed + self._test_modify_app_profile_helper( + app_profile_id_1, + instance, + routing_policy_type=routing, + description=description, + cluster_id=ALT_CLUSTER_ID_2, + allow_transactional_writes=True, + ignore_warnings=True, + ) + + # Test create app profile with single cluster routing policy + description = "routing policy-single" + app_profile_id_2 = "app_profile_id_2" + routing = enums.RoutingPolicyType.SINGLE + self._test_create_app_profile_helper( + app_profile_id_2, + instance, + routing_policy_type=routing, + description=description, + cluster_id=ALT_CLUSTER_ID_2, + allow_transactional_writes=False, + ) + app_profiles_to_delete.append(app_profile_id_2) + + # Test list app profiles + self._test_list_app_profiles_helper( + instance, [app_profile_id_1, app_profile_id_2] + ) + + # Test modify app profile app_profile_id_2 to + # allow transactional writes + # Note: no need to set ``ignore_warnings`` to True + # since we are not restrictings anything with this modification. + self._test_modify_app_profile_helper( + app_profile_id_2, + instance, + routing_policy_type=routing, + description=description, + cluster_id=ALT_CLUSTER_ID_2, + allow_transactional_writes=True, + ) + + # Test modify app profile app_profile_id_2 routing policy + # to multi_cluster_routing policy + # modify description + description = "to routing policy-multy" + routing = enums.RoutingPolicyType.ANY + self._test_modify_app_profile_helper( + app_profile_id_2, + instance, + routing_policy_type=routing, + description=description, + allow_transactional_writes=False, + ignore_warnings=True, + ) + + # Test delete app profiles + for app_profile_id in app_profiles_to_delete: + self._test_delete_app_profile_helper(app_profile_id, instance) + + @pytest.mark.skipif( + not KMS_KEY_NAME, reason="requires KMS_KEY_NAME environment variable" + ) + def test_create_instance_w_two_clusters_cmek(self): + from google.cloud.bigtable import enums + from google.cloud.bigtable.table import ClusterState + + _PRODUCTION = enums.Instance.Type.PRODUCTION + ALT_INSTANCE_ID = "dif-cmek" + UNIQUE_SUFFIX + instance = Config.CLIENT.instance( + ALT_INSTANCE_ID, instance_type=_PRODUCTION, labels=LABELS + ) + + ALT_CLUSTER_ID_1 = ALT_INSTANCE_ID + "-c1" + ALT_CLUSTER_ID_2 = ALT_INSTANCE_ID + "-c2" + LOCATION_ID_2 = "us-central1-f" + STORAGE_TYPE = enums.StorageType.HDD + serve_nodes = 1 + cluster_1 = instance.cluster( + ALT_CLUSTER_ID_1, + location_id=LOCATION_ID, + serve_nodes=serve_nodes, + default_storage_type=STORAGE_TYPE, + kms_key_name=KMS_KEY_NAME, + ) + cluster_2 = instance.cluster( + ALT_CLUSTER_ID_2, + location_id=LOCATION_ID_2, + serve_nodes=serve_nodes, + default_storage_type=STORAGE_TYPE, + kms_key_name=KMS_KEY_NAME, + ) + operation = instance.create(clusters=[cluster_1, cluster_2]) + + # Make sure this instance gets deleted after the test case. + self.instances_to_delete.append(instance) + + # We want to make sure the operation completes. + operation.result(timeout=120) + + # Create a new instance instance and make sure it is the same. + instance_alt = Config.CLIENT.instance(ALT_INSTANCE_ID) + instance_alt.reload() + + self.assertEqual(instance, instance_alt) + self.assertEqual(instance.display_name, instance_alt.display_name) + self.assertEqual(instance.type_, instance_alt.type_) + + clusters, failed_locations = instance_alt.list_clusters() + self.assertEqual(failed_locations, []) + + clusters.sort(key=lambda x: x.name) + alt_cluster_1, alt_cluster_2 = clusters + + self.assertEqual(cluster_1.location_id, alt_cluster_1.location_id) + self.assertEqual(alt_cluster_1.state, enums.Cluster.State.READY) + self.assertEqual(cluster_1.serve_nodes, alt_cluster_1.serve_nodes) + self.assertEqual( + cluster_1.default_storage_type, alt_cluster_1.default_storage_type + ) + self.assertEqual(cluster_2.location_id, alt_cluster_2.location_id) + self.assertEqual(alt_cluster_2.state, enums.Cluster.State.READY) + self.assertEqual(cluster_2.serve_nodes, alt_cluster_2.serve_nodes) + self.assertEqual( + cluster_2.default_storage_type, alt_cluster_2.default_storage_type + ) + + # Test list clusters in project via 'client.list_clusters' + clusters, failed_locations = Config.CLIENT.list_clusters() + self.assertFalse(failed_locations) + found = set([cluster.name for cluster in clusters]) + self.assertTrue( + {alt_cluster_1.name, alt_cluster_2.name, Config.CLUSTER.name}.issubset( + found + ) + ) + + temp_table_id = "test-get-cluster-states" + temp_table = instance.table(temp_table_id) + temp_table.create() + + encryption_info = temp_table.get_encryption_info() + self.assertEqual( + encryption_info[ALT_CLUSTER_ID_1][0].encryption_type, + enums.EncryptionInfo.EncryptionType.CUSTOMER_MANAGED_ENCRYPTION, + ) + self.assertEqual( + encryption_info[ALT_CLUSTER_ID_2][0].encryption_type, + enums.EncryptionInfo.EncryptionType.CUSTOMER_MANAGED_ENCRYPTION, + ) + result = temp_table.get_cluster_states() ReplicationState = enums.Table.ReplicationState expected_results = [ @@ -843,6 +1060,7 @@ def test_backup(self): self.skipTest("backups are not supported in the emulator") from google.cloud._helpers import _datetime_to_pb_timestamp + from google.cloud.bigtable import enums temp_table_id = "test-backup-table" temp_table = Config.INSTANCE_DATA.table(temp_table_id) @@ -879,6 +1097,10 @@ def test_backup(self): self.assertEqual(temp_backup_id, temp_table_backup.backup_id) self.assertEqual(CLUSTER_ID_DATA, temp_table_backup.cluster) self.assertEqual(expire, temp_table_backup.expire_time.seconds) + self.assertEqual( + temp_table_backup.encryption_info.encryption_type, + enums.EncryptionInfo.EncryptionType.GOOGLE_DEFAULT_ENCRYPTION, + ) # Testing `Backup.update_expire_time()` method expire += 3600 # A one-hour change in the `expire_time` parameter @@ -1213,13 +1435,13 @@ def test_read_with_label_applied(self): row.commit() # Combine a label with column 1. - label1 = u"label-red" + label1 = "label-red" label1_filter = ApplyLabelFilter(label1) col1_filter = ColumnQualifierRegexFilter(COL_NAME1) chain1 = RowFilterChain(filters=[col1_filter, label1_filter]) # Combine a label with column 2. - label2 = u"label-blue" + label2 = "label-blue" label2_filter = ApplyLabelFilter(label2) col2_filter = ColumnQualifierRegexFilter(COL_NAME2) chain2 = RowFilterChain(filters=[col2_filter, label2_filter]) diff --git a/tests/unit/test_backup.py b/tests/unit/test_backup.py index 02efef492..0a5ba74c1 100644 --- a/tests/unit/test_backup.py +++ b/tests/unit/test_backup.py @@ -66,6 +66,7 @@ def test_constructor_defaults(self): self.assertIsNone(backup._end_time) self.assertIsNone(backup._size_bytes) self.assertIsNone(backup._state) + self.assertIsNone(backup._encryption_info) def test_constructor_non_defaults(self): instance = _Instance(self.INSTANCE_NAME) @@ -77,6 +78,7 @@ def test_constructor_non_defaults(self): cluster_id=self.CLUSTER_ID, table_id=self.TABLE_ID, expire_time=expire_time, + encryption_info="encryption_info", ) self.assertEqual(backup.backup_id, self.BACKUP_ID) @@ -84,6 +86,7 @@ def test_constructor_non_defaults(self): self.assertIs(backup._cluster, self.CLUSTER_ID) self.assertEqual(backup.table_id, self.TABLE_ID) self.assertEqual(backup._expire_time, expire_time) + self.assertEqual(backup._encryption_info, "encryption_info") self.assertIsNone(backup._parent) self.assertIsNone(backup._source_table) @@ -128,14 +131,20 @@ def test_from_pb_bad_name(self): klasse.from_pb(backup_pb, instance) def test_from_pb_success(self): + from google.cloud.bigtable.encryption_info import EncryptionInfo + from google.cloud.bigtable.error import Status from google.cloud.bigtable_admin_v2.types import table from google.cloud._helpers import _datetime_to_pb_timestamp + from google.rpc.code_pb2 import Code client = _Client() instance = _Instance(self.INSTANCE_NAME, client) timestamp = _datetime_to_pb_timestamp(self._make_timestamp()) size_bytes = 1234 state = table.Backup.State.READY + GOOGLE_DEFAULT_ENCRYPTION = ( + table.EncryptionInfo.EncryptionType.GOOGLE_DEFAULT_ENCRYPTION + ) backup_pb = table.Backup( name=self.BACKUP_NAME, source_table=self.TABLE_NAME, @@ -144,6 +153,11 @@ def test_from_pb_success(self): end_time=timestamp, size_bytes=size_bytes, state=state, + encryption_info=table.EncryptionInfo( + encryption_type=GOOGLE_DEFAULT_ENCRYPTION, + encryption_status=_StatusPB(Code.OK, "Status OK"), + kms_key_version="2", + ), ) klasse = self._get_target_class() @@ -159,6 +173,14 @@ def test_from_pb_success(self): self.assertEqual(backup.end_time, timestamp) self.assertEqual(backup._size_bytes, size_bytes) self.assertEqual(backup._state, state) + self.assertEqual( + backup.encryption_info, + EncryptionInfo( + encryption_type=GOOGLE_DEFAULT_ENCRYPTION, + encryption_status=Status(_StatusPB(Code.OK, "Status OK")), + kms_key_version="2", + ), + ) def test_property_name(self): from google.cloud.bigtable.client import Client @@ -862,3 +884,13 @@ def __init__(self, name, client=None): self.name = name self.instance_id = name.rsplit("/", 1)[1] self._client = client + + +def _StatusPB(code, message): + from google.rpc import status_pb2 + + status_pb = status_pb2.Status() + status_pb.code = code + status_pb.message = message + + return status_pb diff --git a/tests/unit/test_cluster.py b/tests/unit/test_cluster.py index d5f731eb6..49a32ea56 100644 --- a/tests/unit/test_cluster.py +++ b/tests/unit/test_cluster.py @@ -16,6 +16,7 @@ import unittest import mock +import pytest from ._testing import _make_credentials @@ -60,6 +61,9 @@ class TestCluster(unittest.TestCase): OP_NAME = "operations/projects/{}/instances/{}/clusters/{}/operations/{}".format( PROJECT, INSTANCE_ID, CLUSTER_ID, OP_ID ) + KEY_RING_ID = "key-ring-id" + CRYPTO_KEY_ID = "crypto-key-id" + KMS_KEY_NAME = f"{LOCATION_PATH}/keyRings/{KEY_RING_ID}/cryptoKeys/{CRYPTO_KEY_ID}" @staticmethod def _get_target_class(): @@ -90,6 +94,7 @@ def test_constructor_defaults(self): self.assertIsNone(cluster.state) self.assertIsNone(cluster.serve_nodes) self.assertIsNone(cluster.default_storage_type) + self.assertIsNone(cluster.kms_key_name) def test_constructor_non_default(self): from google.cloud.bigtable.enums import StorageType @@ -107,6 +112,7 @@ def test_constructor_non_default(self): _state=STATE, serve_nodes=self.SERVE_NODES, default_storage_type=STORAGE_TYPE_SSD, + kms_key_name=self.KMS_KEY_NAME, ) self.assertEqual(cluster.cluster_id, self.CLUSTER_ID) self.assertIs(cluster._instance, instance) @@ -114,6 +120,7 @@ def test_constructor_non_default(self): self.assertEqual(cluster.state, STATE) self.assertEqual(cluster.serve_nodes, self.SERVE_NODES) self.assertEqual(cluster.default_storage_type, STORAGE_TYPE_SSD) + self.assertEqual(cluster.kms_key_name, self.KMS_KEY_NAME) def test_name_property(self): credentials = _make_credentials() @@ -125,6 +132,18 @@ def test_name_property(self): self.assertEqual(cluster.name, self.CLUSTER_NAME) + def test_kms_key_name_property(self): + client = _Client(self.PROJECT) + instance = _Instance(self.INSTANCE_ID, client) + + cluster = self._make_one( + self.CLUSTER_ID, instance, kms_key_name=self.KMS_KEY_NAME + ) + + self.assertEqual(cluster.kms_key_name, self.KMS_KEY_NAME) + with pytest.raises(AttributeError): + cluster.kms_key_name = "I'm read only" + def test_from_pb_success(self): from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2 from google.cloud.bigtable import enums @@ -141,6 +160,9 @@ def test_from_pb_success(self): state=state, serve_nodes=self.SERVE_NODES, default_storage_type=storage_type, + encryption_config=data_v2_pb2.Cluster.EncryptionConfig( + kms_key_name=self.KMS_KEY_NAME, + ), ) klass = self._get_target_class() @@ -152,6 +174,7 @@ def test_from_pb_success(self): self.assertEqual(cluster.state, state) self.assertEqual(cluster.serve_nodes, self.SERVE_NODES) self.assertEqual(cluster.default_storage_type, storage_type) + self.assertEqual(cluster.kms_key_name, self.KMS_KEY_NAME) def test_from_pb_bad_cluster_name(self): from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2 @@ -243,6 +266,7 @@ def test_reload(self): location_id=self.LOCATION_ID, serve_nodes=self.SERVE_NODES, default_storage_type=STORAGE_TYPE_SSD, + kms_key_name=self.KMS_KEY_NAME, ) # Create response_pb @@ -281,6 +305,7 @@ def test_reload(self): self.assertEqual(cluster.state, STATE) self.assertEqual(cluster.serve_nodes, SERVE_NODES_FROM_SERVER) self.assertEqual(cluster.default_storage_type, STORAGE_TYPE_FROM_SERVER) + self.assertEqual(cluster.kms_key_name, None) def test_exists(self): from google.cloud.bigtable_admin_v2.services.bigtable_instance_admin import ( @@ -392,6 +417,80 @@ def test_create(self): ].kwargs self.assertEqual(actual_request, expected_request) + def test_create_w_cmek(self): + import datetime + from google.longrunning import operations_pb2 + from google.protobuf.any_pb2 import Any + from google.cloud.bigtable_admin_v2.types import ( + bigtable_instance_admin as messages_v2_pb2, + ) + from google.cloud._helpers import _datetime_to_pb_timestamp + from google.cloud.bigtable.instance import Instance + from google.cloud.bigtable_admin_v2.services.bigtable_instance_admin import ( + BigtableInstanceAdminClient, + ) + from google.cloud.bigtable_admin_v2.types import instance as instance_v2_pb2 + from google.cloud.bigtable.enums import StorageType + + NOW = datetime.datetime.utcnow() + NOW_PB = _datetime_to_pb_timestamp(NOW) + credentials = _make_credentials() + client = self._make_client( + project=self.PROJECT, credentials=credentials, admin=True + ) + STORAGE_TYPE_SSD = StorageType.SSD + LOCATION = self.LOCATION_PATH + self.LOCATION_ID + instance = Instance(self.INSTANCE_ID, client) + cluster = self._make_one( + self.CLUSTER_ID, + instance, + location_id=self.LOCATION_ID, + serve_nodes=self.SERVE_NODES, + default_storage_type=STORAGE_TYPE_SSD, + kms_key_name=self.KMS_KEY_NAME, + ) + expected_request_cluster = instance_v2_pb2.Cluster( + location=LOCATION, + serve_nodes=cluster.serve_nodes, + default_storage_type=cluster.default_storage_type, + encryption_config=instance_v2_pb2.Cluster.EncryptionConfig( + kms_key_name=self.KMS_KEY_NAME, + ), + ) + expected_request = { + "request": { + "parent": instance.name, + "cluster_id": self.CLUSTER_ID, + "cluster": expected_request_cluster, + } + } + name = instance.name + metadata = messages_v2_pb2.CreateClusterMetadata(request_time=NOW_PB) + type_url = "type.googleapis.com/{}".format( + messages_v2_pb2.CreateClusterMetadata._meta._pb.DESCRIPTOR.full_name + ) + response_pb = operations_pb2.Operation( + name=self.OP_NAME, + metadata=Any(type_url=type_url, value=metadata._pb.SerializeToString()), + ) + + # Patch the stub used by the API method. + api = mock.create_autospec(BigtableInstanceAdminClient) + api.common_location_path.return_value = LOCATION + client._instance_admin_client = api + cluster._instance._client = client + cluster._instance._client.instance_admin_client.instance_path.return_value = ( + name + ) + client._instance_admin_client.create_cluster.return_value = response_pb + # Perform the method and check the result. + cluster.create() + + actual_request = client._instance_admin_client.create_cluster.call_args_list[ + 0 + ].kwargs + self.assertEqual(actual_request, expected_request) + def test_update(self): import datetime from google.longrunning import operations_pb2 diff --git a/tests/unit/test_table.py b/tests/unit/test_table.py index c52119192..ccb8350a3 100644 --- a/tests/unit/test_table.py +++ b/tests/unit/test_table.py @@ -534,6 +534,87 @@ def test_get_cluster_states(self): result = table.get_cluster_states() self.assertEqual(result, expected_result) + def test_get_encryption_info(self): + from google.rpc.code_pb2 import Code + from google.cloud.bigtable_admin_v2.services.bigtable_table_admin import ( + client as bigtable_table_admin, + ) + from google.cloud.bigtable.encryption_info import EncryptionInfo + from google.cloud.bigtable.enums import EncryptionInfo as enum_crypto + from google.cloud.bigtable.error import Status + + ENCRYPTION_TYPE_UNSPECIFIED = ( + enum_crypto.EncryptionType.ENCRYPTION_TYPE_UNSPECIFIED + ) + GOOGLE_DEFAULT_ENCRYPTION = enum_crypto.EncryptionType.GOOGLE_DEFAULT_ENCRYPTION + CUSTOMER_MANAGED_ENCRYPTION = ( + enum_crypto.EncryptionType.CUSTOMER_MANAGED_ENCRYPTION + ) + + table_api = mock.create_autospec(bigtable_table_admin.BigtableTableAdminClient) + credentials = _make_credentials() + client = self._make_client( + project="project-id", credentials=credentials, admin=True + ) + instance = client.instance(instance_id=self.INSTANCE_ID) + table = self._make_one(self.TABLE_ID, instance) + + response_pb = _TablePB( + cluster_states={ + "cluster-id1": _ClusterStateEncryptionInfoPB( + encryption_type=ENCRYPTION_TYPE_UNSPECIFIED, + encryption_status=_StatusPB(Code.OK, "Status OK"), + ), + "cluster-id2": _ClusterStateEncryptionInfoPB( + encryption_type=GOOGLE_DEFAULT_ENCRYPTION, + ), + "cluster-id3": _ClusterStateEncryptionInfoPB( + encryption_type=CUSTOMER_MANAGED_ENCRYPTION, + encryption_status=_StatusPB( + Code.UNKNOWN, "Key version is not yet known." + ), + kms_key_version="UNKNOWN", + ), + } + ) + + # Patch the stub used by the API method. + client._table_admin_client = table_api + bigtable_table_stub = client._table_admin_client + + bigtable_table_stub.get_table.side_effect = [response_pb] + + # build expected result + expected_result = { + "cluster-id1": ( + EncryptionInfo( + encryption_type=ENCRYPTION_TYPE_UNSPECIFIED, + encryption_status=Status(_StatusPB(Code.OK, "Status OK")), + kms_key_version="", + ), + ), + "cluster-id2": ( + EncryptionInfo( + encryption_type=GOOGLE_DEFAULT_ENCRYPTION, + encryption_status=Status(_StatusPB(0, "")), + kms_key_version="", + ), + ), + "cluster-id3": ( + EncryptionInfo( + encryption_type=CUSTOMER_MANAGED_ENCRYPTION, + encryption_status=Status( + _StatusPB(Code.UNKNOWN, "Key version is not yet known.") + ), + kms_key_version="UNKNOWN", + ), + ), + } + + # Perform the method and check the result. + result = table.get_encryption_info() + self.assertEqual(result, expected_result) + def _read_row_helper(self, chunks, expected_result, app_profile_id=None): from google.cloud._testing import _Monkey @@ -2257,5 +2338,31 @@ def _ClusterStatePB(replication_state): return table_v2_pb2.Table.ClusterState(replication_state=replication_state) +def _ClusterStateEncryptionInfoPB( + encryption_type, encryption_status=None, kms_key_version=None +): + from google.cloud.bigtable_admin_v2.types import table as table_v2_pb2 + + return table_v2_pb2.Table.ClusterState( + encryption_info=( + table_v2_pb2.EncryptionInfo( + encryption_type=encryption_type, + encryption_status=encryption_status, + kms_key_version=kms_key_version, + ), + ) + ) + + +def _StatusPB(code, message): + from google.rpc import status_pb2 + + status_pb = status_pb2.Status() + status_pb.code = code + status_pb.message = message + + return status_pb + + def _read_rows_retry_exception(exc): return isinstance(exc, DeadlineExceeded)