Fix utcnow deprecations in samples and improve backup test reliability

Author: aseering

samples/samples/archived/backup_snippet.py

@@ -19,7 +19,8 @@
 """
 
 import time
-from datetime import datetime, timedelta
+from datetime import datetime, timedelta, timezone
+import time
 
 from google.cloud import spanner
 
@@ -30,7 +31,7 @@ def cancel_backup(instance_id, database_id, backup_id):
     instance = spanner_client.instance(instance_id)
     database = instance.database(database_id)
 
-    expire_time = datetime.utcnow() + timedelta(days=30)
+    expire_time = datetime.now(timezone.utc) + timedelta(days=30)
 
     # Create a backup.
     backup = instance.backup(backup_id, database=database, expire_time=expire_time)
@@ -63,14 +64,14 @@ def copy_backup(instance_id, backup_id, source_backup_path):
     instance = spanner_client.instance(instance_id)
 
     # Create a backup object and wait for copy backup operation to complete.
-    expire_time = datetime.utcnow() + timedelta(days=14)
+    expire_time = datetime.now(timezone.utc) + timedelta(days=14)
     copy_backup = instance.copy_backup(
         backup_id=backup_id, source_backup=source_backup_path, expire_time=expire_time
     )
     operation = copy_backup.create()
 
     # Wait for copy backup operation to complete.
-    operation.result(2100)
+    operation.result(3600)
 
     # Verify that the copy backup is ready.
     copy_backup.reload()
@@ -97,14 +98,14 @@ def create_backup(instance_id, database_id, backup_id, version_time):
     database = instance.database(database_id)
 
     # Create a backup
-    expire_time = datetime.utcnow() + timedelta(days=14)
+    expire_time = datetime.now(timezone.utc) + timedelta(days=14)
     backup = instance.backup(
         backup_id, database=database, expire_time=expire_time, version_time=version_time
     )
     operation = backup.create()
 
     # Wait for backup operation to complete.
-    operation.result(2100)
+    operation.result(3600)
 
     # Verify that the backup is ready.
     backup.reload()
@@ -127,15 +128,14 @@ 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
+    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)
+    expire_time = datetime.now(timezone.utc) + timedelta(days=14)
     encryption_config = {
         "encryption_type": CreateBackupEncryptionConfig.EncryptionType.CUSTOMER_MANAGED_ENCRYPTION,
         "kms_key_name": kms_key_name,
@@ -149,7 +149,7 @@ def create_backup_with_encryption_key(
     operation = backup.create()
 
     # Wait for backup operation to complete.
-    operation.result(2100)
+    operation.result(3600)
 
     # Verify that the backup is ready.
     backup.reload()
@@ -295,7 +295,7 @@ def list_backups(instance_id, database_id, backup_id):
         print(backup.name)
 
     # List all backups that expire before a timestamp.
-    expire_time = datetime.utcnow().replace(microsecond=0) + timedelta(days=30)
+    expire_time = datetime.now(timezone.utc).replace(microsecond=0) + timedelta(days=30)
     print(
         'All backups with expire_time before "{}-{}-{}T{}:{}:{}Z":'.format(
             *expire_time.timetuple()
@@ -312,7 +312,7 @@ def list_backups(instance_id, database_id, backup_id):
         print(backup.name)
 
     # List backups that were created after a timestamp that are also ready.
-    create_time = datetime.utcnow().replace(microsecond=0) - timedelta(days=1)
+    create_time = datetime.now(timezone.utc).replace(microsecond=0) - timedelta(days=1)
     print(
         'All backups created after "{}-{}-{}T{}:{}:{}Z" and are READY:'.format(
             *create_time.timetuple()
@@ -396,8 +396,7 @@ 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
+    from google.cloud.spanner_admin_database_v1 import RestoreDatabaseEncryptionConfig
 
     spanner_client = spanner.Client()
     instance = spanner_client.instance(instance_id)

samples/samples/archived/backup_snippet_test.py

@@ -41,6 +41,16 @@ def unique_backup_id():
 RETENTION_DATABASE_ID = unique_database_id()
 RETENTION_PERIOD = "7d"
 COPY_BACKUP_ID = unique_backup_id()
+CMEK_DATABASE_ID = unique_database_id()
+
+
+@pytest.fixture(scope="module")
+def cmek_database(spanner_client, sample_instance, database_dialect):
+    from conftest import create_sample_database
+
+    yield from create_sample_database(
+        spanner_client, sample_instance, CMEK_DATABASE_ID, [], database_dialect
+    )
 
 
 @pytest.mark.dependency(name="create_backup")
@@ -77,12 +87,12 @@ def test_copy_backup(capsys, instance_id, spanner_client):
 def test_create_backup_with_encryption_key(
     capsys,
     instance_id,
-    sample_database,
+    cmek_database,
     kms_key_name,
 ):
     backup_snippet.create_backup_with_encryption_key(
         instance_id,
-        sample_database.database_id,
+        cmek_database.database_id,
         CMEK_BACKUP_ID,
         kms_key_name,
     )
@@ -91,8 +101,10 @@ 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.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):
@@ -103,21 +115,23 @@ 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.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(
     capsys,
     instance_id,
-    sample_database,
+    cmek_database,
     kms_key_name,
 ):
     backup_snippet.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_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

samples/samples/backup_sample.py

@@ -19,7 +19,7 @@
 """
 
 import argparse
