rebase branch with main

googleapis · Feb 20, 2024 · 99367f2 · 99367f2
1 parent 4ec3c25
commit 99367f2
Show file tree

Hide file tree

Showing 6 changed files with 48 additions and 18 deletions.
diff --git a/samples/samples/admin/backup_sample.py → samples/samples/admin/backup_snippet.py b/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]
diff --git a/samples/samples/admin/backup_sample_test.py → samples/samples/admin/backup_snippet_test.py b/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
diff --git a/samples/samples/admin/pg_samples.py b/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
 

diff --git a/samples/samples/admin/pg_samples_test.py b/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,

diff --git a/samples/samples/admin/samples.py b/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
 

diff --git a/samples/samples/admin/samples_test.py b/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