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 <57220027+harshachinta@users.noreply.github.com> Co-authored-by: surbhigarg92 <surbhigarg.92@gmail.com> Co-authored-by: Sri Harsha CH <sriharshach@google.com>
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

Lines changed: 156 additions & 24 deletions
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

Lines changed: 66 additions & 1 deletion
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)
pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy