From 0b7844c31b43948d1977becaf079d692d8bec769 Mon Sep 17 00:00:00 2001 From: larkee Date: Wed, 25 Nov 2020 16:06:51 +1100 Subject: [PATCH 01/11] feat: add support for creating databases with CMEK --- google/cloud/spanner_v1/backup.py | 10 +++++++ google/cloud/spanner_v1/database.py | 34 ++++++++++++++++++++-- google/cloud/spanner_v1/instance.py | 18 +++++++++++- tests/unit/test_backup.py | 14 +++++++++ tests/unit/test_database.py | 44 +++++++++++++++++++++++++++++ tests/unit/test_instance.py | 9 +++++- 6 files changed, 125 insertions(+), 4 deletions(-) diff --git a/google/cloud/spanner_v1/backup.py b/google/cloud/spanner_v1/backup.py index 2277a33fce..350e01f4c9 100644 --- a/google/cloud/spanner_v1/backup.py +++ b/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. diff --git a/google/cloud/spanner_v1/database.py b/google/cloud/spanner_v1/database.py index 92c797b987..cab0b92a76 100644 --- a/google/cloud/spanner_v1/database.py +++ b/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, @@ -108,12 +108,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 @@ -127,6 +140,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() @@ -242,6 +262,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. @@ -330,6 +358,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 @@ -372,6 +401,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. diff --git a/google/cloud/spanner_v1/instance.py b/google/cloud/spanner_v1/instance.py index db729d9527..c1dcc61e7f 100644 --- a/google/cloud/spanner_v1/instance.py +++ b/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. diff --git a/tests/unit/test_backup.py b/tests/unit/test_backup.py index bf6ce68a84..8dda2ecb36 100644 --- a/tests/unit/test_backup.py +++ b/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, diff --git a/tests/unit/test_database.py b/tests/unit/test_database.py index 148bb79b0e..8afc16d2d2 100644 --- a/tests/unit/test_database.py +++ b/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, diff --git a/tests/unit/test_instance.py b/tests/unit/test_instance.py index c1d02c5728..335a9730a2 100644 --- a/tests/unit/test_instance.py +++ b/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 From 1b2eb67b69b0f9e121cb61407dae138366943ef0 Mon Sep 17 00:00:00 2001 From: larkee Date: Wed, 25 Nov 2020 16:18:34 +1100 Subject: [PATCH 02/11] refactor: use kwargs for EncryptionConfig conversion --- google/cloud/spanner_v1/database.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/google/cloud/spanner_v1/database.py b/google/cloud/spanner_v1/database.py index cab0b92a76..215660d776 100644 --- a/google/cloud/spanner_v1/database.py +++ b/google/cloud/spanner_v1/database.py @@ -141,9 +141,7 @@ def __init__( self._logger = logger if type(encryption_config) == dict: - self._encryption_config = EncryptionConfig( - kms_key_name=encryption_config["kms_key_name"] - ) + self._encryption_config = EncryptionConfig(**encryption_config) else: self._encryption_config = encryption_config From 5c69d0925e4adb536f6aa70fd6fc86fbde107218 Mon Sep 17 00:00:00 2001 From: larkee Date: Wed, 25 Nov 2020 19:55:41 +1100 Subject: [PATCH 03/11] feat: add support for creating backups with CMEK --- google/cloud/spanner_v1/backup.py | 31 ++++++++++- google/cloud/spanner_v1/instance.py | 12 ++++- tests/unit/test_backup.py | 84 +++++++++++++++++++++++++++-- tests/unit/test_instance.py | 8 ++- 4 files changed, 126 insertions(+), 9 deletions(-) diff --git a/google/cloud/spanner_v1/backup.py b/google/cloud/spanner_v1/backup.py index 350e01f4c9..57697cfc25 100644 --- a/google/cloud/spanner_v1/backup.py +++ b/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 diff --git a/google/cloud/spanner_v1/instance.py b/google/cloud/spanner_v1/instance.py index c1dcc61e7f..a6c062bddf 100644 --- a/google/cloud/spanner_v1/instance.py +++ b/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): diff --git a/tests/unit/test_backup.py b/tests/unit/test_backup.py index 8dda2ecb36..29b67e6750 100644 --- a/tests/unit/test_backup.py +++ b/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)], ) diff --git a/tests/unit/test_instance.py b/tests/unit/test_instance.py index 335a9730a2..152a032519 100644 --- a/tests/unit/test_instance.py +++ b/tests/unit/test_instance.py @@ -610,15 +610,20 @@ 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) @@ -626,6 +631,7 @@ def test_backup_factory_explicit(self): 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 From 27b5f3056fb47b18ba06757c62a3c78a064374fd Mon Sep 17 00:00:00 2001 From: larkee Date: Mon, 30 Nov 2020 11:54:47 +1100 Subject: [PATCH 04/11] feat: add support for restore a database with CMEK --- google/cloud/spanner_v1/database.py | 27 ++++--- google/cloud/spanner_v1/instance.py | 8 +- tests/unit/test_database.py | 117 +++++++++++++++++++++++----- tests/unit/test_instance.py | 3 +- 4 files changed, 122 insertions(+), 33 deletions(-) diff --git a/google/cloud/spanner_v1/database.py b/google/cloud/spanner_v1/database.py index 215660d776..f244ecdfdf 100644 --- a/google/cloud/spanner_v1/database.py +++ b/google/cloud/spanner_v1/database.py @@ -47,6 +47,8 @@ ) 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 RestoreDatabaseEncryptionConfig +from google.cloud.spanner_admin_database_v1 import RestoreDatabaseRequest from google.cloud.spanner_admin_database_v1 import UpdateDatabaseDdlRequest from google.cloud.spanner_v1 import ( ExecuteSqlRequest, @@ -113,8 +115,9 @@ class Database(object): 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` + If a dict is provided, it must be of the same form as either of the protobuf + messages :class:`~google.cloud.spanner_admin_database_v1.types.EncryptionConfig` + or :class:`~google.cloud.spanner_admin_database_v1.types.RestoreDatabaseEncryptionConfig` """ _spanner_api = None @@ -139,11 +142,7 @@ def __init__( self._earliest_version_time = None self.log_commit_stats = False self._logger = logger - - if type(encryption_config) == dict: - self._encryption_config = EncryptionConfig(**encryption_config) - else: - self._encryption_config = encryption_config + self._encryption_config = encryption_config if pool is None: pool = BurstyPool() @@ -351,6 +350,8 @@ def create(self): db_name = self.database_id if "-" in db_name: db_name = "`%s`" % (db_name,) + if type(self._encryption_config) == dict: + self._encryption_config = EncryptionConfig(**self._encryption_config) request = CreateDatabaseRequest( parent=self._instance.name, @@ -616,8 +617,8 @@ def run_in_transaction(self, func, *args, **kw): def restore(self, source): """Restore from a backup to this database. - :type backup: :class:`~google.cloud.spanner_v1.backup.Backup` - :param backup: the path of the backup being restored from. + :type source: :class:`~google.cloud.spanner_v1.backup.Backup` + :param source: the path of the source being restored from. :rtype: :class:`~google.api_core.operation.Operation` :returns: a future used to poll the status of the create request @@ -631,10 +632,16 @@ def restore(self, source): raise ValueError("Restore source not specified") api = self._instance._client.database_admin_api metadata = _metadata_with_prefix(self.name) - future = api.restore_database( + if type(self._encryption_config) == dict: + self._encryption_config = RestoreDatabaseEncryptionConfig(**self._encryption_config) + request = RestoreDatabaseRequest( parent=self._instance.name, database_id=self.database_id, backup=source.name, + encryption_config=self._encryption_config + ) + future = api.restore_database( + request=request, metadata=metadata, ) return future diff --git a/google/cloud/spanner_v1/instance.py b/google/cloud/spanner_v1/instance.py index a6c062bddf..9f71cc4b4c 100644 --- a/google/cloud/spanner_v1/instance.py +++ b/google/cloud/spanner_v1/instance.py @@ -378,12 +378,14 @@ def database(self, database_id, ddl_statements=(), pool=None, logger=None, encry to stdout. :type encryption_config: - :class:`~google.cloud.spanner_admin_database_v1.types.EncryptionConfig` + :class:`~google.cloud.spanner_admin_database_v1.types.EncryptionConfig` or + :class:`~google.cloud.spanner_admin_database_v1.types.RestoreDatabaseEncryptionConfig` 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 + If a dict is provided, it must be of the same form as either of the protobuf + messages :class:`~google.cloud.spanner_admin_database_v1.types.EncryptionConfig` + or :class:`~google.cloud.spanner_admin_database_v1.types.RestoreDatabaseEncryptionConfig` :rtype: :class:`~google.cloud.spanner_v1.database.Database` :returns: a database owned by this instance. diff --git a/tests/unit/test_database.py b/tests/unit/test_database.py index 8afc16d2d2..341921ac06 100644 --- a/tests/unit/test_database.py +++ b/tests/unit/test_database.py @@ -171,19 +171,6 @@ def test_ctor_w_encryption_config(self): 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 @@ -532,6 +519,7 @@ def test_create_instance_not_found(self): def test_create_success(self): from tests._fixtures import DDL_STATEMENTS from google.cloud.spanner_admin_database_v1 import CreateDatabaseRequest + from google.cloud.spanner_admin_database_v1 import EncryptionConfig op_future = object() client = _Client() @@ -539,8 +527,9 @@ def test_create_success(self): api.create_database.return_value = op_future instance = _Instance(self.INSTANCE_NAME, client=client) pool = _Pool() + encryption_config = EncryptionConfig(kms_key_name="kms_key_name") database = self._make_one( - self.DATABASE_ID, instance, ddl_statements=DDL_STATEMENTS, pool=pool + self.DATABASE_ID, instance, ddl_statements=DDL_STATEMENTS, pool=pool, encryption_config=encryption_config ) future = database.create() @@ -551,7 +540,40 @@ def test_create_success(self): parent=self.INSTANCE_NAME, create_statement="CREATE DATABASE {}".format(self.DATABASE_ID), extra_statements=DDL_STATEMENTS, - encryption_config=None, + encryption_config=encryption_config, + ) + + api.create_database.assert_called_once_with( + request=expected_request, + metadata=[("google-cloud-resource-prefix", database.name)], + ) + + def test_create_success_w_encryption_config_dict(self): + from tests._fixtures import DDL_STATEMENTS + from google.cloud.spanner_admin_database_v1 import CreateDatabaseRequest + from google.cloud.spanner_admin_database_v1 import EncryptionConfig + + op_future = object() + client = _Client() + api = client.database_admin_api = self._make_database_admin_api() + api.create_database.return_value = op_future + instance = _Instance(self.INSTANCE_NAME, client=client) + pool = _Pool() + encryption_config = {"kms_key_name": "kms_key_name"} + database = self._make_one( + self.DATABASE_ID, instance, ddl_statements=DDL_STATEMENTS, pool=pool, encryption_config=encryption_config + ) + + future = database.create() + + self.assertIs(future, op_future) + + expected_encryption_config = EncryptionConfig(**encryption_config) + expected_request = CreateDatabaseRequest( + parent=self.INSTANCE_NAME, + create_statement="CREATE DATABASE {}".format(self.DATABASE_ID), + extra_statements=DDL_STATEMENTS, + encryption_config=expected_encryption_config, ) api.create_database.assert_called_once_with( @@ -1172,6 +1194,7 @@ def test_restore_backup_unspecified(self): def test_restore_grpc_error(self): from google.api_core.exceptions import Unknown + from google.cloud.spanner_admin_database_v1 import RestoreDatabaseRequest client = _Client() api = client.database_admin_api = self._make_database_admin_api() @@ -1184,15 +1207,20 @@ def test_restore_grpc_error(self): with self.assertRaises(Unknown): database.restore(backup) - api.restore_database.assert_called_once_with( + expected_request = RestoreDatabaseRequest( parent=self.INSTANCE_NAME, database_id=self.DATABASE_ID, backup=self.BACKUP_NAME, + ) + + api.restore_database.assert_called_once_with( + request=expected_request, metadata=[("google-cloud-resource-prefix", database.name)], ) def test_restore_not_found(self): from google.api_core.exceptions import NotFound + from google.cloud.spanner_admin_database_v1 import RestoreDatabaseRequest client = _Client() api = client.database_admin_api = self._make_database_admin_api() @@ -1205,31 +1233,84 @@ def test_restore_not_found(self): with self.assertRaises(NotFound): database.restore(backup) - api.restore_database.assert_called_once_with( + expected_request = RestoreDatabaseRequest( parent=self.INSTANCE_NAME, database_id=self.DATABASE_ID, backup=self.BACKUP_NAME, + ) + + api.restore_database.assert_called_once_with( + request=expected_request, metadata=[("google-cloud-resource-prefix", database.name)], ) def test_restore_success(self): + from google.cloud.spanner_admin_database_v1 import RestoreDatabaseEncryptionConfig + from google.cloud.spanner_admin_database_v1 import RestoreDatabaseRequest + op_future = object() client = _Client() api = client.database_admin_api = self._make_database_admin_api() api.restore_database.return_value = op_future instance = _Instance(self.INSTANCE_NAME, client=client) pool = _Pool() - database = self._make_one(self.DATABASE_ID, instance, pool=pool) + encryption_config = RestoreDatabaseEncryptionConfig( + encryption_type=RestoreDatabaseEncryptionConfig.EncryptionType.CUSTOMER_MANAGED_ENCRYPTION, + kms_key_name="kms_key_name" + ) + database = self._make_one(self.DATABASE_ID, instance, pool=pool, encryption_config=encryption_config) backup = _Backup(self.BACKUP_NAME) future = database.restore(backup) self.assertIs(future, op_future) + expected_request = RestoreDatabaseRequest( + parent=self.INSTANCE_NAME, + database_id=self.DATABASE_ID, + backup=self.BACKUP_NAME, + encryption_config=encryption_config + ) + api.restore_database.assert_called_once_with( + request=expected_request, + metadata=[("google-cloud-resource-prefix", database.name)], + ) + + def test_restore_success_w_encryption_config_dict(self): + from google.cloud.spanner_admin_database_v1 import RestoreDatabaseEncryptionConfig + from google.cloud.spanner_admin_database_v1 import RestoreDatabaseRequest + + op_future = object() + client = _Client() + api = client.database_admin_api = self._make_database_admin_api() + api.restore_database.return_value = op_future + instance = _Instance(self.INSTANCE_NAME, client=client) + pool = _Pool() + encryption_config = { + 'encryption_type': RestoreDatabaseEncryptionConfig.EncryptionType.CUSTOMER_MANAGED_ENCRYPTION, + 'kms_key_name': 'kms_key_name' + } + database = self._make_one(self.DATABASE_ID, instance, pool=pool, encryption_config=encryption_config) + backup = _Backup(self.BACKUP_NAME) + + future = database.restore(backup) + + self.assertIs(future, op_future) + + expected_encryption_config = RestoreDatabaseEncryptionConfig( + encryption_type=RestoreDatabaseEncryptionConfig.EncryptionType.CUSTOMER_MANAGED_ENCRYPTION, + kms_key_name="kms_key_name" + ) + expected_request = RestoreDatabaseRequest( parent=self.INSTANCE_NAME, database_id=self.DATABASE_ID, backup=self.BACKUP_NAME, + encryption_config=expected_encryption_config + ) + + api.restore_database.assert_called_once_with( + request=expected_request, metadata=[("google-cloud-resource-prefix", database.name)], ) diff --git a/tests/unit/test_instance.py b/tests/unit/test_instance.py index 152a032519..3207e241d9 100644 --- a/tests/unit/test_instance.py +++ b/tests/unit/test_instance.py @@ -490,7 +490,6 @@ 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 @@ -499,7 +498,7 @@ 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") + encryption_config = {"kms_key_name": "kms_key_name"} database = instance.database( DATABASE_ID, From 2ca28220d279fc5c63a7fe09435988e753a07f70 Mon Sep 17 00:00:00 2001 From: larkee Date: Tue, 23 Feb 2021 21:35:55 +1100 Subject: [PATCH 05/11] style: fix lint --- google/cloud/spanner_v1/backup.py | 13 +++++--- google/cloud/spanner_v1/database.py | 26 ++++++++-------- google/cloud/spanner_v1/instance.py | 23 ++++++++++---- tests/unit/test_backup.py | 48 ++++++++++------------------- tests/unit/test_database.py | 40 ++++++++++++++++-------- tests/unit/test_instance.py | 7 +++-- 6 files changed, 88 insertions(+), 69 deletions(-) diff --git a/google/cloud/spanner_v1/backup.py b/google/cloud/spanner_v1/backup.py index 57697cfc25..0b22a2dd16 100644 --- a/google/cloud/spanner_v1/backup.py +++ b/google/cloud/spanner_v1/backup.py @@ -70,7 +70,13 @@ class Backup(object): """ def __init__( - self, backup_id, instance, database="", expire_time=None, version_time=None, encryption_config=None + self, + backup_id, + instance, + database="", + expire_time=None, + version_time=None, + encryption_config=None, ): self.backup_id = backup_id self._instance = instance @@ -253,10 +259,7 @@ def create(self): encryption_config=self._encryption_config, ) - future = api.create_backup( - request=request, - metadata=metadata, - ) + future = api.create_backup(request=request, metadata=metadata,) return future def exists(self): diff --git a/google/cloud/spanner_v1/database.py b/google/cloud/spanner_v1/database.py index f244ecdfdf..4cd127100d 100644 --- a/google/cloud/spanner_v1/database.py +++ b/google/cloud/spanner_v1/database.py @@ -17,6 +17,7 @@ import copy import functools import grpc +import logging import re import threading @@ -123,13 +124,13 @@ class Database(object): _spanner_api = None def __init__( - self, - database_id, - instance, - ddl_statements=(), - pool=None, - logger=None, - encryption_config=None, + self, + database_id, + instance, + ddl_statements=(), + pool=None, + logger=None, + encryption_config=None, ): self.database_id = database_id self._instance = instance @@ -633,17 +634,16 @@ def restore(self, source): api = self._instance._client.database_admin_api metadata = _metadata_with_prefix(self.name) if type(self._encryption_config) == dict: - self._encryption_config = RestoreDatabaseEncryptionConfig(**self._encryption_config) + self._encryption_config = RestoreDatabaseEncryptionConfig( + **self._encryption_config + ) request = RestoreDatabaseRequest( parent=self._instance.name, database_id=self.database_id, backup=source.name, - encryption_config=self._encryption_config - ) - future = api.restore_database( - request=request, - metadata=metadata, + encryption_config=self._encryption_config, ) + future = api.restore_database(request=request, metadata=metadata,) return future def is_ready(self): diff --git a/google/cloud/spanner_v1/instance.py b/google/cloud/spanner_v1/instance.py index 9f71cc4b4c..1d955164ca 100644 --- a/google/cloud/spanner_v1/instance.py +++ b/google/cloud/spanner_v1/instance.py @@ -357,7 +357,14 @@ def delete(self): api.delete_instance(name=self.name, metadata=metadata) - def database(self, database_id, ddl_statements=(), pool=None, logger=None, encryption_config=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 @@ -390,9 +397,6 @@ def database(self, database_id, ddl_statements=(), pool=None, logger=None, encry :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, @@ -426,7 +430,14 @@ def list_databases(self, page_size=None): ) return page_iter - def backup(self, backup_id, database="", expire_time=None, version_time=None, encryption_config=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 @@ -472,7 +483,7 @@ def backup(self, backup_id, database="", expire_time=None, version_time=None, en database=database, expire_time=expire_time, version_time=version_time, - encryption_config=encryption_config + encryption_config=encryption_config, ) def list_backups(self, filter_="", page_size=None): diff --git a/tests/unit/test_backup.py b/tests/unit/test_backup.py index 29b67e6750..9d3e1cc472 100644 --- a/tests/unit/test_backup.py +++ b/tests/unit/test_backup.py @@ -63,19 +63,20 @@ def test_ctor_defaults(self): 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" + kms_key_name="key_name", ) backup = self._make_one( self.BACKUP_ID, instance, database=self.DATABASE_NAME, expire_time=timestamp, - encryption_config=encryption_config + encryption_config=encryption_config, ) self.assertEqual(backup.backup_id, self.BACKUP_ID) @@ -87,19 +88,17 @@ def test_ctor_non_defaults(self): 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" - } + 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 + encryption_config=encryption_config, ) expected_encryption_config = CreateBackupEncryptionConfig(**encryption_config) @@ -222,7 +221,7 @@ def test_encryption_config_property(self): 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" + kms_key_name="kms_key_name", ) self.assertEqual(backup.encryption_config, expected) @@ -248,14 +247,11 @@ def test_create_grpc_error(self): backup.create() request = CreateBackupRequest( - parent=self.INSTANCE_NAME, - backup_id=self.BACKUP_ID, - backup=backup_pb, + 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)], + request=request, metadata=[("google-cloud-resource-prefix", backup.name)], ) def test_create_already_exists(self): @@ -279,14 +275,11 @@ def test_create_already_exists(self): backup.create() request = CreateBackupRequest( - parent=self.INSTANCE_NAME, - backup_id=self.BACKUP_ID, - backup=backup_pb, + 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)], + request=request, metadata=[("google-cloud-resource-prefix", backup.name)], ) def test_create_instance_not_found(self): @@ -310,14 +303,11 @@ def test_create_instance_not_found(self): backup.create() request = CreateBackupRequest( - parent=self.INSTANCE_NAME, - backup_id=self.BACKUP_ID, - backup=backup_pb, + 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)], + request=request, metadata=[("google-cloud-resource-prefix", backup.name)], ) def test_create_expire_time_not_set(self): @@ -352,17 +342,14 @@ 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" - } + 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 + encryption_config=encryption_config, ) backup_pb = Backup( @@ -379,12 +366,11 @@ def test_create_success(self): parent=self.INSTANCE_NAME, backup_id=self.BACKUP_ID, backup=backup_pb, - encryption_config=expected_encryption_config + encryption_config=expected_encryption_config, ) api.create_backup.assert_called_once_with( - request=request, - metadata=[("google-cloud-resource-prefix", backup.name)], + request=request, metadata=[("google-cloud-resource-prefix", backup.name)], ) def test_exists_grpc_error(self): diff --git a/tests/unit/test_database.py b/tests/unit/test_database.py index 341921ac06..604cecbaa0 100644 --- a/tests/unit/test_database.py +++ b/tests/unit/test_database.py @@ -529,7 +529,11 @@ def test_create_success(self): pool = _Pool() encryption_config = EncryptionConfig(kms_key_name="kms_key_name") database = self._make_one( - self.DATABASE_ID, instance, ddl_statements=DDL_STATEMENTS, pool=pool, encryption_config=encryption_config + self.DATABASE_ID, + instance, + ddl_statements=DDL_STATEMENTS, + pool=pool, + encryption_config=encryption_config, ) future = database.create() @@ -561,7 +565,11 @@ def test_create_success_w_encryption_config_dict(self): pool = _Pool() encryption_config = {"kms_key_name": "kms_key_name"} database = self._make_one( - self.DATABASE_ID, instance, ddl_statements=DDL_STATEMENTS, pool=pool, encryption_config=encryption_config + self.DATABASE_ID, + instance, + ddl_statements=DDL_STATEMENTS, + pool=pool, + encryption_config=encryption_config, ) future = database.create() @@ -1245,7 +1253,9 @@ def test_restore_not_found(self): ) def test_restore_success(self): - from google.cloud.spanner_admin_database_v1 import RestoreDatabaseEncryptionConfig + from google.cloud.spanner_admin_database_v1 import ( + RestoreDatabaseEncryptionConfig, + ) from google.cloud.spanner_admin_database_v1 import RestoreDatabaseRequest op_future = object() @@ -1256,9 +1266,11 @@ def test_restore_success(self): pool = _Pool() encryption_config = RestoreDatabaseEncryptionConfig( encryption_type=RestoreDatabaseEncryptionConfig.EncryptionType.CUSTOMER_MANAGED_ENCRYPTION, - kms_key_name="kms_key_name" + kms_key_name="kms_key_name", + ) + database = self._make_one( + self.DATABASE_ID, instance, pool=pool, encryption_config=encryption_config ) - database = self._make_one(self.DATABASE_ID, instance, pool=pool, encryption_config=encryption_config) backup = _Backup(self.BACKUP_NAME) future = database.restore(backup) @@ -1269,7 +1281,7 @@ def test_restore_success(self): parent=self.INSTANCE_NAME, database_id=self.DATABASE_ID, backup=self.BACKUP_NAME, - encryption_config=encryption_config + encryption_config=encryption_config, ) api.restore_database.assert_called_once_with( @@ -1278,7 +1290,9 @@ def test_restore_success(self): ) def test_restore_success_w_encryption_config_dict(self): - from google.cloud.spanner_admin_database_v1 import RestoreDatabaseEncryptionConfig + from google.cloud.spanner_admin_database_v1 import ( + RestoreDatabaseEncryptionConfig, + ) from google.cloud.spanner_admin_database_v1 import RestoreDatabaseRequest op_future = object() @@ -1288,10 +1302,12 @@ def test_restore_success_w_encryption_config_dict(self): instance = _Instance(self.INSTANCE_NAME, client=client) pool = _Pool() encryption_config = { - 'encryption_type': RestoreDatabaseEncryptionConfig.EncryptionType.CUSTOMER_MANAGED_ENCRYPTION, - 'kms_key_name': 'kms_key_name' + "encryption_type": RestoreDatabaseEncryptionConfig.EncryptionType.CUSTOMER_MANAGED_ENCRYPTION, + "kms_key_name": "kms_key_name", } - database = self._make_one(self.DATABASE_ID, instance, pool=pool, encryption_config=encryption_config) + database = self._make_one( + self.DATABASE_ID, instance, pool=pool, encryption_config=encryption_config + ) backup = _Backup(self.BACKUP_NAME) future = database.restore(backup) @@ -1300,13 +1316,13 @@ def test_restore_success_w_encryption_config_dict(self): expected_encryption_config = RestoreDatabaseEncryptionConfig( encryption_type=RestoreDatabaseEncryptionConfig.EncryptionType.CUSTOMER_MANAGED_ENCRYPTION, - kms_key_name="kms_key_name" + kms_key_name="kms_key_name", ) expected_request = RestoreDatabaseRequest( parent=self.INSTANCE_NAME, database_id=self.DATABASE_ID, backup=self.BACKUP_NAME, - encryption_config=expected_encryption_config + encryption_config=expected_encryption_config, ) api.restore_database.assert_called_once_with( diff --git a/tests/unit/test_instance.py b/tests/unit/test_instance.py index 3207e241d9..2ed777b25b 100644 --- a/tests/unit/test_instance.py +++ b/tests/unit/test_instance.py @@ -618,11 +618,14 @@ def test_backup_factory_explicit(self): timestamp = datetime.datetime.utcnow().replace(tzinfo=UTC) encryption_config = CreateBackupEncryptionConfig( encryption_type=CreateBackupEncryptionConfig.EncryptionType.CUSTOMER_MANAGED_ENCRYPTION, - kms_key_name="kms_key_name" + kms_key_name="kms_key_name", ) backup = instance.backup( - BACKUP_ID, database=DATABASE_NAME, expire_time=timestamp, encryption_config=encryption_config + BACKUP_ID, + database=DATABASE_NAME, + expire_time=timestamp, + encryption_config=encryption_config, ) self.assertIsInstance(backup, Backup) From 50e2d8f3b2e85f829583dc5bf7a2f8047dc00e93 Mon Sep 17 00:00:00 2001 From: larkee Date: Fri, 12 Mar 2021 20:33:35 +1100 Subject: [PATCH 06/11] fix: verify that correct encryption type is used when using a key --- google/cloud/spanner_v1/backup.py | 7 +++++++ google/cloud/spanner_v1/database.py | 11 +++++++++-- tests/unit/test_backup.py | 21 +++++++++++++++++++++ tests/unit/test_database.py | 20 ++++++++++++++++++++ 4 files changed, 57 insertions(+), 2 deletions(-) diff --git a/google/cloud/spanner_v1/backup.py b/google/cloud/spanner_v1/backup.py index 0b22a2dd16..6be72e8833 100644 --- a/google/cloud/spanner_v1/backup.py +++ b/google/cloud/spanner_v1/backup.py @@ -244,6 +244,13 @@ def create(self): raise ValueError("expire_time not set") if not self._database: raise ValueError("database not set") + if ( + self.encryption_config + and self.encryption_config.kms_key_name + and self.encryption_config.encryption_type + != CreateBackupEncryptionConfig.EncryptionType.CUSTOMER_MANAGED_ENCRYPTION + ): + raise ValueError("kms_key_name only used with CUSTOMER_MANAGED_ENCRYPTION") api = self._instance._client.database_admin_api metadata = _metadata_with_prefix(self.name) backup = BackupPB( diff --git a/google/cloud/spanner_v1/database.py b/google/cloud/spanner_v1/database.py index 4cd127100d..b42433bc03 100644 --- a/google/cloud/spanner_v1/database.py +++ b/google/cloud/spanner_v1/database.py @@ -631,12 +631,19 @@ def restore(self, source): """ if source is None: raise ValueError("Restore source not specified") - api = self._instance._client.database_admin_api - metadata = _metadata_with_prefix(self.name) if type(self._encryption_config) == dict: self._encryption_config = RestoreDatabaseEncryptionConfig( **self._encryption_config ) + if ( + self.encryption_config + and self.encryption_config.kms_key_name + and self.encryption_config.encryption_type + != RestoreDatabaseEncryptionConfig.EncryptionType.CUSTOMER_MANAGED_ENCRYPTION + ): + raise ValueError("kms_key_name only used with CUSTOMER_MANAGED_ENCRYPTION") + api = self._instance._client.database_admin_api + metadata = _metadata_with_prefix(self.name) request = RestoreDatabaseRequest( parent=self._instance.name, database_id=self.database_id, diff --git a/tests/unit/test_backup.py b/tests/unit/test_backup.py index 9d3e1cc472..335ccb564b 100644 --- a/tests/unit/test_backup.py +++ b/tests/unit/test_backup.py @@ -373,6 +373,27 @@ def test_create_success(self): request=request, metadata=[("google-cloud-resource-prefix", backup.name)], ) + def test_create_w_invalid_encryption_config(self): + from google.cloud.spanner_admin_database_v1 import CreateBackupEncryptionConfig + + client = _Client() + instance = _Instance(self.INSTANCE_NAME, client=client) + expire_timestamp = self._make_timestamp() + encryption_config = { + "encryption_type": CreateBackupEncryptionConfig.EncryptionType.GOOGLE_DEFAULT_ENCRYPTION, + "kms_key_name": "key_name", + } + backup = self._make_one( + self.BACKUP_ID, + instance, + database=self.DATABASE_NAME, + expire_time=expire_timestamp, + encryption_config=encryption_config, + ) + + with self.assertRaises(ValueError): + backup.create() + def test_exists_grpc_error(self): from google.api_core.exceptions import Unknown diff --git a/tests/unit/test_database.py b/tests/unit/test_database.py index 604cecbaa0..4bd7f7659e 100644 --- a/tests/unit/test_database.py +++ b/tests/unit/test_database.py @@ -1330,6 +1330,26 @@ def test_restore_success_w_encryption_config_dict(self): metadata=[("google-cloud-resource-prefix", database.name)], ) + def test_restore_w_invalid_encryption_config_dict(self): + from google.cloud.spanner_admin_database_v1 import ( + RestoreDatabaseEncryptionConfig, + ) + + client = _Client() + instance = _Instance(self.INSTANCE_NAME, client=client) + pool = _Pool() + encryption_config = { + "encryption_type": RestoreDatabaseEncryptionConfig.EncryptionType.GOOGLE_DEFAULT_ENCRYPTION, + "kms_key_name": "kms_key_name", + } + database = self._make_one( + self.DATABASE_ID, instance, pool=pool, encryption_config=encryption_config + ) + backup = _Backup(self.BACKUP_NAME) + + with self.assertRaises(ValueError): + database.restore(backup) + def test_is_ready(self): from google.cloud.spanner_admin_database_v1 import Database From ce8189c62676902104bbd462b6a72ffcd7cfc366 Mon Sep 17 00:00:00 2001 From: larkee Date: Wed, 17 Mar 2021 14:25:07 +1100 Subject: [PATCH 07/11] test: use non-default encryption for backup tests to test CMEK support --- tests/system/test_system.py | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/tests/system/test_system.py b/tests/system/test_system.py index 575f79746e..4d6e9cd268 100644 --- a/tests/system/test_system.py +++ b/tests/system/test_system.py @@ -738,6 +738,10 @@ def test_create_invalid(self): op.result() def test_backup_workflow(self): + from google.cloud.spanner_admin_database_v1 import ( + CreateBackupEncryptionConfig, + RestoreDatabaseEncryptionConfig, + ) from datetime import datetime from datetime import timedelta from pytz import UTC @@ -746,6 +750,9 @@ def test_backup_workflow(self): backup_id = "backup_id" + unique_resource_id("_") expire_time = datetime.utcnow() + timedelta(days=3) expire_time = expire_time.replace(tzinfo=UTC) + encryption_config = CreateBackupEncryptionConfig( + encryption_type=CreateBackupEncryptionConfig.EncryptionType.GOOGLE_DEFAULT_ENCRYPTION, + ) # Create backup. backup = instance.backup( @@ -753,6 +760,7 @@ def test_backup_workflow(self): database=self._db, expire_time=expire_time, version_time=self.database_version_time, + encryption_config=encryption_config, ) operation = backup.create() self.to_delete.append(backup) @@ -771,6 +779,7 @@ def test_backup_workflow(self): self.assertEqual(self.database_version_time, backup.version_time) self.assertIsNotNone(backup.size_bytes) self.assertIsNotNone(backup.state) + self.assertEqual(encryption_config, backup.encryption_config) # Update with valid argument. valid_expire_time = datetime.utcnow() + timedelta(days=7) @@ -780,7 +789,10 @@ def test_backup_workflow(self): # Restore database to same instance. restored_id = "restored_db" + unique_resource_id("_") - database = instance.database(restored_id) + encryption_config = RestoreDatabaseEncryptionConfig( + encryption_type=RestoreDatabaseEncryptionConfig.EncryptionType.GOOGLE_DEFAULT_ENCRYPTION, + ) + database = instance.database(restored_id, encryption_config=encryption_config) self.to_drop.append(database) operation = database.restore(source=backup) restored_db = operation.result() @@ -791,6 +803,7 @@ def test_backup_workflow(self): metadata = operation.metadata self.assertEqual(self.database_version_time, metadata.backup_info.version_time) + self.assertEqual(encryption_config, backup.encryption_config) database.drop() backup.delete() From d233412f66c9ab2ac82e2e8d8623ee382c47f8b8 Mon Sep 17 00:00:00 2001 From: larkee Date: Wed, 17 Mar 2021 15:38:19 +1100 Subject: [PATCH 08/11] test: fix encryption assertion --- tests/system/test_system.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/system/test_system.py b/tests/system/test_system.py index 4d6e9cd268..f3c74c321e 100644 --- a/tests/system/test_system.py +++ b/tests/system/test_system.py @@ -803,7 +803,8 @@ def test_backup_workflow(self): metadata = operation.metadata self.assertEqual(self.database_version_time, metadata.backup_info.version_time) - self.assertEqual(encryption_config, backup.encryption_config) + database.reload() + self.assertEqual(encryption_config, database.encryption_config) database.drop() backup.delete() From 24daedea7bd083f3e9a94ea65e69b5e0e00a80ce Mon Sep 17 00:00:00 2001 From: larkee Date: Thu, 18 Mar 2021 15:36:40 +1100 Subject: [PATCH 09/11] test: fix encryption type for assertion --- tests/system/test_system.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/system/test_system.py b/tests/system/test_system.py index f3c74c321e..8be207ef06 100644 --- a/tests/system/test_system.py +++ b/tests/system/test_system.py @@ -740,6 +740,7 @@ def test_create_invalid(self): def test_backup_workflow(self): from google.cloud.spanner_admin_database_v1 import ( CreateBackupEncryptionConfig, + EncryptionConfig, RestoreDatabaseEncryptionConfig, ) from datetime import datetime @@ -804,7 +805,8 @@ def test_backup_workflow(self): metadata = operation.metadata self.assertEqual(self.database_version_time, metadata.backup_info.version_time) database.reload() - self.assertEqual(encryption_config, database.encryption_config) + expected_encryption_config = EncryptionConfig() + self.assertEqual(expected_encryption_config, database.encryption_config) database.drop() backup.delete() From 311995c7401d2b3a5c01e50ff6eabba7938d6634 Mon Sep 17 00:00:00 2001 From: larkee Date: Thu, 18 Mar 2021 16:58:41 +1100 Subject: [PATCH 10/11] docs: fix docstring types --- google/cloud/spanner_v1/backup.py | 4 ++-- google/cloud/spanner_v1/database.py | 3 ++- google/cloud/spanner_v1/instance.py | 4 ++-- 3 files changed, 6 insertions(+), 5 deletions(-) diff --git a/google/cloud/spanner_v1/backup.py b/google/cloud/spanner_v1/backup.py index beb2ea6a87..c702266a66 100644 --- a/google/cloud/spanner_v1/backup.py +++ b/google/cloud/spanner_v1/backup.py @@ -180,7 +180,7 @@ def referencing_databases(self): @property def encryption_info(self): """Encryption info for this backup. - :rtype: :class:`~google.clod.spanner_admin_database_v1.proto.common_pb2.EncryptionInfo` + :rtype: :class:`~google.clod.spanner_admin_database_v1.types.EncryptionInfo` :returns: a class representing the encryption info """ return self._encryption_info @@ -188,7 +188,7 @@ def encryption_info(self): @property def encryption_config(self): """Encryption config for this database. - :rtype: :class:`~google.cloud.spanner_admin_instance_v1.CreateBackupEncryptionConfig` + :rtype: :class:`~google.cloud.spanner_admin_instance_v1.types.CreateBackupEncryptionConfig` :returns: an object representing the restore info for this database """ return self._encryption_config diff --git a/google/cloud/spanner_v1/database.py b/google/cloud/spanner_v1/database.py index ac53248946..f77458dca9 100644 --- a/google/cloud/spanner_v1/database.py +++ b/google/cloud/spanner_v1/database.py @@ -113,6 +113,7 @@ class Database(object): to stdout. :type encryption_config: :class:`~google.cloud.spanner_admin_database_v1.types.EncryptionConfig` + or :class:`~google.cloud.spanner_admin_database_v1.types.RestoreDatabaseEncryptionConfig` or :class:`dict` :param encryption_config: (Optional) Encryption information about the database. @@ -263,7 +264,7 @@ def earliest_version_time(self): @property def encryption_config(self): """Encryption config for this database. - :rtype: :class:`~google.cloud.spanner_admin_instance_v1.proto.common_pb2.EncryptionConfig` + :rtype: :class:`~google.cloud.spanner_admin_instance_v1.types.EncryptionConfig` :returns: an object representing the restore info for this database """ return self._encryption_config diff --git a/google/cloud/spanner_v1/instance.py b/google/cloud/spanner_v1/instance.py index b6f7cbedff..a80e1bed36 100644 --- a/google/cloud/spanner_v1/instance.py +++ b/google/cloud/spanner_v1/instance.py @@ -385,8 +385,8 @@ def database( to stdout. :type encryption_config: - :class:`~google.cloud.spanner_admin_database_v1.types.EncryptionConfig` or - :class:`~google.cloud.spanner_admin_database_v1.types.RestoreDatabaseEncryptionConfig` + :class:`~google.cloud.spanner_admin_database_v1.types.EncryptionConfig` + or :class:`~google.cloud.spanner_admin_database_v1.types.RestoreDatabaseEncryptionConfig` or :class:`dict` :param encryption_config: (Optional) Encryption information about the database. From ad06371e2db051f43399cfa0117bc3896ea7236e Mon Sep 17 00:00:00 2001 From: larkee Date: Thu, 18 Mar 2021 17:28:34 +1100 Subject: [PATCH 11/11] docs: update docstring descriptions --- google/cloud/spanner_v1/backup.py | 2 +- google/cloud/spanner_v1/database.py | 4 ++-- google/cloud/spanner_v1/instance.py | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/google/cloud/spanner_v1/backup.py b/google/cloud/spanner_v1/backup.py index c702266a66..9068816705 100644 --- a/google/cloud/spanner_v1/backup.py +++ b/google/cloud/spanner_v1/backup.py @@ -189,7 +189,7 @@ def encryption_info(self): def encryption_config(self): """Encryption config for this database. :rtype: :class:`~google.cloud.spanner_admin_instance_v1.types.CreateBackupEncryptionConfig` - :returns: an object representing the restore info for this database + :returns: an object representing the encryption config for this database """ return self._encryption_config diff --git a/google/cloud/spanner_v1/database.py b/google/cloud/spanner_v1/database.py index f77458dca9..3b367445e9 100644 --- a/google/cloud/spanner_v1/database.py +++ b/google/cloud/spanner_v1/database.py @@ -116,7 +116,7 @@ class Database(object): or :class:`~google.cloud.spanner_admin_database_v1.types.RestoreDatabaseEncryptionConfig` or :class:`dict` :param encryption_config: - (Optional) Encryption information about the database. + (Optional) Encryption configuration for the database. If a dict is provided, it must be of the same form as either of the protobuf messages :class:`~google.cloud.spanner_admin_database_v1.types.EncryptionConfig` or :class:`~google.cloud.spanner_admin_database_v1.types.RestoreDatabaseEncryptionConfig` @@ -265,7 +265,7 @@ def earliest_version_time(self): def encryption_config(self): """Encryption config for this database. :rtype: :class:`~google.cloud.spanner_admin_instance_v1.types.EncryptionConfig` - :returns: an object representing the restore info for this database + :returns: an object representing the encryption config for this database """ return self._encryption_config diff --git a/google/cloud/spanner_v1/instance.py b/google/cloud/spanner_v1/instance.py index a80e1bed36..5a9cf95f5a 100644 --- a/google/cloud/spanner_v1/instance.py +++ b/google/cloud/spanner_v1/instance.py @@ -389,7 +389,7 @@ def database( or :class:`~google.cloud.spanner_admin_database_v1.types.RestoreDatabaseEncryptionConfig` or :class:`dict` :param encryption_config: - (Optional) Encryption information about the database. + (Optional) Encryption configuration for the database. If a dict is provided, it must be of the same form as either of the protobuf messages :class:`~google.cloud.spanner_admin_database_v1.types.EncryptionConfig` or :class:`~google.cloud.spanner_admin_database_v1.types.RestoreDatabaseEncryptionConfig`