From 68551c20cd101045f3d3fe948d04b99388f28c26 Mon Sep 17 00:00:00 2001 From: Htut Khine Htay Win Date: Mon, 28 Oct 2024 05:26:16 -0700 Subject: [PATCH] feat: allow multiple KMS keys to create CMEK database/backup (#1191) * Removed api files * Fixed lint errors * Fixed integration tests. * Fixed lint in snippets_test.py * Resolved comments from reviewer * chore: skip tests since KMS keys are not added to test project --------- Co-authored-by: Sri Harsha CH <57220027+harshachinta@users.noreply.github.com> Co-authored-by: surbhigarg92 Co-authored-by: Sri Harsha CH --- samples/samples/backup_sample.py | 180 ++++++++++++++++++++++---- samples/samples/backup_sample_test.py | 67 +++++++++- samples/samples/conftest.py | 43 +++++- samples/samples/snippets.py | 125 ++++++++++-------- samples/samples/snippets_test.py | 19 +++ 5 files changed, 352 insertions(+), 82 deletions(-) diff --git a/samples/samples/backup_sample.py b/samples/samples/backup_sample.py index d3c2c667c5..e3a2b6957d 100644 --- a/samples/samples/backup_sample.py +++ b/samples/samples/backup_sample.py @@ -19,8 +19,8 @@ """ import argparse -import time from datetime import datetime, timedelta +import time from google.api_core import protobuf_helpers from google.cloud import spanner @@ -31,8 +31,7 @@ def create_backup(instance_id, database_id, backup_id, version_time): """Creates a backup for a database.""" - from google.cloud.spanner_admin_database_v1.types import \ - backup as backup_pb + from google.cloud.spanner_admin_database_v1.types import backup as backup_pb spanner_client = spanner.Client() database_admin_api = spanner_client.database_admin_api @@ -76,10 +75,8 @@ def create_backup_with_encryption_key( ): """Creates a backup for a database using a Customer Managed Encryption Key (CMEK).""" - from google.cloud.spanner_admin_database_v1 import \ - CreateBackupEncryptionConfig - from google.cloud.spanner_admin_database_v1.types import \ - backup as backup_pb + from google.cloud.spanner_admin_database_v1 import CreateBackupEncryptionConfig + from google.cloud.spanner_admin_database_v1.types import backup as backup_pb spanner_client = spanner.Client() database_admin_api = spanner_client.database_admin_api @@ -119,6 +116,53 @@ def create_backup_with_encryption_key( # [END spanner_create_backup_with_encryption_key] +# [START spanner_create_backup_with_MR_CMEK] +def create_backup_with_multiple_kms_keys( + instance_id, database_id, backup_id, kms_key_names +): + """Creates a backup for a database using multiple KMS keys(CMEK).""" + + from google.cloud.spanner_admin_database_v1 import CreateBackupEncryptionConfig + from google.cloud.spanner_admin_database_v1.types import backup as backup_pb + + spanner_client = spanner.Client() + database_admin_api = spanner_client.database_admin_api + + # Create a backup + expire_time = datetime.utcnow() + timedelta(days=14) + encryption_config = { + "encryption_type": CreateBackupEncryptionConfig.EncryptionType.CUSTOMER_MANAGED_ENCRYPTION, + "kms_key_names": kms_key_names, + } + request = backup_pb.CreateBackupRequest( + parent=database_admin_api.instance_path(spanner_client.project, instance_id), + backup_id=backup_id, + backup=backup_pb.Backup( + database=database_admin_api.database_path( + spanner_client.project, instance_id, database_id + ), + expire_time=expire_time, + ), + encryption_config=encryption_config, + ) + operation = database_admin_api.create_backup(request) + + # Wait for backup operation to complete. + backup = operation.result(2100) + + # Verify that the backup is ready. + assert backup.state == backup_pb.Backup.State.READY + + # Get the name, create time, backup size and encryption key. + print( + "Backup {} of size {} bytes was created at {} using encryption key {}".format( + backup.name, backup.size_bytes, backup.create_time, kms_key_names + ) + ) + + +# [END spanner_create_backup_with_MR_CMEK] + # [START spanner_restore_backup] def restore_database(instance_id, new_database_id, backup_id): @@ -162,7 +206,9 @@ def restore_database_with_encryption_key( ): """Restores a database from a backup using a Customer Managed Encryption Key (CMEK).""" from google.cloud.spanner_admin_database_v1 import ( - RestoreDatabaseEncryptionConfig, RestoreDatabaseRequest) + RestoreDatabaseEncryptionConfig, + RestoreDatabaseRequest, + ) spanner_client = spanner.Client() database_admin_api = spanner_client.database_admin_api @@ -200,11 +246,56 @@ def restore_database_with_encryption_key( # [END spanner_restore_backup_with_encryption_key] +# [START spanner_restore_backup_with_MR_CMEK] +def restore_database_with_multiple_kms_keys( + instance_id, new_database_id, backup_id, kms_key_names +): + """Restores a database from a backup using a Customer Managed Encryption Key (CMEK).""" + from google.cloud.spanner_admin_database_v1 import ( + RestoreDatabaseEncryptionConfig, + RestoreDatabaseRequest, + ) + + spanner_client = spanner.Client() + database_admin_api = spanner_client.database_admin_api + + # Start restoring an existing backup to a new database. + encryption_config = { + "encryption_type": RestoreDatabaseEncryptionConfig.EncryptionType.CUSTOMER_MANAGED_ENCRYPTION, + "kms_key_names": kms_key_names, + } + + request = RestoreDatabaseRequest( + parent=database_admin_api.instance_path(spanner_client.project, instance_id), + database_id=new_database_id, + backup=database_admin_api.backup_path( + spanner_client.project, instance_id, backup_id + ), + encryption_config=encryption_config, + ) + operation = database_admin_api.restore_database(request) + + # Wait for restore operation to complete. + db = operation.result(1600) + + # Newly created database has restore information. + restore_info = db.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, + db.encryption_config.kms_key_names, + ) + ) + + +# [END spanner_restore_backup_with_MR_CMEK] + # [START spanner_cancel_backup_create] def cancel_backup(instance_id, database_id, backup_id): - from google.cloud.spanner_admin_database_v1.types import \ - backup as backup_pb + from google.cloud.spanner_admin_database_v1.types import backup as backup_pb spanner_client = spanner.Client() database_admin_api = spanner_client.database_admin_api @@ -259,8 +350,7 @@ def cancel_backup(instance_id, database_id, backup_id): # [START spanner_list_backup_operations] def list_backup_operations(instance_id, database_id, backup_id): - from google.cloud.spanner_admin_database_v1.types import \ - backup as backup_pb + from google.cloud.spanner_admin_database_v1.types import backup as backup_pb spanner_client = spanner.Client() database_admin_api = spanner_client.database_admin_api @@ -314,8 +404,7 @@ def list_backup_operations(instance_id, database_id, backup_id): # [START spanner_list_database_operations] def list_database_operations(instance_id): - from google.cloud.spanner_admin_database_v1.types import \ - spanner_database_admin + from google.cloud.spanner_admin_database_v1.types import spanner_database_admin spanner_client = spanner.Client() database_admin_api = spanner_client.database_admin_api @@ -346,8 +435,7 @@ def list_database_operations(instance_id): # [START spanner_list_backups] def list_backups(instance_id, database_id, backup_id): - from google.cloud.spanner_admin_database_v1.types import \ - backup as backup_pb + from google.cloud.spanner_admin_database_v1.types import backup as backup_pb spanner_client = spanner.Client() database_admin_api = spanner_client.database_admin_api @@ -444,8 +532,7 @@ def list_backups(instance_id, database_id, backup_id): # [START spanner_delete_backup] def delete_backup(instance_id, backup_id): - from google.cloud.spanner_admin_database_v1.types import \ - backup as backup_pb + from google.cloud.spanner_admin_database_v1.types import backup as backup_pb spanner_client = spanner.Client() database_admin_api = spanner_client.database_admin_api @@ -486,8 +573,7 @@ def delete_backup(instance_id, backup_id): # [START spanner_update_backup] def update_backup(instance_id, backup_id): - from google.cloud.spanner_admin_database_v1.types import \ - backup as backup_pb + from google.cloud.spanner_admin_database_v1.types import backup as backup_pb spanner_client = spanner.Client() database_admin_api = spanner_client.database_admin_api @@ -526,8 +612,7 @@ def create_database_with_version_retention_period( ): """Creates a database with a version retention period.""" - from google.cloud.spanner_admin_database_v1.types import \ - spanner_database_admin + from google.cloud.spanner_admin_database_v1.types import spanner_database_admin spanner_client = spanner.Client() database_admin_api = spanner_client.database_admin_api @@ -578,8 +663,7 @@ def create_database_with_version_retention_period( def copy_backup(instance_id, backup_id, source_backup_path): """Copies a backup.""" - from google.cloud.spanner_admin_database_v1.types import \ - backup as backup_pb + from google.cloud.spanner_admin_database_v1.types import backup as backup_pb spanner_client = spanner.Client() database_admin_api = spanner_client.database_admin_api @@ -613,6 +697,54 @@ def copy_backup(instance_id, backup_id, source_backup_path): # [END spanner_copy_backup] +# [START spanner_copy_backup_with_MR_CMEK] +def copy_backup_with_multiple_kms_keys( + instance_id, backup_id, source_backup_path, kms_key_names +): + """Copies a backup.""" + + from google.cloud.spanner_admin_database_v1.types import backup as backup_pb + from google.cloud.spanner_admin_database_v1 import CopyBackupEncryptionConfig + + spanner_client = spanner.Client() + database_admin_api = spanner_client.database_admin_api + + encryption_config = { + "encryption_type": CopyBackupEncryptionConfig.EncryptionType.CUSTOMER_MANAGED_ENCRYPTION, + "kms_key_names": kms_key_names, + } + + # Create a backup object and wait for copy backup operation to complete. + expire_time = datetime.utcnow() + timedelta(days=14) + request = backup_pb.CopyBackupRequest( + parent=database_admin_api.instance_path(spanner_client.project, instance_id), + backup_id=backup_id, + source_backup=source_backup_path, + expire_time=expire_time, + encryption_config=encryption_config, + ) + + operation = database_admin_api.copy_backup(request) + + # Wait for backup operation to complete. + copy_backup = operation.result(2100) + + # Verify that the copy backup is ready. + assert copy_backup.state == backup_pb.Backup.State.READY + + print( + "Backup {} of size {} bytes was created at {} with version time {} using encryption keys {}".format( + copy_backup.name, + copy_backup.size_bytes, + copy_backup.create_time, + copy_backup.version_time, + copy_backup.encryption_information, + ) + ) + + +# [END spanner_copy_backup_with_MR_CMEK] + if __name__ == "__main__": # noqa: C901 parser = argparse.ArgumentParser( diff --git a/samples/samples/backup_sample_test.py b/samples/samples/backup_sample_test.py index 6d656c5545..5ab1e747ab 100644 --- a/samples/samples/backup_sample_test.py +++ b/samples/samples/backup_sample_test.py @@ -13,8 +13,8 @@ # limitations under the License. import uuid -import pytest from google.api_core.exceptions import DeadlineExceeded +import pytest from test_utils.retry import RetryErrors import backup_sample @@ -93,6 +93,49 @@ def test_create_backup_with_encryption_key( assert kms_key_name in out +@pytest.mark.skip(reason="skipped since the KMS keys are not added on test " + "project") +@pytest.mark.dependency(name="create_backup_with_multiple_kms_keys") +def test_create_backup_with_multiple_kms_keys( + capsys, + multi_region_instance, + multi_region_instance_id, + sample_multi_region_database, + kms_key_names, +): + backup_sample.create_backup_with_multiple_kms_keys( + multi_region_instance_id, + sample_multi_region_database.database_id, + CMEK_BACKUP_ID, + kms_key_names, + ) + out, _ = capsys.readouterr() + assert CMEK_BACKUP_ID in out + assert kms_key_names[0] in out + assert kms_key_names[1] in out + assert kms_key_names[2] in out + + +@pytest.mark.skip(reason="skipped since the KMS keys are not added on test " + "project") +@pytest.mark.dependency(depends=["create_backup_with_multiple_kms_keys"]) +def test_copy_backup_with_multiple_kms_keys( + capsys, multi_region_instance_id, spanner_client, kms_key_names +): + source_backup_path = ( + spanner_client.project_name + + "/instances/" + + multi_region_instance_id + + "/backups/" + + CMEK_BACKUP_ID + ) + backup_sample.copy_backup_with_multiple_kms_keys( + multi_region_instance_id, COPY_BACKUP_ID, source_backup_path, kms_key_names + ) + out, _ = capsys.readouterr() + assert COPY_BACKUP_ID in out + + @pytest.mark.dependency(depends=["create_backup"]) @RetryErrors(exception=DeadlineExceeded, max_tries=2) def test_restore_database(capsys, instance_id, sample_database): @@ -121,6 +164,28 @@ def test_restore_database_with_encryption_key( assert kms_key_name in out +@pytest.mark.skip(reason="skipped since the KMS keys are not added on test " + "project") +@pytest.mark.dependency(depends=["create_backup_with_multiple_kms_keys"]) +@RetryErrors(exception=DeadlineExceeded, max_tries=2) +def test_restore_database_with_multiple_kms_keys( + capsys, + multi_region_instance_id, + sample_multi_region_database, + kms_key_names, +): + backup_sample.restore_database_with_multiple_kms_keys( + multi_region_instance_id, CMEK_RESTORE_DB_ID, CMEK_BACKUP_ID, kms_key_names + ) + out, _ = capsys.readouterr() + assert (sample_multi_region_database.database_id + " restored to ") in out + assert (CMEK_RESTORE_DB_ID + " from backup ") in out + assert CMEK_BACKUP_ID in out + assert kms_key_names[0] in out + assert kms_key_names[1] in out + assert kms_key_names[2] in out + + @pytest.mark.dependency(depends=["create_backup", "copy_backup"]) def test_list_backup_operations(capsys, instance_id, sample_database): backup_sample.list_backup_operations( diff --git a/samples/samples/conftest.py b/samples/samples/conftest.py index 2d72db62f3..b34e9d16b1 100644 --- a/samples/samples/conftest.py +++ b/samples/samples/conftest.py @@ -16,11 +16,11 @@ import time import uuid -import pytest from google.api_core import exceptions from google.cloud import spanner_admin_database_v1 from google.cloud.spanner_admin_database_v1.types.common import DatabaseDialect from google.cloud.spanner_v1 import backup, client, database, instance +import pytest from test_utils import retry from google.cloud.spanner_admin_instance_v1.types import spanner_instance_admin @@ -248,8 +248,7 @@ def database_ddl(): return [] -@pytest.fixture(scope="module") -def sample_database( +def create_sample_database( spanner_client, sample_instance, database_id, database_ddl, database_dialect ): if database_dialect == DatabaseDialect.POSTGRESQL: @@ -289,6 +288,28 @@ def sample_database( sample_database.drop() +@pytest.fixture(scope="module") +def sample_database( + spanner_client, sample_instance, database_id, database_ddl, database_dialect +): + yield from create_sample_database( + spanner_client, sample_instance, database_id, database_ddl, database_dialect + ) + + +@pytest.fixture(scope="module") +def sample_multi_region_database( + spanner_client, multi_region_instance, database_id, database_ddl, database_dialect +): + yield from create_sample_database( + spanner_client, + multi_region_instance, + database_id, + database_ddl, + database_dialect, + ) + + @pytest.fixture(scope="module") def bit_reverse_sequence_database( spanner_client, sample_instance, bit_reverse_sequence_database_id, database_dialect @@ -329,3 +350,19 @@ def kms_key_name(spanner_client): "spanner-test-keyring", "spanner-test-cmek", ) + + +@pytest.fixture(scope="module") +def kms_key_names(spanner_client): + kms_key_names_list = [] + # this list of cloud-regions correspond to `nam3` + for cloud_region in ["us-east1", "us-east4", "us-central1"]: + kms_key_names_list.append( + "projects/{}/locations/{}/keyRings/{}/cryptoKeys/{}".format( + spanner_client.project, + cloud_region, + "spanner-test-keyring", + "spanner-test-cmek", + ) + ) + return kms_key_names_list diff --git a/samples/samples/snippets.py b/samples/samples/snippets.py index 984b8588a7..c958a66822 100644 --- a/samples/samples/snippets.py +++ b/samples/samples/snippets.py @@ -33,6 +33,7 @@ from google.cloud.spanner_v1 import DirectedReadOptions, param_types from google.cloud.spanner_v1.data_types import JsonObject from google.protobuf import field_mask_pb2 # type: ignore + from testdata import singer_pb2 OPERATION_TIMEOUT_SECONDS = 240 @@ -41,8 +42,7 @@ # [START spanner_create_instance] def create_instance(instance_id): """Creates an instance.""" - from google.cloud.spanner_admin_instance_v1.types import \ - spanner_instance_admin + from google.cloud.spanner_admin_instance_v1.types import spanner_instance_admin spanner_client = spanner.Client() @@ -107,8 +107,7 @@ def update_instance(instance_id): # [START spanner_create_instance_with_processing_units] def create_instance_with_processing_units(instance_id, processing_units): """Creates an instance.""" - from google.cloud.spanner_admin_instance_v1.types import \ - spanner_instance_admin + from google.cloud.spanner_admin_instance_v1.types import spanner_instance_admin spanner_client = spanner.Client() @@ -168,8 +167,7 @@ def get_instance_config(instance_config): # [START spanner_list_instance_configs] def list_instance_config(): """Lists the available instance configurations.""" - from google.cloud.spanner_admin_instance_v1.types import \ - spanner_instance_admin + from google.cloud.spanner_admin_instance_v1.types import spanner_instance_admin spanner_client = spanner.Client() @@ -192,8 +190,7 @@ def list_instance_config(): # [START spanner_create_instance_partition] def create_instance_partition(instance_id, instance_partition_id): """Creates an instance partition.""" - from google.cloud.spanner_admin_instance_v1.types import \ - spanner_instance_admin + from google.cloud.spanner_admin_instance_v1.types import spanner_instance_admin spanner_client = spanner.Client() instance_admin_api = spanner_client.instance_admin_api @@ -222,8 +219,7 @@ def create_instance_partition(instance_id, instance_partition_id): # [START spanner_list_databases] def list_databases(instance_id): """Lists databases and their leader options.""" - from google.cloud.spanner_admin_database_v1.types import \ - spanner_database_admin + from google.cloud.spanner_admin_database_v1.types import spanner_database_admin spanner_client = spanner.Client() database_admin_api = spanner_client.database_admin_api @@ -246,8 +242,7 @@ def list_databases(instance_id): # [START spanner_create_database] def create_database(instance_id, database_id): """Creates a database and tables for sample data.""" - from google.cloud.spanner_admin_database_v1.types import \ - spanner_database_admin + from google.cloud.spanner_admin_database_v1.types import spanner_database_admin spanner_client = spanner.Client() database_admin_api = spanner_client.database_admin_api @@ -293,8 +288,7 @@ def create_database(instance_id, database_id): # [START spanner_update_database] def update_database(instance_id, database_id): """Updates the drop protection setting for a database.""" - from google.cloud.spanner_admin_database_v1.types import \ - spanner_database_admin + from google.cloud.spanner_admin_database_v1.types import spanner_database_admin spanner_client = spanner.Client() database_admin_api = spanner_client.database_admin_api @@ -332,8 +326,7 @@ def update_database(instance_id, database_id): 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).""" from google.cloud.spanner_admin_database_v1 import EncryptionConfig - from google.cloud.spanner_admin_database_v1.types import \ - spanner_database_admin + from google.cloud.spanner_admin_database_v1.types import spanner_database_admin spanner_client = spanner.Client() database_admin_api = spanner_client.database_admin_api @@ -372,12 +365,53 @@ def create_database_with_encryption_key(instance_id, database_id, kms_key_name): # [END spanner_create_database_with_encryption_key] +# [START spanner_create_database_with_MR_CMEK] +def create_database_with_multiple_kms_keys(instance_id, database_id, kms_key_names): + """Creates a database with tables using multiple KMS keys(CMEK).""" + from google.cloud.spanner_admin_database_v1 import EncryptionConfig + from google.cloud.spanner_admin_database_v1.types import spanner_database_admin + + spanner_client = spanner.Client() + database_admin_api = spanner_client.database_admin_api + + request = spanner_database_admin.CreateDatabaseRequest( + parent=database_admin_api.instance_path(spanner_client.project, instance_id), + create_statement=f"CREATE DATABASE `{database_id}`", + extra_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=EncryptionConfig(kms_key_names=kms_key_names), + ) + + operation = database_admin_api.create_database(request=request) + + print("Waiting for operation to complete...") + database = operation.result(OPERATION_TIMEOUT_SECONDS) + + print( + "Database {} created with multiple KMS keys {}".format( + database.name, database.encryption_config.kms_key_names + ) + ) + + +# [END spanner_create_database_with_MR_CMEK] # [START spanner_create_database_with_default_leader] def create_database_with_default_leader(instance_id, database_id, default_leader): """Creates a database with tables with a default leader.""" - from google.cloud.spanner_admin_database_v1.types import \ - spanner_database_admin + from google.cloud.spanner_admin_database_v1.types import spanner_database_admin spanner_client = spanner.Client() database_admin_api = spanner_client.database_admin_api @@ -420,8 +454,7 @@ def create_database_with_default_leader(instance_id, database_id, default_leader # [START spanner_update_database_with_default_leader] def update_database_with_default_leader(instance_id, database_id, default_leader): """Updates a database with tables with a default leader.""" - from google.cloud.spanner_admin_database_v1.types import \ - spanner_database_admin + from google.cloud.spanner_admin_database_v1.types import spanner_database_admin spanner_client = spanner.Client() database_admin_api = spanner_client.database_admin_api @@ -713,8 +746,7 @@ def query_data_with_new_column(instance_id, database_id): def add_index(instance_id, database_id): """Adds a simple index to the example database.""" - from google.cloud.spanner_admin_database_v1.types import \ - spanner_database_admin + from google.cloud.spanner_admin_database_v1.types import spanner_database_admin spanner_client = spanner.Client() database_admin_api = spanner_client.database_admin_api @@ -817,8 +849,7 @@ def read_data_with_index(instance_id, database_id): def add_storing_index(instance_id, database_id): """Adds an storing index to the example database.""" - from google.cloud.spanner_admin_database_v1.types import \ - spanner_database_admin + from google.cloud.spanner_admin_database_v1.types import spanner_database_admin spanner_client = spanner.Client() database_admin_api = spanner_client.database_admin_api @@ -881,8 +912,7 @@ def read_data_with_storing_index(instance_id, database_id): def add_column(instance_id, database_id): """Adds a new column to the Albums table in the example database.""" - from google.cloud.spanner_admin_database_v1.types import \ - spanner_database_admin + from google.cloud.spanner_admin_database_v1.types import spanner_database_admin spanner_client = spanner.Client() database_admin_api = spanner_client.database_admin_api @@ -1044,8 +1074,7 @@ def read_only_transaction(instance_id, database_id): def create_table_with_timestamp(instance_id, database_id): """Creates a table with a COMMIT_TIMESTAMP column.""" - from google.cloud.spanner_admin_database_v1.types import \ - spanner_database_admin + from google.cloud.spanner_admin_database_v1.types import spanner_database_admin spanner_client = spanner.Client() database_admin_api = spanner_client.database_admin_api @@ -1112,8 +1141,7 @@ def insert_data_with_timestamp(instance_id, database_id): def add_timestamp_column(instance_id, database_id): """Adds a new TIMESTAMP column to the Albums table in the example database.""" - from google.cloud.spanner_admin_database_v1.types import \ - spanner_database_admin + from google.cloud.spanner_admin_database_v1.types import spanner_database_admin spanner_client = spanner.Client() database_admin_api = spanner_client.database_admin_api @@ -1216,8 +1244,7 @@ def query_data_with_timestamp(instance_id, database_id): def add_numeric_column(instance_id, database_id): """Adds a new NUMERIC column to the Venues table in the example database.""" - from google.cloud.spanner_admin_database_v1.types import \ - spanner_database_admin + from google.cloud.spanner_admin_database_v1.types import spanner_database_admin spanner_client = spanner.Client() database_admin_api = spanner_client.database_admin_api @@ -1284,8 +1311,7 @@ def add_json_column(instance_id, database_id): # instance_id = "your-spanner-instance" # database_id = "your-spanner-db-id" - from google.cloud.spanner_admin_database_v1.types import \ - spanner_database_admin + from google.cloud.spanner_admin_database_v1.types import spanner_database_admin spanner_client = spanner.Client() database_admin_api = spanner_client.database_admin_api @@ -2018,8 +2044,7 @@ def create_table_with_datatypes(instance_id, database_id): # instance_id = "your-spanner-instance" # database_id = "your-spanner-db-id" - from google.cloud.spanner_admin_database_v1.types import \ - spanner_database_admin + from google.cloud.spanner_admin_database_v1.types import spanner_database_admin spanner_client = spanner.Client() database_admin_api = spanner_client.database_admin_api @@ -2613,8 +2638,7 @@ def add_and_drop_database_roles(instance_id, database_id): # instance_id = "your-spanner-instance" # database_id = "your-spanner-db-id" - from google.cloud.spanner_admin_database_v1.types import \ - spanner_database_admin + from google.cloud.spanner_admin_database_v1.types import spanner_database_admin spanner_client = spanner.Client() database_admin_api = spanner_client.database_admin_api @@ -2680,8 +2704,7 @@ def list_database_roles(instance_id, database_id): # [START spanner_list_database_roles] # instance_id = "your-spanner-instance" # database_id = "your-spanner-db-id" - from google.cloud.spanner_admin_database_v1.types import \ - spanner_database_admin + from google.cloud.spanner_admin_database_v1.types import spanner_database_admin spanner_client = spanner.Client() database_admin_api = spanner_client.database_admin_api @@ -2763,8 +2786,7 @@ def enable_fine_grained_access( def create_table_with_foreign_key_delete_cascade(instance_id, database_id): """Creates a table with foreign key delete cascade action""" - from google.cloud.spanner_admin_database_v1.types import \ - spanner_database_admin + from google.cloud.spanner_admin_database_v1.types import spanner_database_admin spanner_client = spanner.Client() database_admin_api = spanner_client.database_admin_api @@ -2811,8 +2833,7 @@ def create_table_with_foreign_key_delete_cascade(instance_id, database_id): def alter_table_with_foreign_key_delete_cascade(instance_id, database_id): """Alters a table with foreign key delete cascade action""" - from google.cloud.spanner_admin_database_v1.types import \ - spanner_database_admin + from google.cloud.spanner_admin_database_v1.types import spanner_database_admin spanner_client = spanner.Client() database_admin_api = spanner_client.database_admin_api @@ -2850,8 +2871,7 @@ def alter_table_with_foreign_key_delete_cascade(instance_id, database_id): def drop_foreign_key_constraint_delete_cascade(instance_id, database_id): """Alter table to drop foreign key delete cascade action""" - from google.cloud.spanner_admin_database_v1.types import \ - spanner_database_admin + from google.cloud.spanner_admin_database_v1.types import spanner_database_admin spanner_client = spanner.Client() database_admin_api = spanner_client.database_admin_api @@ -2886,8 +2906,7 @@ def drop_foreign_key_constraint_delete_cascade(instance_id, database_id): def create_sequence(instance_id, database_id): """Creates the Sequence and insert data""" - from google.cloud.spanner_admin_database_v1.types import \ - spanner_database_admin + from google.cloud.spanner_admin_database_v1.types import spanner_database_admin spanner_client = spanner.Client() database_admin_api = spanner_client.database_admin_api @@ -2945,8 +2964,7 @@ def insert_customers(transaction): def alter_sequence(instance_id, database_id): """Alters the Sequence and insert data""" - from google.cloud.spanner_admin_database_v1.types import \ - spanner_database_admin + from google.cloud.spanner_admin_database_v1.types import spanner_database_admin spanner_client = spanner.Client() database_admin_api = spanner_client.database_admin_api @@ -3000,8 +3018,7 @@ def insert_customers(transaction): def drop_sequence(instance_id, database_id): """Drops the Sequence""" - from google.cloud.spanner_admin_database_v1.types import \ - spanner_database_admin + from google.cloud.spanner_admin_database_v1.types import spanner_database_admin spanner_client = spanner.Client() database_admin_api = spanner_client.database_admin_api @@ -3150,8 +3167,7 @@ def set_custom_timeout_and_retry(instance_id, database_id): # [START spanner_create_instance_with_autoscaling_config] def create_instance_with_autoscaling_config(instance_id): """Creates a Cloud Spanner instance with an autoscaling configuration.""" - from google.cloud.spanner_admin_instance_v1.types import \ - spanner_instance_admin + from google.cloud.spanner_admin_instance_v1.types import spanner_instance_admin spanner_client = spanner.Client() @@ -3214,6 +3230,7 @@ def add_proto_type_columns(instance_id, database_id): """Adds a new Proto Message column and Proto Enum column to the Singers table.""" import os + from google.cloud.spanner_admin_database_v1.types import spanner_database_admin dirname = os.path.dirname(__file__) diff --git a/samples/samples/snippets_test.py b/samples/samples/snippets_test.py index 6938aa1cd7..ba3c0bbfe7 100644 --- a/samples/samples/snippets_test.py +++ b/samples/samples/snippets_test.py @@ -233,6 +233,25 @@ def test_create_database_with_encryption_config( assert kms_key_name in out +@pytest.mark.skip(reason="skipped since the KMS keys are not added on test " + "project") +def test_create_database_with_multiple_kms_keys( + capsys, + multi_region_instance, + multi_region_instance_id, + cmek_database_id, + kms_key_names, +): + snippets.create_database_with_multiple_kms_keys( + multi_region_instance_id, cmek_database_id, kms_key_names + ) + out, _ = capsys.readouterr() + assert cmek_database_id in out + assert kms_key_names[0] in out + assert kms_key_names[1] in out + assert kms_key_names[2] in out + + def test_get_instance_config(capsys): instance_config = "nam6" snippets.get_instance_config(instance_config)