From f8d9bd33e04675a8dca148c2fae4a9133beebbca Mon Sep 17 00:00:00 2001 From: larkee <31196561+larkee@users.noreply.github.com> Date: Tue, 30 Mar 2021 16:10:22 +1100 Subject: [PATCH] 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 --- samples/samples/backup_sample.py | 72 +++++++++++++++++++++++++++ samples/samples/backup_sample_test.py | 34 ++++++++++++- samples/samples/snippets.py | 37 ++++++++++++++ samples/samples/snippets_test.py | 11 ++++ 4 files changed, 153 insertions(+), 1 deletion(-) diff --git a/samples/samples/backup_sample.py b/samples/samples/backup_sample.py index f0d5ce363d..196cfbe04b 100644 --- a/samples/samples/backup_sample.py +++ b/samples/samples/backup_sample.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() diff --git a/samples/samples/backup_sample_test.py b/samples/samples/backup_sample_test.py index 7118d98bed..8d1d95ff51 100644 --- a/samples/samples/backup_sample_test.py +++ b/samples/samples/backup_sample_test.py @@ -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) diff --git a/samples/samples/snippets.py b/samples/samples/snippets.py index 9a94e85a9b..10fc6413c2 100644 --- a/samples/samples/snippets.py +++ b/samples/samples/snippets.py @@ -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. diff --git a/samples/samples/snippets_test.py b/samples/samples/snippets_test.py index ee8c6ebe23..28d13fa330 100644 --- a/samples/samples/snippets_test.py +++ b/samples/samples/snippets_test.py @@ -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()