From e959f9e256f649f93abb09eaa0afe388319263ac Mon Sep 17 00:00:00 2001 From: Sakthivel Subramanian Date: Tue, 19 May 2026 06:38:00 +0000 Subject: [PATCH] test(spanner): handle AlreadyExists idempotently on restore test retries When a restore database test times out on the client-side, Spanner continues the restore operation in the background. When the test runner retries the test, the call would fail with AlreadyExists. This change catches the AlreadyExists exception, retrieves the database metadata, and asserts that it is restoring from the correct source backup and achieves a READY or READY_OPTIMIZING state. This allows the test to execute idempotently without namespace collisions. --- .../samples/samples/backup_sample_test.py | 116 ++++++++++++++---- 1 file changed, 91 insertions(+), 25 deletions(-) diff --git a/packages/google-cloud-spanner/samples/samples/backup_sample_test.py b/packages/google-cloud-spanner/samples/samples/backup_sample_test.py index 7dfb16589b95..9bcdc3a2b3d7 100644 --- a/packages/google-cloud-spanner/samples/samples/backup_sample_test.py +++ b/packages/google-cloud-spanner/samples/samples/backup_sample_test.py @@ -13,7 +13,7 @@ # limitations under the License. import uuid -from google.api_core.exceptions import DeadlineExceeded +from google.api_core.exceptions import AlreadyExists, DeadlineExceeded import pytest from test_utils.retry import RetryErrors @@ -35,6 +35,49 @@ def unique_backup_id(): return f"test-backup-{uuid.uuid4().hex[:10]}" +def _verify_restored_database( + spanner_client, + instance_id, + database_id, + backup_id, + kms_key_name=None, + kms_key_names=None, +): + from google.cloud.spanner_admin_database_v1.types import spanner_database_admin + import time + + database_admin_api = spanner_client.database_admin_api + db_path = database_admin_api.database_path( + spanner_client.project, instance_id, database_id + ) + db = database_admin_api.get_database(name=db_path) + + # Verify it was restored from the correct backup + restore_info = db.restore_info + assert restore_info is not None + assert backup_id in restore_info.backup_info.backup + + if kms_key_name: + assert kms_key_name in db.encryption_config.kms_key_name + if kms_key_names: + assert kms_key_names[0] in db.encryption_config.kms_key_names[0] + + # Verify/Wait until the state is either READY or READY_OPTIMIZING + for _ in range(40): # Wait up to 20 minutes + if db.state in ( + spanner_database_admin.Database.State.READY, + spanner_database_admin.Database.State.READY_OPTIMIZING, + ): + break + time.sleep(30) + db = database_admin_api.get_database(name=db_path) + + assert db.state in ( + spanner_database_admin.Database.State.READY, + spanner_database_admin.Database.State.READY_OPTIMIZING, + ) + + RESTORE_DB_ID = unique_database_id() BACKUP_ID = unique_backup_id() CMEK_RESTORE_DB_ID = unique_database_id() @@ -136,12 +179,15 @@ def test_copy_backup_with_multiple_kms_keys( @pytest.mark.dependency(depends=["create_backup"]) @RetryErrors(exception=(DeadlineExceeded, TimeoutError), max_tries=2) -def test_restore_database(capsys, instance_id, sample_database): - backup_sample.restore_database(instance_id, RESTORE_DB_ID, BACKUP_ID) - out, _ = capsys.readouterr() - assert (sample_database.database_id + " restored to ") in out - assert (RESTORE_DB_ID + " from backup ") in out - assert BACKUP_ID in out +def test_restore_database(capsys, instance_id, sample_database, spanner_client): + try: + backup_sample.restore_database(instance_id, RESTORE_DB_ID, BACKUP_ID) + out, _ = capsys.readouterr() + assert (sample_database.database_id + " restored to ") in out + assert (RESTORE_DB_ID + " from backup ") in out + assert BACKUP_ID in out + except AlreadyExists: + _verify_restored_database(spanner_client, instance_id, RESTORE_DB_ID, BACKUP_ID) @pytest.mark.dependency(depends=["create_backup_with_encryption_key"]) @@ -151,15 +197,25 @@ def test_restore_database_with_encryption_key( instance_id, sample_database, kms_key_name, + spanner_client, ): - backup_sample.restore_database_with_encryption_key( - instance_id, CMEK_RESTORE_DB_ID, CMEK_BACKUP_ID, kms_key_name - ) - out, _ = capsys.readouterr() - assert (sample_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_name in out + try: + backup_sample.restore_database_with_encryption_key( + instance_id, CMEK_RESTORE_DB_ID, CMEK_BACKUP_ID, kms_key_name + ) + out, _ = capsys.readouterr() + assert (sample_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_name in out + except AlreadyExists: + _verify_restored_database( + spanner_client, + instance_id, + CMEK_RESTORE_DB_ID, + CMEK_BACKUP_ID, + kms_key_name=kms_key_name, + ) @pytest.mark.skip(reason="skipped since the KMS keys are not added on test project") @@ -170,17 +226,27 @@ def test_restore_database_with_multiple_kms_keys( multi_region_instance_id, sample_multi_region_database, kms_key_names, + spanner_client, ): - 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 + try: + 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 + except AlreadyExists: + _verify_restored_database( + spanner_client, + multi_region_instance_id, + CMEK_RESTORE_DB_ID, + CMEK_BACKUP_ID, + kms_key_names=kms_key_names, + ) @pytest.mark.dependency(depends=["create_backup", "copy_backup"])