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

* feat: expose customer managed encryption key for ML models

* feat: add encryptionConfiguration in _PROPERTY_TO_API_FIELD

* changes in condition

* change in document and parameter

* create a new file for class EncryptionConfiguration

* feat(bigquery): refactor test class of encryption configuration and change location in key

* feat(bigquery): add unit test in table class

* feat(bigquery): add apache license in file and test to show previous location works

Author: HemangChothani

bigquery/docs/reference.rst

@@ -88,7 +88,6 @@ Table
     table.TableReference
     table.Row
     table.RowIterator
-    table.EncryptionConfiguration
     table.TimePartitioning
     table.TimePartitioningType
 
@@ -173,6 +172,13 @@ Enums
 
     enums.StandardSqlDataTypes
 
+Encryption Configuration
+========================
+
+.. autosummary::
+    :toctree: generated
+
+    encryption_configuration.EncryptionConfiguration
 
 Additional Types
 ================

bigquery/google/cloud/bigquery/__init__.py

@@ -73,12 +73,12 @@
 from google.cloud.bigquery.routine import RoutineArgument
 from google.cloud.bigquery.routine import RoutineReference
 from google.cloud.bigquery.schema import SchemaField
-from google.cloud.bigquery.table import EncryptionConfiguration
 from google.cloud.bigquery.table import Table
 from google.cloud.bigquery.table import TableReference
 from google.cloud.bigquery.table import Row
 from google.cloud.bigquery.table import TimePartitioningType
 from google.cloud.bigquery.table import TimePartitioning
+from google.cloud.bigquery.encryption_configuration import EncryptionConfiguration
 
 __all__ = [
     "__version__",
@@ -94,7 +94,6 @@
     "DatasetReference",
     "AccessEntry",
     # Tables
-    "EncryptionConfiguration",
     "Table",
     "TableReference",
     "Row",
@@ -136,6 +135,8 @@
     "StandardSqlDataTypes",
     "SourceFormat",
     "WriteDisposition",
+    # EncryptionConfiguration
+    "EncryptionConfiguration",
 ]
 
 

bigquery/google/cloud/bigquery/encryption_configuration.py

@@ -0,0 +1,84 @@
+# Copyright 2015 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.
+
+"""Define class for the custom encryption configuration."""
+
+import copy
+
+
+class EncryptionConfiguration(object):
+    """Custom encryption configuration (e.g., Cloud KMS keys).
+
+    Args:
+        kms_key_name (str): resource ID of Cloud KMS key used for encryption
+    """
+
+    def __init__(self, kms_key_name=None):
+        self._properties = {}
+        if kms_key_name is not None:
+            self._properties["kmsKeyName"] = kms_key_name
+
+    @property
+    def kms_key_name(self):
+        """str: Resource ID of Cloud KMS key
+
+        Resource ID of Cloud KMS key or :data:`None` if using default
+        encryption.
+        """
+        return self._properties.get("kmsKeyName")
+
+    @kms_key_name.setter
+    def kms_key_name(self, value):
+        self._properties["kmsKeyName"] = value
+
+    @classmethod
+    def from_api_repr(cls, resource):
+        """Construct an encryption configuration from its API representation
+
+        Args:
+            resource (Dict[str, object]):
+                An encryption configuration representation as returned from
+                the API.
+
+        Returns:
+            google.cloud.bigquery.table.EncryptionConfiguration:
+                An encryption configuration parsed from ``resource``.
+        """
+        config = cls()
+        config._properties = copy.deepcopy(resource)
+        return config
+
+    def to_api_repr(self):
+        """Construct the API resource representation of this encryption
+        configuration.
+
+        Returns:
+            Dict[str, object]:
+                Encryption configuration as represented as an API resource
+        """
+        return copy.deepcopy(self._properties)
+
+    def __eq__(self, other):
+        if not isinstance(other, EncryptionConfiguration):
+            return NotImplemented
+        return self.kms_key_name == other.kms_key_name
+
+    def __ne__(self, other):
+        return not self == other
+
+    def __hash__(self):
+        return hash(self.kms_key_name)
+
+    def __repr__(self):
+        return "EncryptionConfiguration({})".format(self.kms_key_name)

bigquery/google/cloud/bigquery/job.py

@@ -37,12 +37,12 @@
 from google.cloud.bigquery.routine import RoutineReference
 from google.cloud.bigquery.schema import SchemaField
 from google.cloud.bigquery.table import _EmptyRowIterator
-from google.cloud.bigquery.table import EncryptionConfiguration
 from google.cloud.bigquery.table import _table_arg_to_table_ref
 from google.cloud.bigquery.table import TableReference
 from google.cloud.bigquery.table import Table
 from google.cloud.bigquery.table import TimePartitioning
 from google.cloud.bigquery import _helpers
+from google.cloud.bigquery.encryption_configuration import EncryptionConfiguration
 
 _DONE_STATE = "DONE"
 _STOPPED_REASON = "stopped"
@@ -1040,7 +1040,7 @@ def create_disposition(self, value):
 
     @property
     def destination_encryption_configuration(self):
-        """google.cloud.bigquery.table.EncryptionConfiguration: Custom
+        """google.cloud.bigquery.encryption_configuration.EncryptionConfiguration: Custom
         encryption configuration for the destination table.
 
         Custom encryption configuration (e.g., Cloud KMS keys) or :data:`None`
@@ -1434,7 +1434,7 @@ def schema(self):
 
     @property
     def destination_encryption_configuration(self):
