Merge branch 'main' of https://github.com/hwin16/python-spanner

googleapis · Sep 5, 2024 · 907b4a8 · 907b4a8
2 parents 57fdb72 + 13c3b43
commit 907b4a8
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 5 deletions.
diff --git a/google/cloud/spanner_v1/backup.py b/google/cloud/spanner_v1/backup.py
@@ -280,16 +280,24 @@ def create(self):
         if (
             (
                 self._encryption_config
-                and self._encryption_config.kms_key_name
+                and (
+                    self._encryption_config.kms_key_name
+                    or self._encryption_config.kms_key_names
+                )
                 and self._encryption_config.encryption_type
                 != CreateBackupEncryptionConfig.EncryptionType.CUSTOMER_MANAGED_ENCRYPTION
             )
             and self._encryption_config
-            and self._encryption_config.kms_key_name
+            and (
+                self._encryption_config.kms_key_name
+                or self._encryption_config.kms_key_names
+            )
             and self._encryption_config.encryption_type
             != CopyBackupEncryptionConfig.EncryptionType.CUSTOMER_MANAGED_ENCRYPTION
         ):
-            raise ValueError("kms_key_name only used with CUSTOMER_MANAGED_ENCRYPTION")
+            raise ValueError(
+                "kms_key_name or kms_key_names only used with CUSTOMER_MANAGED_ENCRYPTION"
+            )
 
         api = self._instance._client.database_admin_api
         metadata = _metadata_with_prefix(self.name)

diff --git a/google/cloud/spanner_v1/database.py b/google/cloud/spanner_v1/database.py
@@ -911,11 +911,16 @@ def restore(self, source):
             )
         if (
             self.encryption_config
-            and self.encryption_config.kms_key_name
+            and (
+                self.encryption_config.kms_key_name
+                or self.encryption_config.kms_key_names
+            )
             and self.encryption_config.encryption_type
             != RestoreDatabaseEncryptionConfig.EncryptionType.CUSTOMER_MANAGED_ENCRYPTION
         ):
-            raise ValueError("kms_key_name only used with CUSTOMER_MANAGED_ENCRYPTION")
+            raise ValueError(
+                "kms_key_name or kms_key_names only used with CUSTOMER_MANAGED_ENCRYPTION"
+            )
         api = self._instance._client.database_admin_api
         metadata = _metadata_with_prefix(self.name)
         request = RestoreDatabaseRequest(

diff --git a/tests/unit/test_database.py b/tests/unit/test_database.py
@@ -210,6 +210,18 @@ def test_ctor_w_encryption_config(self):
         self.assertIs(database._instance, instance)
         self.assertEqual(database._encryption_config, encryption_config)
 
+    def test_ctor_w_multiple_kms_keys(self):
+        from google.cloud.spanner_admin_database_v1 import EncryptionConfig
+
+        instance = _Instance(self.INSTANCE_NAME)
+        encryption_config = EncryptionConfig(kms_key_names=["kms_key1", "kms_key2"])
+        database = self._make_one(
+            self.DATABASE_ID, instance, encryption_config=encryption_config
+        )
+        self.assertEqual(database.database_id, self.DATABASE_ID)
+        self.assertIs(database._instance, instance)
+        self.assertEqual(database._encryption_config, encryption_config)
+
     def test_ctor_w_directed_read_options(self):
         client = _Client(directed_read_options=DIRECTED_READ_OPTIONS)
         instance = _Instance(self.INSTANCE_NAME, client=client)