Skip to content
Permalink
Browse files
feat: add support for CMEK (#105)
* feat: add support for creating databases with CMEK

* refactor: use kwargs for EncryptionConfig conversion

* feat: add support for creating backups with CMEK

* feat: add support for restore a database with CMEK

* style: fix lint

* fix: verify that correct encryption type is used when using a key

* test: use non-default encryption for backup tests to test CMEK support

* test: fix encryption assertion

* test: fix encryption type for assertion

* docs: fix docstring types

* docs: update docstring descriptions

Co-authored-by: larkee <larkee@users.noreply.github.com>
  • Loading branch information
larkee and larkee committed Mar 19, 2021
1 parent 801ddc8 commit e990ff70342e7c2e27059e82c8d74cce39eb85d0
@@ -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,24 @@ 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
@@ -71,6 +87,11 @@ def __init__(
self._size_bytes = None
self._state = None
self._referencing_databases = None
self._encryption_info = None
if type(encryption_config) == dict:
self._encryption_config = CreateBackupEncryptionConfig(**encryption_config)
else:
self._encryption_config = encryption_config

@property
def name(self):
@@ -156,6 +177,22 @@ 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.types.EncryptionInfo`
:returns: a class representing the encryption info
"""
return self._encryption_info

@property
def encryption_config(self):
"""Encryption config for this database.
:rtype: :class:`~google.cloud.spanner_admin_instance_v1.types.CreateBackupEncryptionConfig`
:returns: an object representing the encryption config 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.
@@ -207,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(
@@ -215,12 +259,14 @@ 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,
metadata=metadata,
encryption_config=self._encryption_config,
)

future = api.create_backup(request=request, metadata=metadata,)
return future

def exists(self):
@@ -255,6 +301,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.
@@ -47,6 +47,9 @@
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 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,
@@ -108,12 +111,27 @@ 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:`~google.cloud.spanner_admin_database_v1.types.RestoreDatabaseEncryptionConfig`
or :class:`dict`
:param encryption_config:
(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`
"""

_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
@@ -126,6 +144,7 @@ def __init__(
self._earliest_version_time = None
self.log_commit_stats = False
self._logger = logger
self._encryption_config = encryption_config

if pool is None:
pool = BurstyPool()
@@ -242,6 +261,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.types.EncryptionConfig`
:returns: an object representing the encryption config for this database
"""
return self._encryption_config

@property
def ddl_statements(self):
"""DDL Statements used to define database schema.
@@ -325,11 +352,14 @@ 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,
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 +402,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.
@@ -588,8 +619,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
@@ -601,14 +632,26 @@ def restore(self, source):
"""
if source is None:
raise ValueError("Restore source not specified")
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)
future = api.restore_database(
request = RestoreDatabaseRequest(
parent=self._instance.name,
database_id=self.database_id,
backup=source.name,
metadata=metadata,
encryption_config=self._encryption_config,
)
future = api.restore_database(request=request, metadata=metadata,)
return future

def is_ready(self):
@@ -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):
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,11 +384,26 @@ 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:`~google.cloud.spanner_admin_database_v1.types.RestoreDatabaseEncryptionConfig`
or :class:`dict`
:param encryption_config:
(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`
: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
database_id,
self,
ddl_statements=ddl_statements,
pool=pool,
logger=logger,
encryption_config=encryption_config,
)

def list_databases(self, page_size=None):
@@ -408,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):
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
@@ -430,6 +459,14 @@ def backup(self, backup_id, database="", expire_time=None, version_time=None):
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`
:rtype: :class:`~google.cloud.spanner_v1.backup.Backup`
:returns: a backup owned by this instance.
"""
@@ -440,6 +477,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(
@@ -448,6 +486,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):
@@ -738,6 +738,11 @@ def test_create_invalid(self):
op.result()

def test_backup_workflow(self):
from google.cloud.spanner_admin_database_v1 import (
CreateBackupEncryptionConfig,
EncryptionConfig,
RestoreDatabaseEncryptionConfig,
)
from datetime import datetime
from datetime import timedelta
from pytz import UTC
@@ -746,13 +751,17 @@ 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(
backup_id,
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 +780,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 +790,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 +804,9 @@ def test_backup_workflow(self):

metadata = operation.metadata
self.assertEqual(self.database_version_time, metadata.backup_info.version_time)
database.reload()
expected_encryption_config = EncryptionConfig()
self.assertEqual(expected_encryption_config, database.encryption_config)

database.drop()
backup.delete()

0 comments on commit e990ff7

Please sign in to comment.