feat(bigquery): implement defaultEncryptionConfiguration on datasets (#9489)

HemangChothani · tswast · commit d92cd6c25b27 · 2019-10-18T16:32:38.000-07:00
* feat: add customer managed encryption key for dataset

* change as recommended.

* change paramter name as per document

* cosmetic changes

* feat(bigquery): refactor class imports as it moved to new file

* feat(bigquery): location name updated in key as suggested
diff --git a/google/cloud/bigquery/dataset.py b/google/cloud/bigquery/dataset.py
@@ -24,6 +24,7 @@
 from google.cloud.bigquery.model import ModelReference
 from google.cloud.bigquery.routine import RoutineReference
 from google.cloud.bigquery.table import TableReference
+from google.cloud.bigquery.encryption_configuration import EncryptionConfiguration
 
 
 def _get_table_reference(self, table_id):
@@ -361,6 +362,7 @@ class Dataset(object):
         "default_partition_expiration_ms": "defaultPartitionExpirationMs",
         "default_table_expiration_ms": "defaultTableExpirationMs",
         "friendly_name": "friendlyName",
+        "default_encryption_configuration": "defaultEncryptionConfiguration",
     }
 
     def __init__(self, dataset_ref):
@@ -573,6 +575,30 @@ def labels(self, value):
             raise ValueError("Pass a dict")
         self._properties["labels"] = value
 
+    @property
+    def default_encryption_configuration(self):
+        """google.cloud.bigquery.encryption_configuration.EncryptionConfiguration: Custom
+        encryption configuration for all tables in the dataset.
+
+        Custom encryption configuration (e.g., Cloud KMS keys) or :data:`None`
+        if using default encryption.
+
+        See `protecting data with Cloud KMS keys
+        <https://cloud.google.com/bigquery/docs/customer-managed-encryption>`_
+        in the BigQuery documentation.
+        """
+        prop = self._properties.get("defaultEncryptionConfiguration")
+        if prop:
+            prop = EncryptionConfiguration.from_api_repr(prop)
+        return prop
+
+    @default_encryption_configuration.setter
+    def default_encryption_configuration(self, value):
+        api_repr = value
+        if value:
+            api_repr = value.to_api_repr()
+        self._properties["defaultEncryptionConfiguration"] = api_repr
+
     @classmethod
     def from_string(cls, full_dataset_id):
         """Construct a dataset from fully-qualified dataset ID.
diff --git a/tests/unit/test_dataset.py b/tests/unit/test_dataset.py
@@ -275,6 +275,7 @@ class TestDataset(unittest.TestCase):
     PROJECT = "project"
     DS_ID = "dataset-id"
     DS_REF = DatasetReference(PROJECT, DS_ID)
+    KMS_KEY_NAME = "projects/1/locations/us/keyRings/1/cryptoKeys/1"
 
     @staticmethod
     def _get_target_class():
@@ -314,6 +315,7 @@ def _make_resource(self):
                 {"role": "WRITER", "specialGroup": "projectWriters"},
                 {"role": "READER", "specialGroup": "projectReaders"},
             ],
+            "defaultEncryptionConfiguration": {"kmsKeyName": self.KMS_KEY_NAME},
         }
 
     def _verify_access_entry(self, access_entries, resource):
@@ -369,6 +371,13 @@ def _verify_resource_properties(self, dataset, resource):
         self.assertEqual(dataset.description, resource.get("description"))
         self.assertEqual(dataset.friendly_name, resource.get("friendlyName"))
         self.assertEqual(dataset.location, resource.get("location"))
+        if "defaultEncryptionConfiguration" in resource:
+            self.assertEqual(
+                dataset.default_encryption_configuration.kms_key_name,
+                resource.get("defaultEncryptionConfiguration")["kmsKeyName"],
+            )
+        else:
+            self.assertIsNone(dataset.default_encryption_configuration)
 
         if "access" in resource:
             self._verify_access_entry(dataset.access_entries, resource)
@@ -558,6 +567,22 @@ def test_to_api_repr_w_custom_field(self):
         }
         self.assertEqual(resource, exp_resource)
 
+    def test_default_encryption_configuration_setter(self):
+        from google.cloud.bigquery.encryption_configuration import (
+            EncryptionConfiguration,
+        )
+
+        dataset = self._make_one(self.DS_REF)
+        encryption_configuration = EncryptionConfiguration(
+            kms_key_name=self.KMS_KEY_NAME
+        )
+        dataset.default_encryption_configuration = encryption_configuration
+        self.assertEqual(
+            dataset.default_encryption_configuration.kms_key_name, self.KMS_KEY_NAME
+        )
+        dataset.default_encryption_configuration = None
+        self.assertIsNone(dataset.default_encryption_configuration)
+
     def test_from_string(self):
         cls = self._get_target_class()
         got = cls.from_string("string-project.string_dataset")
