Skip to content
Permalink
Browse files
feat: add samples for CMEK support (#275)
* feat: add samples for CMEK support

* test: fix backups cleanup

* test: correctly use database id for cmek restore

* test: add clean up for databases

* refactor: remove version time from sample

* refactor: use user-provided key for creating encrypted backup message

Co-authored-by: larkee <larkee@users.noreply.github.com>
  • Loading branch information
larkee and larkee committed Mar 30, 2021
1 parent 75f8340 commit f8d9bd33e04675a8dca148c2fae4a9133beebbca
Showing with 153 additions and 1 deletion.
  1. +72 −0 samples/samples/backup_sample.py
  2. +33 −1 samples/samples/backup_sample_test.py
  3. +37 −0 samples/samples/snippets.py
  4. +11 −0 samples/samples/snippets_test.py
@@ -55,6 +55,42 @@ def create_backup(instance_id, database_id, backup_id, version_time):

# [END spanner_create_backup]

# [START spanner_create_backup_with_encryption_key]
def create_backup_with_encryption_key(instance_id, database_id, backup_id, kms_key_name):
"""Creates a backup for a database using a Customer Managed Encryption Key (CMEK)."""
from google.cloud.spanner_admin_database_v1 import CreateBackupEncryptionConfig

spanner_client = spanner.Client()
instance = spanner_client.instance(instance_id)
database = instance.database(database_id)

# Create a backup
expire_time = datetime.utcnow() + timedelta(days=14)
encryption_config = {
'encryption_type': CreateBackupEncryptionConfig.EncryptionType.CUSTOMER_MANAGED_ENCRYPTION,
'kms_key_name': kms_key_name,
}
backup = instance.backup(backup_id, database=database, expire_time=expire_time, encryption_config=encryption_config)
operation = backup.create()

# Wait for backup operation to complete.
operation.result(1200)

# Verify that the backup is ready.
backup.reload()
assert backup.is_ready() is True

# Get the name, create time, backup size and encryption key.
backup.reload()
print(
"Backup {} of size {} bytes was created at {} using encryption key {}".format(
backup.name, backup.size_bytes, backup.create_time, kms_key_name
)
)


# [END spanner_create_backup_with_encryption_key]


# [START spanner_restore_backup]
def restore_database(instance_id, new_database_id, backup_id):
@@ -87,6 +123,42 @@ def restore_database(instance_id, new_database_id, backup_id):
# [END spanner_restore_backup]


# [START spanner_restore_backup_with_encryption_key]
def restore_database_with_encryption_key(instance_id, new_database_id, backup_id, kms_key_name):
"""Restores a database from a backup using a Customer Managed Encryption Key (CMEK)."""
from google.cloud.spanner_admin_database_v1 import RestoreDatabaseEncryptionConfig

spanner_client = spanner.Client()
instance = spanner_client.instance(instance_id)

# Start restoring an existing backup to a new database.
backup = instance.backup(backup_id)
encryption_config = {
'encryption_type': RestoreDatabaseEncryptionConfig.EncryptionType.CUSTOMER_MANAGED_ENCRYPTION,
'kms_key_name': kms_key_name,
}
new_database = instance.database(new_database_id, encryption_config=encryption_config)
operation = new_database.restore(backup)

# Wait for restore operation to complete.
operation.result(1600)

# Newly created database has restore information.
new_database.reload()
restore_info = new_database.restore_info
print(
"Database {} restored to {} from backup {} with using encryption key {}.".format(
restore_info.backup_info.source_database,
new_database_id,
restore_info.backup_info.backup,
new_database.encryption_config.kms_key_name,
)
)


# [END spanner_restore_backup_with_encryption_key]


# [START spanner_cancel_backup_create]
def cancel_backup(instance_id, database_id, backup_id):
spanner_client = spanner.Client()
@@ -38,9 +38,11 @@ def unique_backup_id():

INSTANCE_ID = unique_instance_id()
DATABASE_ID = unique_database_id()
RETENTION_DATABASE_ID = unique_database_id()
RESTORE_DB_ID = unique_database_id()
BACKUP_ID = unique_backup_id()
CMEK_RESTORE_DB_ID = unique_database_id()
CMEK_BACKUP_ID = unique_backup_id()
RETENTION_DATABASE_ID = unique_database_id()
RETENTION_PERIOD = "7d"


@@ -54,6 +56,12 @@ def spanner_instance():
op = instance.create()
op.result(120) # block until completion
yield instance
for database_pb in instance.list_databases():
database = instance.database(database_pb.name.split("/")[-1])
database.drop()
for backup_pb in instance.list_backups():
backup = instance.backup(backup_pb.name.split("/")[-1])
backup.delete()
instance.delete()


@@ -77,6 +85,16 @@ def test_create_backup(capsys, database):
assert BACKUP_ID in out


def test_create_backup_with_encryption_key(capsys, spanner_instance, database):
kms_key_name = "projects/{}/locations/{}/keyRings/{}/cryptoKeys/{}".format(
spanner_instance._client.project, "us-central1", "spanner-test-keyring", "spanner-test-cmek"
)
backup_sample.create_backup_with_encryption_key(INSTANCE_ID, DATABASE_ID, CMEK_BACKUP_ID, kms_key_name)
out, _ = capsys.readouterr()
assert CMEK_BACKUP_ID in out
assert kms_key_name in out


# Depends on test_create_backup having run first
@RetryErrors(exception=DeadlineExceeded, max_tries=2)
def test_restore_database(capsys):
@@ -87,6 +105,20 @@ def test_restore_database(capsys):
assert BACKUP_ID in out


# Depends on test_create_backup having run first
@RetryErrors(exception=DeadlineExceeded, max_tries=2)
def test_restore_database_with_encryption_key(capsys, spanner_instance):
kms_key_name = "projects/{}/locations/{}/keyRings/{}/cryptoKeys/{}".format(
spanner_instance._client.project, "us-central1", "spanner-test-keyring", "spanner-test-cmek"
)
backup_sample.restore_database_with_encryption_key(INSTANCE_ID, CMEK_RESTORE_DB_ID, CMEK_BACKUP_ID, kms_key_name)
out, _ = capsys.readouterr()
assert (DATABASE_ID + " restored to ") in out
assert (CMEK_RESTORE_DB_ID + " from backup ") in out
assert CMEK_BACKUP_ID in out
assert kms_key_name in out


# Depends on test_create_backup having run first
def test_list_backup_operations(capsys, spanner_instance):
backup_sample.list_backup_operations(INSTANCE_ID, DATABASE_ID)
@@ -92,6 +92,43 @@ def create_database(instance_id, database_id):
# [END spanner_create_database]


# [START spanner_create_database_with_encryption_key]
def create_database_with_encryption_key(instance_id, database_id, kms_key_name):
"""Creates a database with tables using a Customer Managed Encryption Key (CMEK)."""
spanner_client = spanner.Client()
instance = spanner_client.instance(instance_id)

database = instance.database(
database_id,
ddl_statements=[
"""CREATE TABLE Singers (
SingerId INT64 NOT NULL,
FirstName STRING(1024),
LastName STRING(1024),
SingerInfo BYTES(MAX)
) PRIMARY KEY (SingerId)""",
"""CREATE TABLE Albums (
SingerId INT64 NOT NULL,
AlbumId INT64 NOT NULL,
AlbumTitle STRING(MAX)
) PRIMARY KEY (SingerId, AlbumId),
INTERLEAVE IN PARENT Singers ON DELETE CASCADE""",
],
encryption_config={'kms_key_name': kms_key_name},
)

operation = database.create()

print("Waiting for operation to complete...")
operation.result(120)

print("Database {} created with encryption key {}".format(
database.name, database.encryption_config.kms_key_name))


# [END spanner_create_database_with_encryption_key]


# [START spanner_insert_data]
def insert_data(instance_id, database_id):
"""Inserts sample data into the given database.
@@ -33,6 +33,7 @@ def unique_database_id():

INSTANCE_ID = unique_instance_id()
DATABASE_ID = unique_database_id()
CMEK_DATABASE_ID = unique_database_id()


@pytest.fixture(scope="module")
@@ -63,6 +64,16 @@ def test_create_database(database):
database.reload()


def test_create_database_with_encryption_config(capsys, spanner_instance):
kms_key_name = "projects/{}/locations/{}/keyRings/{}/cryptoKeys/{}".format(
spanner_instance._client.project, "us-central1", "spanner-test-keyring", "spanner-test-cmek"
)
snippets.create_database_with_encryption_key(INSTANCE_ID, CMEK_DATABASE_ID, kms_key_name)
out, _ = capsys.readouterr()
assert CMEK_DATABASE_ID in out
assert kms_key_name in out


def test_insert_data(capsys):
snippets.insert_data(INSTANCE_ID, DATABASE_ID)
out, _ = capsys.readouterr()

0 comments on commit f8d9bd3

Please sign in to comment.