rebase branch with main

samples/samples/admin/backup_sample.py → samples/samples/admin/backup_snippet.py

@@ -17,8 +17,8 @@
 For more information, see the README.rst under /spanner.
 """
 
-from datetime import datetime, timedelta
 import time
+from datetime import datetime, timedelta
 
 from google.cloud import spanner
 from google.cloud.spanner_admin_database_v1.types import backup as backup_pb
@@ -85,8 +85,8 @@ def create_backup_with_encryption_key(
         backup=backup_pb.Backup(
             database=database.name,
             expire_time=expire_time,
-            encryption_config=encryption_config,
         ),
+        encryption_config=encryption_config,
     )
     operation = spanner_client.database_admin_api.create_backup(request)
 
@@ -146,7 +146,10 @@ 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,
+        RestoreDatabaseRequest,
+    )
 
     spanner_client = spanner.Client()
     instance = spanner_client.instance(instance_id)
@@ -168,6 +171,20 @@ def restore_database_with_encryption_key(
     # Newly created database has restore information.
     new_database.reload()
     restore_info = new_database.restore_info
+
+    request = RestoreDatabaseRequest(
+        parent=instance.name,
+        database_id=new_database_id,
+        backup="{}/backups/{}".format(instance.name, backup_id),
+        encryption_config=encryption_config,
+    )
+    operation = spanner_client.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,
@@ -458,4 +475,4 @@ def copy_backup(instance_id, backup_id, source_backup_path):
     )
 
 
-# [END spanner_copy_backup
+# [END spanner_copy_backup]

samples/samples/admin/backup_sample_test.py → samples/samples/admin/backup_snippet_test.py

@@ -13,10 +13,9 @@
 # limitations under the License.
 import uuid
 
-from google.api_core.exceptions import DeadlineExceeded
+import backup_snippet
 import pytest
-
-import backup_sample
+from google.api_core.exceptions import DeadlineExceeded
 from test_utils.retry import RetryErrors
 
 
@@ -50,7 +49,7 @@ def test_create_backup(capsys, instance_id, sample_database):
         results = snapshot.execute_sql("SELECT CURRENT_TIMESTAMP()")
         version_time = list(results)[0][0]
 
-    backup_sample.create_backup(
+    backup_snippet.create_backup(
         instance_id,
         sample_database.database_id,
         BACKUP_ID,
@@ -67,7 +66,7 @@ def test_create_backup_with_encryption_key(
     sample_database,
     kms_key_name,
 ):
-    backup_sample.create_backup_with_encryption_key(
+    backup_snippet.create_backup_with_encryption_key(
         instance_id,
         sample_database.database_id,
         CMEK_BACKUP_ID,
@@ -81,8 +80,26 @@ def test_create_backup_with_encryption_key(
 # @pytest.mark.dependency(depends=["create_backup"])
 @RetryErrors(exception=DeadlineExceeded, max_tries=2)
 def test_restore_database(capsys, instance_id, sample_database):
-    backup_sample.restore_database(instance_id, RESTORE_DB_ID, BACKUP_ID)
+    backup_snippet.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
+
+
+@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,
+    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_RESTORE_DB_ID + " from backup ") in out
+    assert CMEK_BACKUP_ID in out
+    assert kms_key_name in out

samples/samples/admin/pg_samples.py

@@ -18,9 +18,6 @@
 Spanner PostgreSql dialect.
 For more information, see the README.rst under /spanner.
 """
-import base64
-import decimal
-
 from google.cloud import spanner
 from google.cloud.spanner_admin_database_v1.types import spanner_database_admin
 

samples/samples/admin/pg_samples_test.py

@@ -14,13 +14,12 @@
 
 import uuid
 
+import pg_samples as samples
+import pytest
 from google.api_core import exceptions
 from google.cloud.spanner_admin_database_v1.types.common import DatabaseDialect
-import pytest
 from test_utils.retry import RetryErrors
 
-import pg_samples as samples
-
 CREATE_TABLE_SINGERS = """\
 CREATE TABLE Singers (
     SingerId     BIGINT NOT NULL,

samples/samples/admin/samples.py

@@ -22,8 +22,8 @@
 import time
 
 from google.cloud import spanner
-from google.cloud.spanner_admin_instance_v1.types import spanner_instance_admin
 from google.cloud.spanner_admin_database_v1.types import spanner_database_admin
+from google.cloud.spanner_admin_instance_v1.types import spanner_instance_admin
 
 OPERATION_TIMEOUT_SECONDS = 240
 

samples/samples/admin/samples_test.py

@@ -21,9 +21,9 @@
 
 import uuid
 
+import pytest
 from google.api_core import exceptions
 from google.cloud.spanner_admin_database_v1.types.common import DatabaseDialect
-import pytest
 from test_utils.retry import RetryErrors
 
 import samples