feat: add support for CMEK

Author: larkee

google/cloud/spanner_v1/backup.py

@@ -63,6 +63,7 @@ def __init__(self, backup_id, instance, database="", expire_time=None):
         self._size_bytes = None
         self._state = None
         self._referencing_databases = None
+        self._encryption_info = None
 
     @property
     def name(self):
@@ -138,6 +139,15 @@ 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.proto.common_pb2.EncryptionInfo`
+        :returns: a class representing the encryption info
+        """
+        return self._encryption_info
+
     @classmethod
     def from_pb(cls, backup_pb, instance):
         """Create an instance of this class from a protobuf message.
@@ -232,6 +242,7 @@ def reload(self):
         self._size_bytes = pb.size_bytes
         self._state = enums.Backup.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

@@ -29,6 +29,7 @@
 
 # pylint: disable=ungrouped-imports
 from google.cloud.spanner_admin_database_v1.gapic import enums
+from google.cloud.spanner_admin_database_v1.proto.common_pb2 import EncryptionConfig
 from google.cloud.spanner_v1._helpers import (
     _make_value_pb,
     _merge_query_options,
@@ -95,7 +96,14 @@ class Database(object):
 
     _spanner_api = None
 
-    def __init__(self, database_id, instance, ddl_statements=(), pool=None):
+    def __init__(
+        self,
+        database_id,
+        instance,
+        ddl_statements=(),
+        pool=None,
+        encryption_config=None,
+    ):
         self.database_id = database_id
         self._instance = instance
         self._ddl_statements = _check_ddl_statements(ddl_statements)
@@ -104,6 +112,13 @@ def __init__(self, database_id, instance, ddl_statements=(), pool=None):
         self._create_time = None
         self._restore_info = None
 
+        if type(encryption_config) == dict:
+            self._encryption_config = EncryptionConfig(
+                kms_key_name=encryption_config["kms_key_name"]
+            )
+        else:
+            self._encryption_config = encryption_config
+
         if pool is None:
             pool = BurstyPool()
 
@@ -200,6 +215,15 @@ def restore_info(self):
         """
         return self._restore_info
 
+    @property
+    def encryption_config(self):
+        """Encryption config for this database.
+
+        :rtype: :class:`~google.cloud.spanner_admin_instance_v1.proto.common_pb2.EncryptionConfig`
+        :returns: an object representing the restore info for this database
+        """
+        return self._encryption_config
+
     @property
     def ddl_statements(self):
         """DDL Statements used to define database schema.
@@ -269,6 +293,7 @@ def create(self):
             parent=self._instance.name,
             create_statement="CREATE DATABASE %s" % (db_name,),
             extra_statements=list(self._ddl_statements),
+            encryption_config=self._encryption_config,
             metadata=metadata,
         )
         return future
@@ -309,6 +334,7 @@ def reload(self):
         self._state = enums.Database.State(response.state)
         self._create_time = _pb_timestamp_to_datetime(response.create_time)
         self._restore_info = response.restore_info
+        self._encryption_config = response.encryption_config
 
     def update_ddl(self, ddl_statements, operation_id=""):
         """Update DDL for this database.

google/cloud/spanner_v1/instance.py

@@ -348,7 +348,9 @@ def delete(self):
 
         api.delete_instance(self.name, metadata=metadata)
 
-    def database(self, database_id, ddl_statements=(), pool=None):
+    def database(
+        self, database_id, ddl_statements=(), pool=None, encryption_config=None
+    ):
         """Factory to create a database within this instance.
 
         :type database_id: str
@@ -365,7 +367,13 @@ def database(self, database_id, ddl_statements=(), pool=None):
         :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)
+        return Database(
+            database_id,
+            self,
+            ddl_statements=ddl_statements,
+            pool=pool,
+            encryption_config=encryption_config,
+        )
 
     def list_databases(self, page_size=None, page_token=None):
         """List databases for the instance.

tests/unit/test_backup.py

@@ -170,6 +170,18 @@ def test_referencing_databases_property(self):
         expected = backup._referencing_databases = [self.DATABASE_NAME]
         self.assertEqual(backup.referencing_databases, expected)
 
+    def test_encrpytion_info_property(self):
+        from google.cloud.spanner_admin_database_v1.proto.common_pb2 import (
+            EncryptionInfo,
+        )
+
+        instance = _Instance(self.INSTANCE_NAME)
+        backup = self._make_one(self.BACKUP_ID, instance)
+        expected = backup._encryption_info = EncryptionInfo(
+            kms_key_version="kms_key_version"
+        )
+        self.assertEqual(backup.encryption_info, expected)
+
     def test_create_grpc_error(self):
         from google.api_core.exceptions import GoogleAPICallError
         from google.api_core.exceptions import Unknown
@@ -436,10 +448,14 @@ def test_reload_not_found(self):
 
     def test_reload_success(self):
         from google.cloud.spanner_admin_database_v1.proto import backup_pb2
+        from google.cloud.spanner_admin_database_v1.proto.common_pb2 import (
+            EncryptionInfo,
+        )
         from google.cloud.spanner_admin_database_v1.gapic import enums
         from google.cloud._helpers import _datetime_to_pb_timestamp
 
         timestamp = self._make_timestamp()
+        encryption_info = EncryptionInfo(kms_key_version="kms_key_version")
 
         client = _Client()
         backup_pb = backup_pb2.Backup(
@@ -450,6 +466,7 @@ def test_reload_success(self):
             size_bytes=10,
             state=1,
             referencing_databases=[],
+            encryption_info=encryption_info,
         )
         api = client.database_admin_api = self._make_database_admin_api()
         api.get_backup.return_value = backup_pb
@@ -464,6 +481,7 @@ def test_reload_success(self):
         self.assertEqual(backup.size_bytes, 10)
         self.assertEqual(backup.state, enums.Backup.State.CREATING)
         self.assertEqual(backup.referencing_databases, [])
+        self.assertEqual(backup.encryption_info, encryption_info)
 
         api.get_backup.assert_called_once_with(
             self.BACKUP_NAME, metadata=[("google-cloud-resource-prefix", backup.name)]

tests/unit/test_database.py

@@ -146,6 +146,35 @@ def test_ctor_w_ddl_statements_ok(self):
         self.assertIs(database._instance, instance)
         self.assertEqual(list(database.ddl_statements), DDL_STATEMENTS)
 
+    def ctor_w_encryption_config(self):
+        from google.cloud.spanner_admin_database_v1.proto.common_pb2 import (
+            EncryptionConfig,
+        )
+
+        instance = _Instance(self.INSTANCE_NAME)
+        encryption_config = EncryptionConfig(kms_key_name="kms_key")
+        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._encyption_config, encryption_config)
+
+    def ctor_w_encryption_config_dict(self):
+        from google.cloud.spanner_admin_database_v1.proto.common_pb2 import (
+            EncryptionConfig,
+        )
+
+        instance = _Instance(self.INSTANCE_NAME)
+        encryption_config_dict = {"kms_key_name": "kms_key"}
+        encryption_config = EncryptionConfig(kms_key_name="kms_key")
+        database = self._make_one(
+            self.DATABASE_ID, instance, encryption_config=encryption_config_dict
+        )
+        self.assertEqual(database.database_id, self.DATABASE_ID)
+        self.assertIs(database._instance, instance)
+        self.assertEqual(database._encyption_config, encryption_config)
+
     def test_from_pb_bad_database_name(self):
         from google.cloud.spanner_admin_database_v1.proto import (
             spanner_database_admin_pb2 as admin_v1_pb2,
@@ -260,6 +289,19 @@ def test_restore_info(self):
         )
         self.assertEqual(database.restore_info, restore_info)
 
+    def test_encryption_config(self):
+        from google.cloud.spanner_admin_database_v1.proto.common_pb2 import (
+            EncryptionConfig,
+        )
+
+        instance = _Instance(self.INSTANCE_NAME)
+        pool = _Pool()
+        database = self._make_one(self.DATABASE_ID, instance, pool=pool)
+        encryption_config = database._encryption_config = mock.create_autospec(
+            EncryptionConfig, instance=True
+        )
+        self.assertEqual(database.encryption_config, encryption_config)
+
     def test_spanner_api_property_w_scopeless_creds(self):
 
         client = _Client()
@@ -396,6 +438,7 @@ def test_create_grpc_error(self):
             parent=self.INSTANCE_NAME,
             create_statement="CREATE DATABASE {}".format(self.DATABASE_ID),
             extra_statements=[],
+            encryption_config=None,
             metadata=[("google-cloud-resource-prefix", database.name)],
         )
 
@@ -417,6 +460,7 @@ def test_create_already_exists(self):
             parent=self.INSTANCE_NAME,
             create_statement="CREATE DATABASE `{}`".format(DATABASE_ID_HYPHEN),
             extra_statements=[],
+            encryption_config=None,
             metadata=[("google-cloud-resource-prefix", database.name)],
         )
 
@@ -437,6 +481,7 @@ def test_create_instance_not_found(self):
             parent=self.INSTANCE_NAME,
             create_statement="CREATE DATABASE {}".format(self.DATABASE_ID),
             extra_statements=[],
+            encryption_config=None,
             metadata=[("google-cloud-resource-prefix", database.name)],
         )
 
@@ -461,6 +506,7 @@ def test_create_success(self):
             parent=self.INSTANCE_NAME,
             create_statement="CREATE DATABASE {}".format(self.DATABASE_ID),
             extra_statements=DDL_STATEMENTS,
+            encryption_config=None,
             metadata=[("google-cloud-resource-prefix", database.name)],
         )
 
@@ -561,6 +607,9 @@ def test_reload_success(self):
             spanner_database_admin_pb2 as admin_v1_pb2,
         )
         from google.cloud.spanner_admin_database_v1.gapic import enums
+        from google.cloud.spanner_admin_database_v1.proto.common_pb2 import (
+            EncryptionConfig,
+        )
         from google.cloud._helpers import _datetime_to_pb_timestamp
         from tests._fixtures import DDL_STATEMENTS
 
@@ -569,12 +618,14 @@ def test_reload_success(self):
 
         client = _Client()
         ddl_pb = admin_v1_pb2.GetDatabaseDdlResponse(statements=DDL_STATEMENTS)
+        encryption_config = EncryptionConfig(kms_key_name="kms_key")
         api = client.database_admin_api = self._make_database_admin_api()
         api.get_database_ddl.return_value = ddl_pb
         db_pb = admin_v1_pb2.Database(
             state=2,
             create_time=_datetime_to_pb_timestamp(timestamp),
             restore_info=restore_info,
+            encryption_config=encryption_config,
         )
         api.get_database.return_value = db_pb
         instance = _Instance(self.INSTANCE_NAME, client=client)
@@ -586,6 +637,7 @@ def test_reload_success(self):
         self.assertEqual(database._create_time, timestamp)
         self.assertEqual(database._restore_info, restore_info)
         self.assertEqual(database._ddl_statements, tuple(DDL_STATEMENTS))
+        self.assertEqual(database._encryption_config, encryption_config)
 
         api.get_database_ddl.assert_called_once_with(
             self.DATABASE_NAME,

tests/unit/test_instance.py

@@ -479,15 +479,22 @@ def test_database_factory_defaults(self):
 
     def test_database_factory_explicit(self):
         from google.cloud.spanner_v1.database import Database
+        from google.cloud.spanner_admin_database_v1.proto.common_pb2 import (
+            EncryptionConfig,
+        )
         from tests._fixtures import DDL_STATEMENTS
 
         client = _Client(self.PROJECT)
         instance = self._make_one(self.INSTANCE_ID, client, self.CONFIG_NAME)
         DATABASE_ID = "database-id"
         pool = _Pool()
+        encryption_config = EncryptionConfig(kms_key_name="kms_key")
 
         database = instance.database(
-            DATABASE_ID, ddl_statements=DDL_STATEMENTS, pool=pool
+            DATABASE_ID,
+            ddl_statements=DDL_STATEMENTS,
+            pool=pool,
+            encryption_config=encryption_config,
         )
 
         self.assertIsInstance(database, Database)
@@ -496,6 +503,7 @@ def test_database_factory_explicit(self):
         self.assertEqual(list(database.ddl_statements), DDL_STATEMENTS)
         self.assertIs(database._pool, pool)
         self.assertIs(pool._bound, database)
+        self.assertIs(database._encryption_config, encryption_config)
 
     def test_list_databases(self):
         from google.cloud.spanner_admin_database_v1.gapic import database_admin_client