-from datetime import datetime, timedelta
+from datetime import datetime, timedelta, timezone
 import time
 
 from google.api_core import protobuf_helpers
@@ -37,7 +37,7 @@ def create_backup(instance_id, database_id, backup_id, version_time):
     database_admin_api = spanner_client.database_admin_api
 
     # Create a backup
-    expire_time = datetime.utcnow() + timedelta(days=14)
+    expire_time = datetime.now(timezone.utc) + timedelta(days=14)
 
     request = backup_pb.CreateBackupRequest(
         parent=database_admin_api.instance_path(spanner_client.project, instance_id),
@@ -54,7 +54,7 @@ def create_backup(instance_id, database_id, backup_id, version_time):
     operation = database_admin_api.create_backup(request)
 
     # Wait for backup operation to complete.
-    backup = operation.result(2100)
+    backup = operation.result(3600)
 
     # Verify that the backup is ready.
     assert backup.state == backup_pb.Backup.State.READY
@@ -82,7 +82,7 @@ def create_backup_with_encryption_key(
     database_admin_api = spanner_client.database_admin_api
 
     # Create a backup
-    expire_time = datetime.utcnow() + timedelta(days=14)
+    expire_time = datetime.now(timezone.utc) + timedelta(days=14)
     encryption_config = {
         "encryption_type": CreateBackupEncryptionConfig.EncryptionType.CUSTOMER_MANAGED_ENCRYPTION,
         "kms_key_name": kms_key_name,
@@ -101,7 +101,7 @@ def create_backup_with_encryption_key(
     operation = database_admin_api.create_backup(request)
 
     # Wait for backup operation to complete.
-    backup = operation.result(2100)
+    backup = operation.result(3600)
 
     # Verify that the backup is ready.
     assert backup.state == backup_pb.Backup.State.READY
@@ -130,7 +130,7 @@ def create_backup_with_multiple_kms_keys(
     database_admin_api = spanner_client.database_admin_api
 
     # Create a backup
-    expire_time = datetime.utcnow() + timedelta(days=14)
+    expire_time = datetime.now(timezone.utc) + timedelta(days=14)
     encryption_config = {
         "encryption_type": CreateBackupEncryptionConfig.EncryptionType.CUSTOMER_MANAGED_ENCRYPTION,
         "kms_key_names": kms_key_names,
@@ -149,7 +149,7 @@ def create_backup_with_multiple_kms_keys(
     operation = database_admin_api.create_backup(request)
 
     # Wait for backup operation to complete.
-    backup = operation.result(2100)
+    backup = operation.result(3600)
 
     # Verify that the backup is ready.
     assert backup.state == backup_pb.Backup.State.READY
@@ -302,7 +302,7 @@ def cancel_backup(instance_id, database_id, backup_id):
     spanner_client = spanner.Client()
     database_admin_api = spanner_client.database_admin_api
 
-    expire_time = datetime.utcnow() + timedelta(days=30)
+    expire_time = datetime.now(timezone.utc) + timedelta(days=30)
 
     # Create a backup.
     request = backup_pb.CreateBackupRequest(
@@ -473,7 +473,7 @@ def list_backups(instance_id, database_id, backup_id):
         print(backup.name)
 
     # List all backups that expire before a timestamp.
-    expire_time = datetime.utcnow().replace(microsecond=0) + timedelta(days=30)
+    expire_time = datetime.now(timezone.utc).replace(microsecond=0) + timedelta(days=30)
     print(
         'All backups with expire_time before "{}-{}-{}T{}:{}:{}Z":'.format(
             *expire_time.timetuple()
@@ -498,7 +498,7 @@ def list_backups(instance_id, database_id, backup_id):
         print(backup.name)
 
     # List backups that were created after a timestamp that are also ready.
-    create_time = datetime.utcnow().replace(microsecond=0) - timedelta(days=1)
+    create_time = datetime.now(timezone.utc).replace(microsecond=0) - timedelta(days=1)
     print(
         'All backups created after "{}-{}-{}T{}:{}:{}Z" and are READY:'.format(
             *create_time.timetuple()
@@ -671,7 +671,7 @@ def copy_backup(instance_id, backup_id, source_backup_path):
     database_admin_api = spanner_client.database_admin_api
 
     # Create a backup object and wait for copy backup operation to complete.
-    expire_time = datetime.utcnow() + timedelta(days=14)
+    expire_time = datetime.now(timezone.utc) + timedelta(days=14)
     request = backup_pb.CopyBackupRequest(
         parent=database_admin_api.instance_path(spanner_client.project, instance_id),
         backup_id=backup_id,
@@ -682,7 +682,7 @@ def copy_backup(instance_id, backup_id, source_backup_path):
     operation = database_admin_api.copy_backup(request)
 
     # Wait for backup operation to complete.
-    copy_backup = operation.result(2100)
+    copy_backup = operation.result(3600)
 
     # Verify that the copy backup is ready.
     assert copy_backup.state == backup_pb.Backup.State.READY
@@ -718,7 +718,7 @@ def copy_backup_with_multiple_kms_keys(
     }
 
     # Create a backup object and wait for copy backup operation to complete.
-    expire_time = datetime.utcnow() + timedelta(days=14)
+    expire_time = datetime.now(timezone.utc) + timedelta(days=14)
     request = backup_pb.CopyBackupRequest(
         parent=database_admin_api.instance_path(spanner_client.project, instance_id),
         backup_id=backup_id,
@@ -730,7 +730,7 @@ def copy_backup_with_multiple_kms_keys(
     operation = database_admin_api.copy_backup(request)
 
     # Wait for backup operation to complete.
-    copy_backup = operation.result(2100)
+    copy_backup = operation.result(3600)
 
     # Verify that the copy backup is ready.
     assert copy_backup.state == backup_pb.Backup.State.READY

samples/samples/backup_sample_test.py

@@ -42,6 +42,16 @@ def unique_backup_id():
 RETENTION_DATABASE_ID = unique_database_id()
 RETENTION_PERIOD = "7d"
 COPY_BACKUP_ID = unique_backup_id()
+CMEK_DATABASE_ID = unique_database_id()
+
+
+@pytest.fixture(scope="module")
+def cmek_database(spanner_client, sample_instance, database_dialect):
+    from conftest import create_sample_database
+
+    yield from create_sample_database(
+        spanner_client, sample_instance, CMEK_DATABASE_ID, [], database_dialect
+    )
 
 
 @pytest.mark.dependency(name="create_backup")
@@ -79,12 +89,12 @@ def test_copy_backup(capsys, instance_id, spanner_client):
 def test_create_backup_with_encryption_key(
     capsys,
     instance_id,
-    sample_database,
+    cmek_database,
     kms_key_name,
 ):
     backup_sample.create_backup_with_encryption_key(
         instance_id,
-        sample_database.database_id,
+        cmek_database.database_id,
         CMEK_BACKUP_ID,
         kms_key_name,
     )
@@ -149,14 +159,14 @@ def test_restore_database(capsys, instance_id, sample_database):
 def test_restore_database_with_encryption_key(
     capsys,
     instance_id,
-    sample_database,
+    cmek_database,
     kms_key_name,
 ):
     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_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

samples/samples/conftest.py

@@ -65,6 +65,9 @@ def scrub_instance_ignore_not_found(to_scrub):
         for backup_pb in to_scrub.list_backups():
             backup.Backup.from_pb(backup_pb, to_scrub).delete()
 
+        for database_pb in to_scrub.list_databases():
+            database.Database.from_pb(database_pb, to_scrub).drop()
+
         retry_429(to_scrub.delete)()
     except exceptions.NotFound:
         pass
@@ -154,12 +157,12 @@ def sample_instance(
 
     yield sample_instance
 
-    for database_pb in sample_instance.list_databases():
-        database.Database.from_pb(database_pb, sample_instance).drop()
-
     for backup_pb in sample_instance.list_backups():
         backup.Backup.from_pb(backup_pb, sample_instance).delete()
 
+    for database_pb in sample_instance.list_databases():
+        database.Database.from_pb(database_pb, sample_instance).drop()
+
     sample_instance.delete()
 
 
@@ -189,12 +192,12 @@ def multi_region_instance(
 
     yield multi_region_instance
 
-    for database_pb in multi_region_instance.list_databases():
-        database.Database.from_pb(database_pb, multi_region_instance).drop()
-
     for backup_pb in multi_region_instance.list_backups():
         backup.Backup.from_pb(backup_pb, multi_region_instance).delete()
 
+    for database_pb in multi_region_instance.list_databases():
+        database.Database.from_pb(database_pb, multi_region_instance).drop()
+
     multi_region_instance.delete()