feat: add support for creating databases with CMEK

Author: larkee

google/cloud/spanner_v1/backup.py

@@ -71,6 +71,7 @@ def __init__(
         self._size_bytes = None
         self._state = None
         self._referencing_databases = None
+        self._encryption_info = None
 
     @property
     def name(self):
@@ -156,6 +157,14 @@ 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.
@@ -255,6 +264,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

@@ -17,7 +17,6 @@
 import copy
 import functools
 import grpc
-import logging
 import re
 import threading
 
@@ -47,6 +46,7 @@
     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 UpdateDatabaseDdlRequest
 from google.cloud.spanner_v1 import ExecuteSqlRequest
 from google.cloud.spanner_v1 import (
@@ -102,12 +102,25 @@ 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:`dict`
+    :param encryption_config:
+        (Optional) Encryption information about the database.
+        If a dict is provided, it must be of the same form as the protobuf
+        message :class:`~google.cloud.spanner_admin_database_v1.types.EncryptionConfig`
     """
 
     _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
@@ -121,6 +134,13 @@ def __init__(
         self.log_commit_stats = False
         self._logger = logger
 
+        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()
 
@@ -236,6 +256,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.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.
@@ -324,6 +352,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,
         )
         future = api.create_database(request=request, metadata=metadata)
         return future
@@ -366,6 +395,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.

google/cloud/spanner_v1/instance.py

@@ -357,7 +357,7 @@ 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,12 +377,28 @@ 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:`dict`
+        :param encryption_config:
+            (Optional) Encryption information about the database.
+            If a dict is provided, it must be of the same form as the protobuf
+            message :class:`~google.cloud.spanner_admin_database_v1.types.EncryptionConfig
+
         :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
         )
+        return Database(
+            database_id,
+            self,
+            ddl_statements=ddl_statements,
+            pool=pool,
+            logger=logger,
+            encryption_config=encryption_config,
+        )
 
     def list_databases(self, page_size=None):
         """List databases for the instance.

tests/unit/test_backup.py

@@ -170,6 +170,16 @@ 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 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
@@ -442,8 +452,10 @@ def test_reload_not_found(self):
 
     def test_reload_success(self):
         from google.cloud.spanner_admin_database_v1 import Backup
+        from google.cloud.spanner_admin_database_v1 import EncryptionInfo
 
         timestamp = self._make_timestamp()
+        encryption_info = EncryptionInfo(kms_key_version="kms_key_version")
 
         client = _Client()
         backup_pb = Backup(
@@ -455,6 +467,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
@@ -470,6 +483,7 @@ def test_reload_success(self):
         self.assertEqual(backup.size_bytes, 10)
         self.assertEqual(backup.state, Backup.State.CREATING)
         self.assertEqual(backup.referencing_databases, [])
+        self.assertEqual(backup.encryption_info, encryption_info)
 
         api.get_backup.assert_called_once_with(
             name=self.BACKUP_NAME,

tests/unit/test_database.py

@@ -159,6 +159,31 @@ def test_ctor_w_explicit_logger(self):
         self.assertFalse(database.log_commit_stats)
         self.assertEqual(database._logger, logger)
 
+    def test_ctor_w_encryption_config(self):
+        from google.cloud.spanner_admin_database_v1 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._encryption_config, encryption_config)
+
+    def test_ctor_w_encryption_config_dict(self):
+        from google.cloud.spanner_admin_database_v1 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._encryption_config, encryption_config)
+
     def test_from_pb_bad_database_name(self):
         from google.cloud.spanner_admin_database_v1 import Database
 
@@ -295,6 +320,17 @@ def test_logger_property_custom(self):
         logger = database._logger = mock.create_autospec(logging.Logger, instance=True)
         self.assertEqual(database.logger, logger)
 
+    def test_encryption_config(self):
+        from google.cloud.spanner_admin_database_v1 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()
@@ -432,6 +468,7 @@ def test_create_grpc_error(self):
             parent=self.INSTANCE_NAME,
             create_statement="CREATE DATABASE {}".format(self.DATABASE_ID),
             extra_statements=[],
+            encryption_config=None,
         )
 
         api.create_database.assert_called_once_with(
@@ -458,6 +495,7 @@ def test_create_already_exists(self):
             parent=self.INSTANCE_NAME,
             create_statement="CREATE DATABASE `{}`".format(DATABASE_ID_HYPHEN),
             extra_statements=[],
+            encryption_config=None,
         )
 
         api.create_database.assert_called_once_with(
@@ -483,6 +521,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,
         )
 
         api.create_database.assert_called_once_with(
@@ -512,6 +551,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,
         )
 
         api.create_database.assert_called_once_with(
@@ -611,6 +651,7 @@ def test_reload_not_found(self):
 
     def test_reload_success(self):
         from google.cloud.spanner_admin_database_v1 import Database
+        from google.cloud.spanner_admin_database_v1 import EncryptionConfig
         from google.cloud.spanner_admin_database_v1 import GetDatabaseDdlResponse
         from google.cloud.spanner_admin_database_v1 import RestoreInfo
         from google.cloud._helpers import _datetime_to_pb_timestamp
@@ -621,6 +662,7 @@ def test_reload_success(self):
 
         client = _Client()
         ddl_pb = 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 = Database(
@@ -629,6 +671,7 @@ def test_reload_success(self):
             restore_info=restore_info,
             version_retention_period="1d",
             earliest_version_time=_datetime_to_pb_timestamp(timestamp),
+            encryption_config=encryption_config,
         )
         api.get_database.return_value = db_pb
         instance = _Instance(self.INSTANCE_NAME, client=client)
@@ -642,6 +685,7 @@ def test_reload_success(self):
         self.assertEqual(database._version_retention_period, "1d")
         self.assertEqual(database._earliest_version_time, timestamp)
         self.assertEqual(database._ddl_statements, tuple(DDL_STATEMENTS))
+        self.assertEqual(database._encryption_config, encryption_config)
 
         api.get_database_ddl.assert_called_once_with(
             database=self.DATABASE_NAME,

tests/unit/test_instance.py

@@ -490,6 +490,7 @@ def test_database_factory_defaults(self):
 
     def test_database_factory_explicit(self):
         from logging import Logger
+        from google.cloud.spanner_admin_database_v1 import EncryptionConfig
         from google.cloud.spanner_v1.database import Database
         from tests._fixtures import DDL_STATEMENTS
 
@@ -498,9 +499,14 @@ def test_database_factory_explicit(self):
         DATABASE_ID = "database-id"
         pool = _Pool()
         logger = mock.create_autospec(Logger, instance=True)
+        encryption_config = EncryptionConfig(kms_key_name="kms_key")
 
         database = instance.database(
-            DATABASE_ID, ddl_statements=DDL_STATEMENTS, pool=pool, logger=logger
+            DATABASE_ID,
+            ddl_statements=DDL_STATEMENTS,
+            pool=pool,
+            logger=logger,
+            encryption_config=encryption_config,
         )
 
         self.assertIsInstance(database, Database)
@@ -510,6 +516,7 @@ def test_database_factory_explicit(self):
         self.assertIs(database._pool, pool)
         self.assertIs(database._logger, logger)
         self.assertIs(pool._bound, database)
+        self.assertIs(database._encryption_config, encryption_config)
 
     def test_list_databases(self):
         from google.cloud.spanner_admin_database_v1 import Database as DatabasePB