feat(bigquery): expose customer managed encryption key for ML models

bigquery/google/cloud/bigquery/model.py

@@ -25,6 +25,7 @@
 from google.api_core import datetime_helpers
 from google.cloud.bigquery import _helpers
 from google.cloud.bigquery_v2 import types
+from google.cloud.bigquery.table import EncryptionConfiguration
 
 
 class Model(object):
@@ -51,6 +52,7 @@ class Model(object):
         # have an exhaustive list of all mutable properties.
         "labels": "labels",
         "description": "description",
+        "encryption_configuration": "encryptionConfiguration",
     }
 
     def __init__(self, model_ref):
@@ -256,6 +258,30 @@ def labels(self, value):
             value = {}
         self._properties["labels"] = value
 
+    @property
+    def encryption_configuration(self):
+        """google.cloud.bigquery.table.EncryptionConfiguration: Custom
+        encryption configuration for the model.
+
+        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("encryptionConfiguration")
+        if prop:
+            prop = EncryptionConfiguration.from_api_repr(prop)
+        return prop
+
+    @encryption_configuration.setter
+    def encryption_configuration(self, value):
+        api_repr = value
+        if value:
+            api_repr = value.to_api_repr()
+        self._properties["encryptionConfiguration"] = api_repr
+
     @classmethod
     def from_api_repr(cls, resource):
         """Factory: construct a model resource given its API representation

bigquery/tests/unit/model/test_model.py

@@ -21,6 +21,8 @@
 import google.cloud._helpers
 from google.cloud.bigquery_v2.gapic import enums
 
+KMS_KEY_NAME = "projects/1/locations/global/keyRings/1/cryptoKeys/1"
+
 
 @pytest.fixture
 def target_class():
@@ -99,6 +101,7 @@ def test_from_api_repr(target_class):
             },
         ],
         "featureColumns": [],
+        "encryptionConfiguration": {"kmsKeyName": KMS_KEY_NAME},
     }
     got = target_class.from_api_repr(resource)
 
@@ -116,6 +119,7 @@ def test_from_api_repr(target_class):
     assert got.friendly_name == u"A friendly name."
     assert got.model_type == enums.Model.ModelType.LOGISTIC_REGRESSION
     assert got.labels == {"greeting": u"こんにちは"}
+    assert got.encryption_configuration.kms_key_name == KMS_KEY_NAME
     assert got.training_runs[0].training_options.initial_learn_rate == 1.0
     assert (
         got.training_runs[0]
@@ -160,6 +164,7 @@ def test_from_api_repr_w_minimal_resource(target_class):
     assert got.friendly_name is None
     assert got.model_type == enums.Model.ModelType.MODEL_TYPE_UNSPECIFIED
     assert got.labels == {}
+    assert got.encryption_configuration is None
     assert len(got.training_runs) == 0
     assert len(got.feature_columns) == 0
     assert len(got.label_columns) == 0
@@ -229,6 +234,17 @@ def test_from_api_repr_w_unknown_fields(target_class):
             ["labels"],
             {"labels": {"a-label": "a-value"}},
         ),
+        (
+            {
+                "friendlyName": "hello",
+                "description": "world",
+                "expirationTime": None,
+                "labels": {"a-label": "a-value"},
+                "encryptionConfiguration": {"kmsKeyName": KMS_KEY_NAME},
+            },
+            ["encryptionConfiguration"],
+            {"encryptionConfiguration": {"kmsKeyName": KMS_KEY_NAME}},
+        ),
     ],
 )
 def test_build_resource(object_under_test, resource, filter_fields, expected):
@@ -283,6 +299,18 @@ def test_replace_labels(object_under_test):
     assert object_under_test.labels == {}
 
 
+def test_set_encryption_configuration(object_under_test):
+    from google.cloud.bigquery.table import EncryptionConfiguration
+
+    assert not object_under_test.encryption_configuration
+    object_under_test.encryption_configuration = EncryptionConfiguration(
+        kms_key_name=KMS_KEY_NAME
+    )
+    assert object_under_test.encryption_configuration.kms_key_name == KMS_KEY_NAME
+    object_under_test.encryption_configuration = None
+    assert not object_under_test.encryption_configuration
+
+
 def test_repr(target_class):
     model = target_class("my-proj.my_dset.my_model")
     got = repr(model)