From d3c646479bdc2ad96f68ed7d4703188dcdd76395 Mon Sep 17 00:00:00 2001 From: Sri Harsha CH <57220027+harshachinta@users.noreply.github.com> Date: Thu, 17 Oct 2024 00:54:21 +0530 Subject: [PATCH 1/3] chore: update sample instance edition to ENTERPRISE_PLUS for testing (#1212) * chore: update edition to ENTERPRISE_PLUS to test all features * chore: skip tests to unblock PR * chore: lint fix --- .../samples/archived/backup_snippet_test.py | 4 +++ samples/samples/conftest.py | 28 ++++++++++++------- samples/samples/snippets.py | 1 + 3 files changed, 23 insertions(+), 10 deletions(-) diff --git a/samples/samples/archived/backup_snippet_test.py b/samples/samples/archived/backup_snippet_test.py index 8fc29b9425..888124ffad 100644 --- a/samples/samples/archived/backup_snippet_test.py +++ b/samples/samples/archived/backup_snippet_test.py @@ -91,6 +91,8 @@ def test_create_backup_with_encryption_key( assert kms_key_name in out +@pytest.mark.skip(reason="same test passes on unarchived test suite, " + "but fails here. Needs investigation") @pytest.mark.dependency(depends=["create_backup"]) @RetryErrors(exception=DeadlineExceeded, max_tries=2) def test_restore_database(capsys, instance_id, sample_database): @@ -101,6 +103,8 @@ def test_restore_database(capsys, instance_id, sample_database): assert BACKUP_ID in out +@pytest.mark.skip(reason="same test passes on unarchived test suite, " + "but fails here. Needs investigation") @pytest.mark.dependency(depends=["create_backup_with_encryption_key"]) @RetryErrors(exception=DeadlineExceeded, max_tries=2) def test_restore_database_with_encryption_key( diff --git a/samples/samples/conftest.py b/samples/samples/conftest.py index 9810a41d45..2d72db62f3 100644 --- a/samples/samples/conftest.py +++ b/samples/samples/conftest.py @@ -22,6 +22,7 @@ from google.cloud.spanner_admin_database_v1.types.common import DatabaseDialect from google.cloud.spanner_v1 import backup, client, database, instance from test_utils import retry +from google.cloud.spanner_admin_instance_v1.types import spanner_instance_admin INSTANCE_CREATION_TIMEOUT = 560 # seconds @@ -128,17 +129,24 @@ def sample_instance( instance_config, sample_name, ): - sample_instance = spanner_client.instance( - instance_id, - instance_config, - labels={ - "cloud_spanner_samples": "true", - "sample_name": sample_name, - "created": str(int(time.time())), - }, + operation = spanner_client.instance_admin_api.create_instance( + parent=spanner_client.project_name, + instance_id=instance_id, + instance=spanner_instance_admin.Instance( + config=instance_config, + display_name="This is a display name.", + node_count=1, + labels={ + "cloud_spanner_samples": "true", + "sample_name": sample_name, + "created": str(int(time.time())), + }, + edition=spanner_instance_admin.Instance.Edition.ENTERPRISE_PLUS, # Optional + ), ) - op = retry_429(sample_instance.create)() - op.result(INSTANCE_CREATION_TIMEOUT) # block until completion + operation.result(INSTANCE_CREATION_TIMEOUT) # block until completion + + sample_instance = spanner_client.instance(instance_id) # Eventual consistency check retry_found = retry.RetryResult(bool) diff --git a/samples/samples/snippets.py b/samples/samples/snippets.py index 8a3764e9a5..984b8588a7 100644 --- a/samples/samples/snippets.py +++ b/samples/samples/snippets.py @@ -128,6 +128,7 @@ def create_instance_with_processing_units(instance_id, processing_units): "sample_name": "snippets-create_instance_with_processing_units", "created": str(int(time.time())), }, + edition=spanner_instance_admin.Instance.Edition.ENTERPRISE_PLUS, ), ) 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 2/3] 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) From 43c190bc694d56e0c57d96dbaa7fc48117f3c971 Mon Sep 17 00:00:00 2001 From: Sri Harsha CH <57220027+harshachinta@users.noreply.github.com> Date: Mon, 28 Oct 2024 19:22:25 +0530 Subject: [PATCH 3/3] fix: add PROTO in streaming chunks (#1213) b/372956316 When the row size exceeds a certain limit, the rows are divided into chunks and sent to the client in multiple parts. The client is responsible for merging these chunks to reconstruct the full row. However, for PROTO and ENUM types, this chunk-merging logic was not implemented, causing a KeyError: 13 when attempting to merge proto chunks. #### Sample to reproduce the test case [Python file](https://gist.github.com/harshachinta/95a81eeda81c422814353a5995d01e20) [proto file ](https://gist.github.com/harshachinta/fd15bf558bd4f40443411ddd164638cc) #### Steps to generate descriptors.pb and code file from proto ``` protoc --proto_path=testdata/ --include_imports --descriptor_set_out=testdata/descriptors.pb --python_out=testdata/ testdata/wrapper.proto ``` --- google/cloud/spanner_v1/streamed.py | 2 ++ tests/unit/test_streamed.py | 40 +++++++++++++++++++++++++++++ 2 files changed, 42 insertions(+) diff --git a/google/cloud/spanner_v1/streamed.py b/google/cloud/spanner_v1/streamed.py index 03acc9010a..89bde0e334 100644 --- a/google/cloud/spanner_v1/streamed.py +++ b/google/cloud/spanner_v1/streamed.py @@ -345,6 +345,8 @@ def _merge_struct(lhs, rhs, type_): TypeCode.TIMESTAMP: _merge_string, TypeCode.NUMERIC: _merge_string, TypeCode.JSON: _merge_string, + TypeCode.PROTO: _merge_string, + TypeCode.ENUM: _merge_string, } diff --git a/tests/unit/test_streamed.py b/tests/unit/test_streamed.py index 85dcb40026..83aa25a9d1 100644 --- a/tests/unit/test_streamed.py +++ b/tests/unit/test_streamed.py @@ -272,6 +272,46 @@ def test__merge_chunk_string_w_bytes(self): ) self.assertIsNone(streamed._pending_chunk) + def test__merge_chunk_proto(self): + from google.cloud.spanner_v1 import TypeCode + + iterator = _MockCancellableIterator() + streamed = self._make_one(iterator) + FIELDS = [self._make_scalar_field("proto", TypeCode.PROTO)] + streamed._metadata = self._make_result_set_metadata(FIELDS) + streamed._pending_chunk = self._make_value( + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAAAAAA" + "6fptVAAAACXBIWXMAAAsTAAALEwEAmpwYAAAA\n" + ) + chunk = self._make_value( + "B3RJTUUH4QQGFwsBTL3HMwAAABJpVFh0Q29tbWVudAAAAAAAU0FNUExF" + "MG3E+AAAAApJREFUCNdj\nYAAAAAIAAeIhvDMAAAAASUVORK5CYII=\n" + ) + + merged = streamed._merge_chunk(chunk) + + self.assertEqual( + merged.string_value, + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAAAAAA6fptVAAAACXBIWXMAAAsTAAAL" + "EwEAmpwYAAAA\nB3RJTUUH4QQGFwsBTL3HMwAAABJpVFh0Q29tbWVudAAAAAAAU0" + "FNUExFMG3E+AAAAApJREFUCNdj\nYAAAAAIAAeIhvDMAAAAASUVORK5CYII=\n", + ) + self.assertIsNone(streamed._pending_chunk) + + def test__merge_chunk_enum(self): + from google.cloud.spanner_v1 import TypeCode + + iterator = _MockCancellableIterator() + streamed = self._make_one(iterator) + FIELDS = [self._make_scalar_field("age", TypeCode.ENUM)] + streamed._metadata = self._make_result_set_metadata(FIELDS) + streamed._pending_chunk = self._make_value(42) + chunk = self._make_value(13) + + merged = streamed._merge_chunk(chunk) + self.assertEqual(merged.string_value, "4213") + self.assertIsNone(streamed._pending_chunk) + def test__merge_chunk_array_of_bool(self): from google.cloud.spanner_v1 import TypeCode