feat: add support for creating backups with CMEK

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,18 @@ 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
@@ -72,6 +82,10 @@ def __init__(
         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):
@@ -165,6 +179,14 @@ def encryption_info(self):
         """
         return self._encryption_info
 
+    @property
+    def encryption_config(self):
+        """Encryption config for this database.
+        :rtype: :class:`~google.cloud.spanner_admin_instance_v1.CreateBackupEncryptionConfig`
+        :returns: an object representing the restore info 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.
@@ -224,10 +246,15 @@ 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,
+            encryption_config=self._encryption_config,
+        )
+
+        future = api.create_backup(
+            request=request,
             metadata=metadata,
         )
         return future

google/cloud/spanner_v1/instance.py

@@ -424,7 +424,7 @@ 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
@@ -445,6 +445,14 @@ def backup(self, backup_id, database="", expire_time=None, version_time=None):
             Optional. The version time that will be used to create 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`
         """
         try:
             return Backup(
@@ -453,6 +461,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(
@@ -461,6 +470,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/unit/test_backup.py

@@ -62,18 +62,53 @@ def test_ctor_defaults(self):
         self.assertIsNone(backup._expire_time)
 
     def test_ctor_non_defaults(self):
+        from google.cloud.spanner_admin_database_v1 import CreateBackupEncryptionConfig
         instance = _Instance(self.INSTANCE_NAME)
         timestamp = self._make_timestamp()
 
+        encryption_config = CreateBackupEncryptionConfig(
+            encryption_type=CreateBackupEncryptionConfig.EncryptionType.CUSTOMER_MANAGED_ENCRYPTION,
+            kms_key_name="key_name"
+        )
         backup = self._make_one(
-            self.BACKUP_ID, instance, database=self.DATABASE_NAME, expire_time=timestamp
+            self.BACKUP_ID,
+            instance,
+            database=self.DATABASE_NAME,
+            expire_time=timestamp,
+            encryption_config=encryption_config
+        )
+
+        self.assertEqual(backup.backup_id, self.BACKUP_ID)
+        self.assertIs(backup._instance, instance)
+        self.assertEqual(backup._database, self.DATABASE_NAME)
+        self.assertIsNotNone(backup._expire_time)
+        self.assertIs(backup._expire_time, timestamp)
+        self.assertEqual(backup.encryption_config, encryption_config)
+
+    def test_ctor_w_encryption_config_dict(self):
+        from google.cloud.spanner_admin_database_v1 import CreateBackupEncryptionConfig
+        instance = _Instance(self.INSTANCE_NAME)
+        timestamp = self._make_timestamp()
+
+        encryption_config = {
+            "encryption_type": 3,
+            "kms_key_name": "key_name"
+        }
+        backup = self._make_one(
+            self.BACKUP_ID,
+            instance,
+            database=self.DATABASE_NAME,
+            expire_time=timestamp,
+            encryption_config=encryption_config
         )
+        expected_encryption_config = CreateBackupEncryptionConfig(**encryption_config)
 
         self.assertEqual(backup.backup_id, self.BACKUP_ID)
         self.assertIs(backup._instance, instance)
         self.assertEqual(backup._database, self.DATABASE_NAME)
         self.assertIsNotNone(backup._expire_time)
         self.assertIs(backup._expire_time, timestamp)
+        self.assertEqual(backup.encryption_config, expected_encryption_config)
 
     def test_from_pb_project_mismatch(self):
         from google.cloud.spanner_admin_database_v1 import Backup
@@ -180,10 +215,22 @@ def test_encrpytion_info_property(self):
         )
         self.assertEqual(backup.encryption_info, expected)
 
+    def test_encryption_config_property(self):
+        from google.cloud.spanner_admin_database_v1 import CreateBackupEncryptionConfig
+
+        instance = _Instance(self.INSTANCE_NAME)
+        backup = self._make_one(self.BACKUP_ID, instance)
+        expected = backup._encryption_config = CreateBackupEncryptionConfig(
+            encryption_type=CreateBackupEncryptionConfig.EncryptionType.CUSTOMER_MANAGED_ENCRYPTION,
+            kms_key_name="kms_key_name"
+        )
+        self.assertEqual(backup.encryption_config, expected)
+
     def test_create_grpc_error(self):
         from google.api_core.exceptions import GoogleAPICallError
         from google.api_core.exceptions import Unknown
         from google.cloud.spanner_admin_database_v1 import Backup
+        from google.cloud.spanner_admin_database_v1 import CreateBackupRequest
 
         client = _Client()
         api = client.database_admin_api = self._make_database_admin_api()
@@ -200,16 +247,21 @@ def test_create_grpc_error(self):
         with self.assertRaises(GoogleAPICallError):
             backup.create()
 
-        api.create_backup.assert_called_once_with(
+        request = CreateBackupRequest(
             parent=self.INSTANCE_NAME,
             backup_id=self.BACKUP_ID,
             backup=backup_pb,
+        )
+
+        api.create_backup.assert_called_once_with(
+            request=request,
             metadata=[("google-cloud-resource-prefix", backup.name)],
         )
 
     def test_create_already_exists(self):
         from google.cloud.exceptions import Conflict
         from google.cloud.spanner_admin_database_v1 import Backup
+        from google.cloud.spanner_admin_database_v1 import CreateBackupRequest
 
         client = _Client()
         api = client.database_admin_api = self._make_database_admin_api()
@@ -226,16 +278,21 @@ def test_create_already_exists(self):
         with self.assertRaises(Conflict):
             backup.create()
 
-        api.create_backup.assert_called_once_with(
+        request = CreateBackupRequest(
             parent=self.INSTANCE_NAME,
             backup_id=self.BACKUP_ID,
             backup=backup_pb,
+        )
+
+        api.create_backup.assert_called_once_with(
+            request=request,
             metadata=[("google-cloud-resource-prefix", backup.name)],
         )
 
     def test_create_instance_not_found(self):
         from google.cloud.exceptions import NotFound
         from google.cloud.spanner_admin_database_v1 import Backup
+        from google.cloud.spanner_admin_database_v1 import CreateBackupRequest
 
         client = _Client()
         api = client.database_admin_api = self._make_database_admin_api()
@@ -252,10 +309,14 @@ def test_create_instance_not_found(self):
         with self.assertRaises(NotFound):
             backup.create()
 
-        api.create_backup.assert_called_once_with(
+        request = CreateBackupRequest(
             parent=self.INSTANCE_NAME,
             backup_id=self.BACKUP_ID,
             backup=backup_pb,
+        )
+
+        api.create_backup.assert_called_once_with(
+            request=request,
             metadata=[("google-cloud-resource-prefix", backup.name)],
         )
 
@@ -276,6 +337,8 @@ def test_create_database_not_set(self):
 
     def test_create_success(self):
         from google.cloud.spanner_admin_database_v1 import Backup
+        from google.cloud.spanner_admin_database_v1 import CreateBackupRequest
+        from google.cloud.spanner_admin_database_v1 import CreateBackupEncryptionConfig
         from datetime import datetime
         from datetime import timedelta
         from pytz import UTC
@@ -289,12 +352,17 @@ def test_create_success(self):
         version_timestamp = datetime.utcnow() - timedelta(minutes=5)
         version_timestamp = version_timestamp.replace(tzinfo=UTC)
         expire_timestamp = self._make_timestamp()
+        encryption_config = {
+            "encryption_type": 3,
+            "kms_key_name": "key_name"
+        }
         backup = self._make_one(
             self.BACKUP_ID,
             instance,
             database=self.DATABASE_NAME,
             expire_time=expire_timestamp,
             version_time=version_timestamp,
+            encryption_config=encryption_config
         )
 
         backup_pb = Backup(
@@ -306,10 +374,16 @@ def test_create_success(self):
         future = backup.create()
         self.assertIs(future, op_future)
 
-        api.create_backup.assert_called_once_with(
+        expected_encryption_config = CreateBackupEncryptionConfig(**encryption_config)
+        request = CreateBackupRequest(
             parent=self.INSTANCE_NAME,
             backup_id=self.BACKUP_ID,
             backup=backup_pb,
+            encryption_config=expected_encryption_config
+        )
+
+        api.create_backup.assert_called_once_with(
+            request=request,
             metadata=[("google-cloud-resource-prefix", backup.name)],
         )
 

tests/unit/test_instance.py

@@ -610,22 +610,28 @@ def test_backup_factory_explicit(self):
         import datetime
         from google.cloud._helpers import UTC
         from google.cloud.spanner_v1.backup import Backup
+        from google.cloud.spanner_admin_database_v1 import CreateBackupEncryptionConfig
 
         client = _Client(self.PROJECT)
         instance = self._make_one(self.INSTANCE_ID, client, self.CONFIG_NAME)
         BACKUP_ID = "backup-id"
         DATABASE_NAME = "database-name"
         timestamp = datetime.datetime.utcnow().replace(tzinfo=UTC)
+        encryption_config = CreateBackupEncryptionConfig(
+            encryption_type=CreateBackupEncryptionConfig.EncryptionType.CUSTOMER_MANAGED_ENCRYPTION,
+            kms_key_name="kms_key_name"
+        )
 
         backup = instance.backup(
-            BACKUP_ID, database=DATABASE_NAME, expire_time=timestamp
+            BACKUP_ID, database=DATABASE_NAME, expire_time=timestamp, encryption_config=encryption_config
         )
 
         self.assertIsInstance(backup, Backup)
         self.assertEqual(backup.backup_id, BACKUP_ID)
         self.assertIs(backup._instance, instance)
         self.assertEqual(backup._database, DATABASE_NAME)
         self.assertIs(backup._expire_time, timestamp)
+        self.assertEqual(backup._encryption_config, encryption_config)
 
     def test_list_backups_defaults(self):
         from google.cloud.spanner_admin_database_v1 import Backup as BackupPB