feat: add support for CMEK (#105)

* feat: add support for creating databases with CMEK

* refactor: use kwargs for EncryptionConfig conversion

* feat: add support for creating backups with CMEK

* feat: add support for restore a database with CMEK

* style: fix lint

* fix: verify that correct encryption type is used when using a key

* test: use non-default encryption for backup tests to test CMEK support

* test: fix encryption assertion

* test: fix encryption type for assertion

* docs: fix docstring types

* docs: update docstring descriptions

Co-authored-by: larkee <[email protected]>

Author: larkee

google/cloud/spanner_v1/backup.py

@@ -19,6 +19,8 @@
 from google.cloud.exceptions import NotFound
 
 from google.cloud.spanner_admin_database_v1 import Backup as BackupPB
+from google.cloud.spanner_admin_database_v1 import CreateBackupEncryptionConfig
+from google.cloud.spanner_admin_database_v1 import CreateBackupRequest
 from google.cloud.spanner_v1._helpers import _metadata_with_prefix
 
 _BACKUP_NAME_RE = re.compile(
@@ -57,10 +59,24 @@ class Backup(object):
                         the externally consistent copy of the database. If
                         not present, it is the same as the `create_time` of
                         the backup.
+
+    :type encryption_config:
+        :class:`~google.cloud.spanner_admin_database_v1.types.CreateBackupEncryptionConfig`
+        or :class:`dict`
+    :param encryption_config:
+        (Optional) Encryption configuration for the backup.
+        If a dict is provided, it must be of the same form as the protobuf
+        message :class:`~google.cloud.spanner_admin_database_v1.types.CreateBackupEncryptionConfig`
     """
 
     def __init__(
-        self, backup_id, instance, database="", expire_time=None, version_time=None
+        self,
+        backup_id,
+        instance,
+        database="",
+        expire_time=None,
+        version_time=None,
+        encryption_config=None,
     ):
         self.backup_id = backup_id
         self._instance = instance
@@ -71,6 +87,11 @@ def __init__(
         self._size_bytes = None
         self._state = None
         self._referencing_databases = None
+        self._encryption_info = None
+        if type(encryption_config) == dict:
+            self._encryption_config = CreateBackupEncryptionConfig(**encryption_config)
+        else:
+            self._encryption_config = encryption_config
 
     @property
     def name(self):
@@ -156,6 +177,22 @@ def referencing_databases(self):
         """
         return self._referencing_databases
 
+    @property
+    def encryption_info(self):
+        """Encryption info for this backup.
+        :rtype: :class:`~google.clod.spanner_admin_database_v1.types.EncryptionInfo`
+        :returns: a class representing the encryption info
+        """
+        return self._encryption_info
+
+    @property
+    def encryption_config(self):
+        """Encryption config for this database.
+        :rtype: :class:`~google.cloud.spanner_admin_instance_v1.types.CreateBackupEncryptionConfig`
+        :returns: an object representing the encryption config for this database
+        """
+        return self._encryption_config
+
     @classmethod
     def from_pb(cls, backup_pb, instance):
         """Create an instance of this class from a protobuf message.
@@ -207,6 +244,13 @@ def create(self):
             raise ValueError("expire_time not set")
         if not self._database:
             raise ValueError("database not set")
+        if (
+            self.encryption_config
+            and self.encryption_config.kms_key_name
+            and self.encryption_config.encryption_type
+            != CreateBackupEncryptionConfig.EncryptionType.CUSTOMER_MANAGED_ENCRYPTION
+        ):
+            raise ValueError("kms_key_name only used with CUSTOMER_MANAGED_ENCRYPTION")
         api = self._instance._client.database_admin_api
         metadata = _metadata_with_prefix(self.name)
         backup = BackupPB(
@@ -215,12 +259,14 @@ def create(self):
             version_time=self.version_time,
         )
 
-        future = api.create_backup(
+        request = CreateBackupRequest(
             parent=self._instance.name,
             backup_id=self.backup_id,
             backup=backup,
-            metadata=metadata,
+            encryption_config=self._encryption_config,
         )
+
+        future = api.create_backup(request=request, metadata=metadata,)
         return future
 
     def exists(self):
@@ -255,6 +301,7 @@ def reload(self):
         self._size_bytes = pb.size_bytes
         self._state = BackupPB.State(pb.state)
         self._referencing_databases = pb.referencing_databases
+        self._encryption_info = pb.encryption_info
 
     def update_expire_time(self, new_expire_time):
         """Update the expire time of this backup.

google/cloud/spanner_v1/database.py

@@ -47,6 +47,9 @@
     SpannerGrpcTransport,
 )
 from google.cloud.spanner_admin_database_v1 import CreateDatabaseRequest
+from google.cloud.spanner_admin_database_v1 import EncryptionConfig
+from google.cloud.spanner_admin_database_v1 import RestoreDatabaseEncryptionConfig
+from google.cloud.spanner_admin_database_v1 import RestoreDatabaseRequest
 from google.cloud.spanner_admin_database_v1 import UpdateDatabaseDdlRequest
 from google.cloud.spanner_v1 import (
     ExecuteSqlRequest,
@@ -108,12 +111,27 @@ class Database(object):
                    is `True` to log commit statistics. If not passed, a logger
                    will be created when needed that will log the commit statistics
                    to stdout.
+    :type encryption_config:
+        :class:`~google.cloud.spanner_admin_database_v1.types.EncryptionConfig`
+        or :class:`~google.cloud.spanner_admin_database_v1.types.RestoreDatabaseEncryptionConfig`
+        or :class:`dict`
+    :param encryption_config:
+        (Optional) Encryption configuration for the database.
+        If a dict is provided, it must be of the same form as either of the protobuf
+        messages :class:`~google.cloud.spanner_admin_database_v1.types.EncryptionConfig`
+        or :class:`~google.cloud.spanner_admin_database_v1.types.RestoreDatabaseEncryptionConfig`
     """
 
     _spanner_api = None
 
     def __init__(
-        self, database_id, instance, ddl_statements=(), pool=None, logger=None
+        self,
+        database_id,
+        instance,
+        ddl_statements=(),
+        pool=None,
+        logger=None,
+        encryption_config=None,
     ):
         self.database_id = database_id
         self._instance = instance
@@ -126,6 +144,7 @@ def __init__(
         self._earliest_version_time = None
         self.log_commit_stats = False
         self._logger = logger
+        self._encryption_config = encryption_config
 
         if pool is None:
             pool = BurstyPool()
@@ -242,6 +261,14 @@ def earliest_version_time(self):
         """
         return self._earliest_version_time
 
+    @property
+    def encryption_config(self):
+        """Encryption config for this database.
+        :rtype: :class:`~google.cloud.spanner_admin_instance_v1.types.EncryptionConfig`
+        :returns: an object representing the encryption config for this database
+        """
+        return self._encryption_config
+
     @property
     def ddl_statements(self):
         """DDL Statements used to define database schema.
@@ -325,11 +352,14 @@ def create(self):
         db_name = self.database_id
         if "-" in db_name:
             db_name = "`%s`" % (db_name,)
+        if type(self._encryption_config) == dict:
+            self._encryption_config = EncryptionConfig(**self._encryption_config)
 
         request = CreateDatabaseRequest(
             parent=self._instance.name,
             create_statement="CREATE DATABASE %s" % (db_name,),
             extra_statements=list(self._ddl_statements),
+            encryption_config=self._encryption_config,
         )
         future = api.create_database(request=request, metadata=metadata)
         return future
@@ -372,6 +402,7 @@ def reload(self):
         self._restore_info = response.restore_info
         self._version_retention_period = response.version_retention_period
         self._earliest_version_time = response.earliest_version_time
+        self._encryption_config = response.encryption_config
 
     def update_ddl(self, ddl_statements, operation_id=""):
         """Update DDL for this database.
@@ -588,8 +619,8 @@ def run_in_transaction(self, func, *args, **kw):
     def restore(self, source):
         """Restore from a backup to this database.
 
-        :type backup: :class:`~google.cloud.spanner_v1.backup.Backup`
-        :param backup: the path of the backup being restored from.
+        :type source: :class:`~google.cloud.spanner_v1.backup.Backup`
+        :param source: the path of the source being restored from.
 
         :rtype: :class:`~google.api_core.operation.Operation`
         :returns: a future used to poll the status of the create request
@@ -601,14 +632,26 @@ def restore(self, source):
         """
         if source is None:
             raise ValueError("Restore source not specified")
+        if type(self._encryption_config) == dict:
+            self._encryption_config = RestoreDatabaseEncryptionConfig(
+                **self._encryption_config
+            )
+        if (
+            self.encryption_config
+            and self.encryption_config.kms_key_name
+            and self.encryption_config.encryption_type
+            != RestoreDatabaseEncryptionConfig.EncryptionType.CUSTOMER_MANAGED_ENCRYPTION
+        ):
+            raise ValueError("kms_key_name only used with CUSTOMER_MANAGED_ENCRYPTION")
         api = self._instance._client.database_admin_api
         metadata = _metadata_with_prefix(self.name)
-        future = api.restore_database(
+        request = RestoreDatabaseRequest(
             parent=self._instance.name,
             database_id=self.database_id,
             backup=source.name,
-            metadata=metadata,
+            encryption_config=self._encryption_config,
         )
+        future = api.restore_database(request=request, metadata=metadata,)
         return future
 
     def is_ready(self):

google/cloud/spanner_v1/instance.py

@@ -357,7 +357,14 @@ def delete(self):
 
         api.delete_instance(name=self.name, metadata=metadata)
 
-    def database(self, database_id, ddl_statements=(), pool=None, logger=None):
+    def database(
+        self,
+        database_id,
+        ddl_statements=(),
+        pool=None,
+        logger=None,
+        encryption_config=None,
+    ):
         """Factory to create a database within this instance.
 
         :type database_id: str
@@ -377,11 +384,26 @@ def database(self, database_id, ddl_statements=(), pool=None, logger=None):
                        will be created when needed that will log the commit statistics
                        to stdout.
 
+        :type encryption_config:
+            :class:`~google.cloud.spanner_admin_database_v1.types.EncryptionConfig`
+            or :class:`~google.cloud.spanner_admin_database_v1.types.RestoreDatabaseEncryptionConfig`
+            or :class:`dict`
+        :param encryption_config:
+            (Optional) Encryption configuration for the database.
+            If a dict is provided, it must be of the same form as either of the protobuf
+            messages :class:`~google.cloud.spanner_admin_database_v1.types.EncryptionConfig`
+            or :class:`~google.cloud.spanner_admin_database_v1.types.RestoreDatabaseEncryptionConfig`
+
         :rtype: :class:`~google.cloud.spanner_v1.database.Database`
         :returns: a database owned by this instance.
         """
         return Database(
-            database_id, self, ddl_statements=ddl_statements, pool=pool, logger=logger
+            database_id,
+            self,
+            ddl_statements=ddl_statements,
+            pool=pool,
+            logger=logger,
+            encryption_config=encryption_config,
         )
 
     def list_databases(self, page_size=None):
@@ -408,7 +430,14 @@ def list_databases(self, page_size=None):
         )
         return page_iter
 
-    def backup(self, backup_id, database="", expire_time=None, version_time=None):
+    def backup(
+        self,
+        backup_id,
+        database="",
+        expire_time=None,
+        version_time=None,
+        encryption_config=None,
+    ):
         """Factory to create a backup within this instance.
 
         :type backup_id: str
@@ -430,6 +459,14 @@ def backup(self, backup_id, database="", expire_time=None, version_time=None):
             consistent copy of the database. If not present, it is the same as
             the `create_time` of the backup.
 
+        :type encryption_config:
+            :class:`~google.cloud.spanner_admin_database_v1.types.CreateBackupEncryptionConfig`
+            or :class:`dict`
+        :param encryption_config:
+            (Optional) Encryption configuration for the backup.
+            If a dict is provided, it must be of the same form as the protobuf
+            message :class:`~google.cloud.spanner_admin_database_v1.types.CreateBackupEncryptionConfig`
+
         :rtype: :class:`~google.cloud.spanner_v1.backup.Backup`
         :returns: a backup owned by this instance.
         """
@@ -440,6 +477,7 @@ def backup(self, backup_id, database="", expire_time=None, version_time=None):
                 database=database.name,
                 expire_time=expire_time,
                 version_time=version_time,
+                encryption_config=encryption_config,
             )
         except AttributeError:
             return Backup(
@@ -448,6 +486,7 @@ def backup(self, backup_id, database="", expire_time=None, version_time=None):
                 database=database,
                 expire_time=expire_time,
                 version_time=version_time,
+                encryption_config=encryption_config,
             )
 
     def list_backups(self, filter_="", page_size=None):

tests/system/test_system.py

@@ -738,6 +738,11 @@ def test_create_invalid(self):
             op.result()
 
     def test_backup_workflow(self):
+        from google.cloud.spanner_admin_database_v1 import (
+            CreateBackupEncryptionConfig,
+            EncryptionConfig,
+            RestoreDatabaseEncryptionConfig,
+        )
         from datetime import datetime
         from datetime import timedelta
         from pytz import UTC
@@ -746,13 +751,17 @@ def test_backup_workflow(self):
         backup_id = "backup_id" + unique_resource_id("_")
         expire_time = datetime.utcnow() + timedelta(days=3)
         expire_time = expire_time.replace(tzinfo=UTC)
+        encryption_config = CreateBackupEncryptionConfig(
+            encryption_type=CreateBackupEncryptionConfig.EncryptionType.GOOGLE_DEFAULT_ENCRYPTION,
+        )
 
         # Create backup.
         backup = instance.backup(
             backup_id,
             database=self._db,
             expire_time=expire_time,
             version_time=self.database_version_time,
+            encryption_config=encryption_config,
         )
         operation = backup.create()
         self.to_delete.append(backup)
@@ -771,6 +780,7 @@ def test_backup_workflow(self):
         self.assertEqual(self.database_version_time, backup.version_time)
         self.assertIsNotNone(backup.size_bytes)
         self.assertIsNotNone(backup.state)
+        self.assertEqual(encryption_config, backup.encryption_config)
 
         # Update with valid argument.
         valid_expire_time = datetime.utcnow() + timedelta(days=7)
@@ -780,7 +790,10 @@ def test_backup_workflow(self):
 
         # Restore database to same instance.
         restored_id = "restored_db" + unique_resource_id("_")
-        database = instance.database(restored_id)
+        encryption_config = RestoreDatabaseEncryptionConfig(
+            encryption_type=RestoreDatabaseEncryptionConfig.EncryptionType.GOOGLE_DEFAULT_ENCRYPTION,
+        )
+        database = instance.database(restored_id, encryption_config=encryption_config)
         self.to_drop.append(database)
         operation = database.restore(source=backup)
         restored_db = operation.result()
@@ -791,6 +804,9 @@ def test_backup_workflow(self):
 
         metadata = operation.metadata
         self.assertEqual(self.database_version_time, metadata.backup_info.version_time)
+        database.reload()
+        expected_encryption_config = EncryptionConfig()
+        self.assertEqual(expected_encryption_config, database.encryption_config)
 
         database.drop()
         backup.delete()