Skip to content

Commit 68551c2

Browse files
hwin16harshachintasurbhigarg92
authored
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 <[email protected]> Co-authored-by: surbhigarg92 <[email protected]> Co-authored-by: Sri Harsha CH <[email protected]>
1 parent d3c6464 commit 68551c2

File tree

5 files changed

+352
-82
lines changed

5 files changed

+352
-82
lines changed

‎samples/samples/backup_sample.py

+156-24
Original file line numberDiff line numberDiff line change
@@ -19,8 +19,8 @@
1919
"""
2020

2121
import argparse
22-
import time
2322
from datetime import datetime, timedelta
23+
import time
2424

2525
from google.api_core import protobuf_helpers
2626
from google.cloud import spanner
@@ -31,8 +31,7 @@
3131
def create_backup(instance_id, database_id, backup_id, version_time):
3232
"""Creates a backup for a database."""
3333

34-
from google.cloud.spanner_admin_database_v1.types import \
35-
backup as backup_pb
34+
from google.cloud.spanner_admin_database_v1.types import backup as backup_pb
3635

3736
spanner_client = spanner.Client()
3837
database_admin_api = spanner_client.database_admin_api
@@ -76,10 +75,8 @@ def create_backup_with_encryption_key(
7675
):
7776
"""Creates a backup for a database using a Customer Managed Encryption Key (CMEK)."""
7877

79-
from google.cloud.spanner_admin_database_v1 import \
80-
CreateBackupEncryptionConfig
81-
from google.cloud.spanner_admin_database_v1.types import \
82-
backup as backup_pb
78+
from google.cloud.spanner_admin_database_v1 import CreateBackupEncryptionConfig
79+
from google.cloud.spanner_admin_database_v1.types import backup as backup_pb
8380

8481
spanner_client = spanner.Client()
8582
database_admin_api = spanner_client.database_admin_api
@@ -119,6 +116,53 @@ def create_backup_with_encryption_key(
119116

120117
# [END spanner_create_backup_with_encryption_key]
121118

119+
# [START spanner_create_backup_with_MR_CMEK]
120+
def create_backup_with_multiple_kms_keys(
121+
instance_id, database_id, backup_id, kms_key_names
122+
):
123+
"""Creates a backup for a database using multiple KMS keys(CMEK)."""
124+
125+
from google.cloud.spanner_admin_database_v1 import CreateBackupEncryptionConfig
126+
from google.cloud.spanner_admin_database_v1.types import backup as backup_pb
127+
128+
spanner_client = spanner.Client()
129+
database_admin_api = spanner_client.database_admin_api
130+
131+
# Create a backup
132+
expire_time = datetime.utcnow() + timedelta(days=14)
133+
encryption_config = {
134+
"encryption_type": CreateBackupEncryptionConfig.EncryptionType.CUSTOMER_MANAGED_ENCRYPTION,
135+
"kms_key_names": kms_key_names,
136+
}
137+
request = backup_pb.CreateBackupRequest(
138+
parent=database_admin_api.instance_path(spanner_client.project, instance_id),
139+
backup_id=backup_id,
140+
backup=backup_pb.Backup(
141+
database=database_admin_api.database_path(
142+
spanner_client.project, instance_id, database_id
143+
),
144+
expire_time=expire_time,
145+
),
146+
encryption_config=encryption_config,
147+
)
148+
operation = database_admin_api.create_backup(request)
149+
150+
# Wait for backup operation to complete.
151+
backup = operation.result(2100)
152+
153+
# Verify that the backup is ready.
154+
assert backup.state == backup_pb.Backup.State.READY
155+
156+
# Get the name, create time, backup size and encryption key.
157+
print(
158+
"Backup {} of size {} bytes was created at {} using encryption key {}".format(
159+
backup.name, backup.size_bytes, backup.create_time, kms_key_names
160+
)
161+
)
162+
163+
164+
# [END spanner_create_backup_with_MR_CMEK]
165+
122166

123167
# [START spanner_restore_backup]
124168
def restore_database(instance_id, new_database_id, backup_id):
@@ -162,7 +206,9 @@ def restore_database_with_encryption_key(
162206
):
163207
"""Restores a database from a backup using a Customer Managed Encryption Key (CMEK)."""
164208
from google.cloud.spanner_admin_database_v1 import (
165-
RestoreDatabaseEncryptionConfig, RestoreDatabaseRequest)
209+
RestoreDatabaseEncryptionConfig,
210+
RestoreDatabaseRequest,
211+
)
166212

167213
spanner_client = spanner.Client()
168214
database_admin_api = spanner_client.database_admin_api
@@ -200,11 +246,56 @@ def restore_database_with_encryption_key(
200246

201247
# [END spanner_restore_backup_with_encryption_key]
202248

249+
# [START spanner_restore_backup_with_MR_CMEK]
250+
def restore_database_with_multiple_kms_keys(
251+
instance_id, new_database_id, backup_id, kms_key_names
252+
):
253+
"""Restores a database from a backup using a Customer Managed Encryption Key (CMEK)."""
254+
from google.cloud.spanner_admin_database_v1 import (
255+
RestoreDatabaseEncryptionConfig,
256+
RestoreDatabaseRequest,
257+
)
258+
259+
spanner_client = spanner.Client()
260+
database_admin_api = spanner_client.database_admin_api
261+
262+
# Start restoring an existing backup to a new database.
263+
encryption_config = {
264+
"encryption_type": RestoreDatabaseEncryptionConfig.EncryptionType.CUSTOMER_MANAGED_ENCRYPTION,
265+
"kms_key_names": kms_key_names,
266+
}
267+
268+
request = RestoreDatabaseRequest(
269+
parent=database_admin_api.instance_path(spanner_client.project, instance_id),
270+
database_id=new_database_id,
271+
backup=database_admin_api.backup_path(
272+
spanner_client.project, instance_id, backup_id
273+
),
274+
encryption_config=encryption_config,
275+
)
276+
operation = database_admin_api.restore_database(request)
277+
278+
# Wait for restore operation to complete.
279+
db = operation.result(1600)
280+
281+
# Newly created database has restore information.
282+
restore_info = db.restore_info
283+
print(
284+
"Database {} restored to {} from backup {} with using encryption key {}.".format(
285+
restore_info.backup_info.source_database,
286+
new_database_id,
287+
restore_info.backup_info.backup,
288+
db.encryption_config.kms_key_names,
289+
)
290+
)
291+
292+
293+
# [END spanner_restore_backup_with_MR_CMEK]
294+
203295

204296
# [START spanner_cancel_backup_create]
205297
def cancel_backup(instance_id, database_id, backup_id):
206-
from google.cloud.spanner_admin_database_v1.types import \
207-
backup as backup_pb
298+
from google.cloud.spanner_admin_database_v1.types import backup as backup_pb
208299

209300
spanner_client = spanner.Client()
210301
database_admin_api = spanner_client.database_admin_api
@@ -259,8 +350,7 @@ def cancel_backup(instance_id, database_id, backup_id):
259350

260351
# [START spanner_list_backup_operations]
261352
def list_backup_operations(instance_id, database_id, backup_id):
262-
from google.cloud.spanner_admin_database_v1.types import \
263-
backup as backup_pb
353+
from google.cloud.spanner_admin_database_v1.types import backup as backup_pb
264354

265355
spanner_client = spanner.Client()
266356
database_admin_api = spanner_client.database_admin_api
@@ -314,8 +404,7 @@ def list_backup_operations(instance_id, database_id, backup_id):
314404

315405
# [START spanner_list_database_operations]
316406
def list_database_operations(instance_id):
317-
from google.cloud.spanner_admin_database_v1.types import \
318-
spanner_database_admin
407+
from google.cloud.spanner_admin_database_v1.types import spanner_database_admin
319408

320409
spanner_client = spanner.Client()
321410
database_admin_api = spanner_client.database_admin_api
@@ -346,8 +435,7 @@ def list_database_operations(instance_id):
346435

347436
# [START spanner_list_backups]
348437
def list_backups(instance_id, database_id, backup_id):
349-
from google.cloud.spanner_admin_database_v1.types import \
350-
backup as backup_pb
438+
from google.cloud.spanner_admin_database_v1.types import backup as backup_pb
351439

352440
spanner_client = spanner.Client()
353441
database_admin_api = spanner_client.database_admin_api
@@ -444,8 +532,7 @@ def list_backups(instance_id, database_id, backup_id):
444532

445533
# [START spanner_delete_backup]
446534
def delete_backup(instance_id, backup_id):
447-
from google.cloud.spanner_admin_database_v1.types import \
448-
backup as backup_pb
535+
from google.cloud.spanner_admin_database_v1.types import backup as backup_pb
449536

450537
spanner_client = spanner.Client()
451538
database_admin_api = spanner_client.database_admin_api
@@ -486,8 +573,7 @@ def delete_backup(instance_id, backup_id):
486573

487574
# [START spanner_update_backup]
488575
def update_backup(instance_id, backup_id):
489-
from google.cloud.spanner_admin_database_v1.types import \
490-
backup as backup_pb
576+
from google.cloud.spanner_admin_database_v1.types import backup as backup_pb
491577

492578
spanner_client = spanner.Client()
493579
database_admin_api = spanner_client.database_admin_api
@@ -526,8 +612,7 @@ def create_database_with_version_retention_period(
526612
):
527613
"""Creates a database with a version retention period."""
528614

529-
from google.cloud.spanner_admin_database_v1.types import \
530-
spanner_database_admin
615+
from google.cloud.spanner_admin_database_v1.types import spanner_database_admin
531616

532617
spanner_client = spanner.Client()
533618
database_admin_api = spanner_client.database_admin_api
@@ -578,8 +663,7 @@ def create_database_with_version_retention_period(
578663
def copy_backup(instance_id, backup_id, source_backup_path):
579664
"""Copies a backup."""
580665

581-
from google.cloud.spanner_admin_database_v1.types import \
582-
backup as backup_pb
666+
from google.cloud.spanner_admin_database_v1.types import backup as backup_pb
583667

584668
spanner_client = spanner.Client()
585669
database_admin_api = spanner_client.database_admin_api
@@ -613,6 +697,54 @@ def copy_backup(instance_id, backup_id, source_backup_path):
613697

614698
# [END spanner_copy_backup]
615699

700+
# [START spanner_copy_backup_with_MR_CMEK]
701+
def copy_backup_with_multiple_kms_keys(
702+
instance_id, backup_id, source_backup_path, kms_key_names
703+
):
704+
"""Copies a backup."""
705+
706+
from google.cloud.spanner_admin_database_v1.types import backup as backup_pb
707+
from google.cloud.spanner_admin_database_v1 import CopyBackupEncryptionConfig
708+
709+
spanner_client = spanner.Client()
710+
database_admin_api = spanner_client.database_admin_api
711+
712+
encryption_config = {
713+
"encryption_type": CopyBackupEncryptionConfig.EncryptionType.CUSTOMER_MANAGED_ENCRYPTION,
714+
"kms_key_names": kms_key_names,
715+
}
716+
717+
# Create a backup object and wait for copy backup operation to complete.
718+
expire_time = datetime.utcnow() + timedelta(days=14)
719+
request = backup_pb.CopyBackupRequest(
720+
parent=database_admin_api.instance_path(spanner_client.project, instance_id),
721+
backup_id=backup_id,
722+
source_backup=source_backup_path,
723+
expire_time=expire_time,
724+
encryption_config=encryption_config,
725+
)
726+
727+
operation = database_admin_api.copy_backup(request)
728+
729+
# Wait for backup operation to complete.
730+
copy_backup = operation.result(2100)
731+
732+
# Verify that the copy backup is ready.
733+
assert copy_backup.state == backup_pb.Backup.State.READY
734+
735+
print(
736+
"Backup {} of size {} bytes was created at {} with version time {} using encryption keys {}".format(
737+
copy_backup.name,
738+
copy_backup.size_bytes,
739+
copy_backup.create_time,
740+
copy_backup.version_time,
741+
copy_backup.encryption_information,
742+
)
743+
)
744+
745+
746+
# [END spanner_copy_backup_with_MR_CMEK]
747+
616748

617749
if __name__ == "__main__": # noqa: C901
618750
parser = argparse.ArgumentParser(

‎samples/samples/backup_sample_test.py

+66-1
Original file line numberDiff line numberDiff line change
@@ -13,8 +13,8 @@
1313
# limitations under the License.
1414
import uuid
1515

16-
import pytest
1716
from google.api_core.exceptions import DeadlineExceeded
17+
import pytest
1818
from test_utils.retry import RetryErrors
1919

2020
import backup_sample
@@ -93,6 +93,49 @@ def test_create_backup_with_encryption_key(
9393
assert kms_key_name in out
9494

9595

96+
@pytest.mark.skip(reason="skipped since the KMS keys are not added on test "
97+
"project")
98+
@pytest.mark.dependency(name="create_backup_with_multiple_kms_keys")
99+
def test_create_backup_with_multiple_kms_keys(
100+
capsys,
101+
multi_region_instance,
102+
multi_region_instance_id,
103+
sample_multi_region_database,
104+
kms_key_names,
105+
):
106+
backup_sample.create_backup_with_multiple_kms_keys(
107+
multi_region_instance_id,
108+
sample_multi_region_database.database_id,
109+
CMEK_BACKUP_ID,
110+
kms_key_names,
111+
)
112+
out, _ = capsys.readouterr()
113+
assert CMEK_BACKUP_ID in out
114+
assert kms_key_names[0] in out
115+
assert kms_key_names[1] in out
116+
assert kms_key_names[2] in out
117+
118+
119+
@pytest.mark.skip(reason="skipped since the KMS keys are not added on test "
120+
"project")
121+
@pytest.mark.dependency(depends=["create_backup_with_multiple_kms_keys"])
122+
def test_copy_backup_with_multiple_kms_keys(
123+
capsys, multi_region_instance_id, spanner_client, kms_key_names
124+
):
125+
source_backup_path = (
126+
spanner_client.project_name
127+
+ "/instances/"
128+
+ multi_region_instance_id
129+
+ "/backups/"
130+
+ CMEK_BACKUP_ID
131+
)
132+
backup_sample.copy_backup_with_multiple_kms_keys(
133+
multi_region_instance_id, COPY_BACKUP_ID, source_backup_path, kms_key_names
134+
)
135+
out, _ = capsys.readouterr()
136+
assert COPY_BACKUP_ID in out
137+
138+
96139
@pytest.mark.dependency(depends=["create_backup"])
97140
@RetryErrors(exception=DeadlineExceeded, max_tries=2)
98141
def test_restore_database(capsys, instance_id, sample_database):
@@ -121,6 +164,28 @@ def test_restore_database_with_encryption_key(
121164
assert kms_key_name in out
122165

123166

167+
@pytest.mark.skip(reason="skipped since the KMS keys are not added on test "
168+
"project")
169+
@pytest.mark.dependency(depends=["create_backup_with_multiple_kms_keys"])
170+
@RetryErrors(exception=DeadlineExceeded, max_tries=2)
171+
def test_restore_database_with_multiple_kms_keys(
172+
capsys,
173+
multi_region_instance_id,
174+
sample_multi_region_database,
175+
kms_key_names,
176+
):
177+
backup_sample.restore_database_with_multiple_kms_keys(
178+
multi_region_instance_id, CMEK_RESTORE_DB_ID, CMEK_BACKUP_ID, kms_key_names
179+
)
180+
out, _ = capsys.readouterr()
181+
assert (sample_multi_region_database.database_id + " restored to ") in out
182+
assert (CMEK_RESTORE_DB_ID + " from backup ") in out
183+
assert CMEK_BACKUP_ID in out
184+
assert kms_key_names[0] in out
185+
assert kms_key_names[1] in out
186+
assert kms_key_names[2] in out
187+
188+
124189
@pytest.mark.dependency(depends=["create_backup", "copy_backup"])
125190
def test_list_backup_operations(capsys, instance_id, sample_database):
126191
backup_sample.list_backup_operations(

0 commit comments

Comments
 (0)