-        """google.cloud.bigquery.table.EncryptionConfiguration: Custom
+        """google.cloud.bigquery.encryption_configuration.EncryptionConfiguration: Custom
         encryption configuration for the destination table.
 
         Custom encryption configuration (e.g., Cloud KMS keys)
@@ -1638,7 +1638,7 @@ def write_disposition(self, value):
 
     @property
     def destination_encryption_configuration(self):
-        """google.cloud.bigquery.table.EncryptionConfiguration: Custom
+        """google.cloud.bigquery.encryption_configuration.EncryptionConfiguration: Custom
         encryption configuration for the destination table.
 
         Custom encryption configuration (e.g., Cloud KMS keys) or :data:`None`
@@ -1709,7 +1709,7 @@ def write_disposition(self):
 
     @property
     def destination_encryption_configuration(self):
-        """google.cloud.bigquery.table.EncryptionConfiguration: Custom
+        """google.cloud.bigquery.encryption_configuration.EncryptionConfiguration: Custom
         encryption configuration for the destination table.
 
         Custom encryption configuration (e.g., Cloud KMS keys) or :data:`None`
@@ -2041,7 +2041,7 @@ def __init__(self, **kwargs):
 
     @property
     def destination_encryption_configuration(self):
-        """google.cloud.bigquery.table.EncryptionConfiguration: Custom
+        """google.cloud.bigquery.encryption_configuration.EncryptionConfiguration: Custom
         encryption configuration for the destination table.
 
         Custom encryption configuration (e.g., Cloud KMS keys) or :data:`None`
@@ -2460,7 +2460,7 @@ def destination(self):
 
     @property
     def destination_encryption_configuration(self):
-        """google.cloud.bigquery.table.EncryptionConfiguration: Custom
+        """google.cloud.bigquery.encryption_configuration.EncryptionConfiguration: Custom
         encryption configuration for the destination table.
 
         Custom encryption configuration (e.g., Cloud KMS keys) or :data:`None`

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.encryption_configuration import EncryptionConfiguration
 
 
 class Model(object):
@@ -48,6 +49,7 @@ class Model(object):
         # have an exhaustive list of all mutable properties.
         "labels": "labels",
         "description": "description",
+        "encryption_configuration": "encryptionConfiguration",
     }
 
     def __init__(self, model_ref):
@@ -253,6 +255,30 @@ def labels(self, value):
             value = {}
         self._properties["labels"] = value
 
+    @property
+    def encryption_configuration(self):
+        """google.cloud.bigquery.encryption_configuration.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/google/cloud/bigquery/table.py

@@ -55,6 +55,7 @@
 from google.cloud.bigquery.schema import _build_schema_resource
 from google.cloud.bigquery.schema import _parse_schema_resource
 from google.cloud.bigquery.external_config import ExternalConfig
+from google.cloud.bigquery.encryption_configuration import EncryptionConfiguration
 
 
 _LOGGER = logging.getLogger(__name__)
@@ -113,73 +114,6 @@ def _view_use_legacy_sql_getter(table):
         return True
 
 
-class EncryptionConfiguration(object):
-    """Custom encryption configuration (e.g., Cloud KMS keys).
-
-    Args:
-        kms_key_name (str): resource ID of Cloud KMS key used for encryption
-    """
-
-    def __init__(self, kms_key_name=None):
-        self._properties = {}
-        if kms_key_name is not None:
-            self._properties["kmsKeyName"] = kms_key_name
-
-    @property
-    def kms_key_name(self):
-        """str: Resource ID of Cloud KMS key
-
-        Resource ID of Cloud KMS key or :data:`None` if using default
-        encryption.
-        """
-        return self._properties.get("kmsKeyName")
-
-    @kms_key_name.setter
-    def kms_key_name(self, value):
-        self._properties["kmsKeyName"] = value
-
-    @classmethod
-    def from_api_repr(cls, resource):
-        """Construct an encryption configuration from its API representation
-
-        Args:
-            resource (Dict[str, object]):
-                An encryption configuration representation as returned from
-                the API.
-
-        Returns:
-            google.cloud.bigquery.table.EncryptionConfiguration:
-                An encryption configuration parsed from ``resource``.
-        """
-        config = cls()
-        config._properties = copy.deepcopy(resource)
-        return config
-
-    def to_api_repr(self):
-        """Construct the API resource representation of this encryption
-        configuration.
-
-        Returns:
-            Dict[str, object]:
-                Encryption configuration as represented as an API resource
-        """
-        return copy.deepcopy(self._properties)
-
-    def __eq__(self, other):
-        if not isinstance(other, EncryptionConfiguration):
-            return NotImplemented
-        return self.kms_key_name == other.kms_key_name
-
-    def __ne__(self, other):
-        return not self == other
-
-    def __hash__(self):
-        return hash(self.kms_key_name)
-
-    def __repr__(self):
-        return "EncryptionConfiguration({})".format(self.kms_key_name)
-
-
 class TableReference(object):
     """TableReferences are pointers to tables.
 
@@ -479,7 +413,7 @@ def labels(self, value):
 
     @property
     def encryption_configuration(self):
-        """google.cloud.bigquery.table.EncryptionConfiguration: Custom
+        """google.cloud.bigquery.encryption_configuration.EncryptionConfiguration: Custom
         encryption configuration for the table.
 
         Custom encryption configuration (e.g., Cloud KMS keys) or :data:`None`