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(
@@ -51,9 +53,17 @@ class Backup(object):
     :param expire_time: (Optional) The expire time that will be used to
                         create the backup. Required if the create method
                         needs to be called.
+
+    :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):
+    def __init__(self, backup_id, instance, database="", expire_time=None, encryption_config=None):
         self.backup_id = backup_id
         self._instance = instance
         self._database = database
@@ -63,6 +73,10 @@ def __init__(self, backup_id, instance, database="", expire_time=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):
@@ -146,6 +160,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.
@@ -200,11 +222,15 @@ def create(self):
         api = self._instance._client.database_admin_api
         metadata = _metadata_with_prefix(self.name)
         backup = BackupPB(database=self._database, expire_time=self.expire_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

@@ -415,7 +415,7 @@ def list_databases(self, page_size=None):
         )
         return page_iter
 
-    def backup(self, backup_id, database="", expire_time=None):
+    def backup(self, backup_id, database="", expire_time=None, encryption_config=None):
         """Factory to create a backup within this instance.
 
         :type backup_id: str
@@ -430,10 +430,18 @@ def backup(self, backup_id, database="", expire_time=None):
         :param expire_time:
             Optional. The expire time that will be used when creating the backup.
             Required if the create method needs to be called.
+
+        :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(
-                backup_id, self, database=database.name, expire_time=expire_time
+                backup_id, self, database=database.name, expire_time=expire_time, encryption_config=encryption_config
             )
         except AttributeError:
             return Backup(backup_id, self, database=database, expire_time=expire_time)

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,22 @@ 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 +279,22 @@ 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,13 +311,18 @@ 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)],
         )
 
+
     def test_create_expire_time_not_set(self):
         instance = _Instance(self.INSTANCE_NAME)
         backup = self._make_one(self.BACKUP_ID, instance, database=self.DATABASE_NAME)
@@ -276,6 +340,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
 
         op_future = object()
         client = _Client()
@@ -284,19 +350,33 @@ def test_create_success(self):
 
         instance = _Instance(self.INSTANCE_NAME, client=client)
         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
+            self.BACKUP_ID,
+            instance,
+            database=self.DATABASE_NAME,
+            expire_time=timestamp,
+            encryption_config=encryption_config
         )
 
-        backup_pb = Backup(database=self.DATABASE_NAME, expire_time=timestamp,)
+        backup_pb = Backup(database=self.DATABASE_NAME, expire_time=timestamp)
 
         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

@@ -605,22 +605,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