Skip to content

Commit 96bee3c

Browse files
larkeelarkee
committed
feat: add samples for CMEK support
1 parent 50e2d8f commit 96bee3c

File tree

4 files changed

+151
-1
lines changed

4 files changed

+151
-1
lines changed

‎samples/samples/backup_sample.py

+73
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,42 @@ def create_backup(instance_id, database_id, backup_id, version_time):
5555

5656
# [END spanner_create_backup]
5757

58+
# [START spanner_create_backup_with_encryption_key]
59+
def create_backup_with_encryption_key(instance_id, database_id, backup_id, kms_key_name):
60+
"""Creates a backup for a database using a Customer Managed Encryption Key (CMEK)."""
61+
from google.cloud.spanner_admin_database_v1 import CreateBackupEncryptionConfig
62+
63+
spanner_client = spanner.Client()
64+
instance = spanner_client.instance(instance_id)
65+
database = instance.database(database_id)
66+
67+
# Create a backup
68+
expire_time = datetime.utcnow() + timedelta(days=14)
69+
encryption_config = {
70+
'encryption_type': CreateBackupEncryptionConfig.EncryptionType.CUSTOMER_MANAGED_ENCRYPTION,
71+
'kms_key_name': kms_key_name,
72+
}
73+
backup = instance.backup(backup_id, database=database, expire_time=expire_time, encryption_config=encryption_config)
74+
operation = backup.create()
75+
76+
# Wait for backup operation to complete.
77+
operation.result(1200)
78+
79+
# Verify that the backup is ready.
80+
backup.reload()
81+
assert backup.is_ready() is True
82+
83+
# Get the name, create time, backup size and encryption key.
84+
backup.reload()
85+
print(
86+
"Backup {} of size {} bytes was created at {} using encryption key {}".format(
87+
backup.name, backup.size_bytes, backup.create_time, backup.encryption_config.kms_key_name
88+
)
89+
)
90+
91+
92+
# [END spanner_create_backup_with_encryption_key]
93+
5894

5995
# [START spanner_restore_backup]
6096
def restore_database(instance_id, new_database_id, backup_id):
@@ -87,6 +123,43 @@ def restore_database(instance_id, new_database_id, backup_id):
87123
# [END spanner_restore_backup]
88124

89125

126+
# [START spanner_restore_backup_with_encryption_key]
127+
def restore_database_with_encryption_key(instance_id, new_database_id, backup_id, kms_key_name):
128+
"""Restores a database from a backup using a Customer Managed Encryption Key (CMEK)."""
129+
from google.cloud.spanner_admin_database_v1 import RestoreDatabaseEncryptionConfig
130+
131+
spanner_client = spanner.Client()
132+
instance = spanner_client.instance(instance_id)
133+
134+
# Start restoring an existing backup to a new database.
135+
backup = instance.backup(backup_id)
136+
encryption_config = {
137+
'encryption_type': RestoreDatabaseEncryptionConfig.EncryptionType.CUSTOMER_MANAGED_ENCRYPTION,
138+
'kms_key_name': kms_key_name,
139+
}
140+
new_database = instance.database(new_database_id, encryption_config=encryption_config)
141+
operation = new_database.restore(backup)
142+
143+
# Wait for restore operation to complete.
144+
operation.result(1600)
145+
146+
# Newly created database has restore information.
147+
new_database.reload()
148+
restore_info = new_database.restore_info
149+
print(
150+
"Database {} restored to {} from backup {} with version time {} using encryption key {}.".format(
151+
restore_info.backup_info.source_database,
152+
new_database_id,
153+
restore_info.backup_info.backup,
154+
restore_info.backup_info.version_time,
155+
new_database.encryption_config.kms_key_name,
156+
)
157+
)
158+
159+
160+
# [END spanner_restore_backup_with_encryption_key]
161+
162+
90163
# [START spanner_cancel_backup_create]
91164
def cancel_backup(instance_id, database_id, backup_id):
92165
spanner_client = spanner.Client()

‎samples/samples/backup_sample_test.py

+30-1
Original file line numberDiff line numberDiff line change
@@ -38,9 +38,11 @@ def unique_backup_id():
3838

3939
INSTANCE_ID = unique_instance_id()
4040
DATABASE_ID = unique_database_id()
41-
RETENTION_DATABASE_ID = unique_database_id()
4241
RESTORE_DB_ID = unique_database_id()
4342
BACKUP_ID = unique_backup_id()
43+
CMEK_RESTORE_DB_ID = unique_backup_id()
44+
CMEK_BACKUP_ID = unique_backup_id()
45+
RETENTION_DATABASE_ID = unique_database_id()
4446
RETENTION_PERIOD = "7d"
4547

4648

@@ -54,6 +56,9 @@ def spanner_instance():
5456
op = instance.create()
5557
op.result(120) # block until completion
5658
yield instance
59+
for backup_pb in instance.list_backups():
60+
backup = instance.backup(backup_pb.name)
61+
backup.delete()
5762
instance.delete()
5863

5964

@@ -77,6 +82,16 @@ def test_create_backup(capsys, database):
7782
assert BACKUP_ID in out
7883

7984

85+
def test_create_backup_with_encryption_key(capsys, spanner_instance, database):
86+
kms_key_name = "projects/{}/locations/{}/keyRings/{}/cryptoKeys/{}".format(
87+
spanner_instance._client.project, "us-central1", "spanner-test-keyring", "spanner-test-key-1568497514"
88+
)
89+
backup_sample.create_backup_with_encryption_key(INSTANCE_ID, DATABASE_ID, CMEK_BACKUP_ID, kms_key_name)
90+
out, _ = capsys.readouterr()
91+
assert CMEK_BACKUP_ID in out
92+
assert kms_key_name in out
93+
94+
8095
# Depends on test_create_backup having run first
8196
@RetryErrors(exception=DeadlineExceeded, max_tries=2)
8297
def test_restore_database(capsys):
@@ -87,6 +102,20 @@ def test_restore_database(capsys):
87102
assert BACKUP_ID in out
88103

89104

105+
# Depends on test_create_backup having run first
106+
@RetryErrors(exception=DeadlineExceeded, max_tries=2)
107+
def test_restore_database_with_encryption_key(capsys, spanner_instance):
108+
kms_key_name = "projects/{}/locations/{}/keyRings/{}/cryptoKeys/{}".format(
109+
spanner_instance._client.project, "us-central1", "spanner-test-keyring", "spanner-test-key-1568497514"
110+
)
111+
backup_sample.restore_database_with_encryption_key(INSTANCE_ID, CMEK_RESTORE_DB_ID, CMEK_BACKUP_ID, kms_key_name)
112+
out, _ = capsys.readouterr()
113+
assert (DATABASE_ID + " restored to ") in out
114+
assert (CMEK_RESTORE_DB_ID + " from backup ") in out
115+
assert CMEK_BACKUP_ID in out
116+
assert kms_key_name in out
117+
118+
90119
# Depends on test_create_backup having run first
91120
def test_list_backup_operations(capsys, spanner_instance):
92121
backup_sample.list_backup_operations(INSTANCE_ID, DATABASE_ID)

‎samples/samples/snippets.py

+37
Original file line numberDiff line numberDiff line change
@@ -92,6 +92,43 @@ def create_database(instance_id, database_id):
9292
# [END spanner_create_database]
9393

9494

95+
# [START spanner_create_database_with_encryption_key]
96+
def create_database_with_encryption_key(instance_id, database_id, kms_key_name):
97+
"""Creates a database with tables using a Customer Managed Encryption Key (CMEK)."""
98+
spanner_client = spanner.Client()
99+
instance = spanner_client.instance(instance_id)
100+
101+
database = instance.database(
102+
database_id,
103+
ddl_statements=[
104+
"""CREATE TABLE Singers (
105+
SingerId INT64 NOT NULL,
106+
FirstName STRING(1024),
107+
LastName STRING(1024),
108+
SingerInfo BYTES(MAX)
109+
) PRIMARY KEY (SingerId)""",
110+
"""CREATE TABLE Albums (
111+
SingerId INT64 NOT NULL,
112+
AlbumId INT64 NOT NULL,
113+
AlbumTitle STRING(MAX)
114+
) PRIMARY KEY (SingerId, AlbumId),
115+
INTERLEAVE IN PARENT Singers ON DELETE CASCADE""",
116+
],
117+
encryption_config={'kms_key_name': kms_key_name},
118+
)
119+
120+
operation = database.create()
121+
122+
print("Waiting for operation to complete...")
123+
operation.result(120)
124+
125+
print("Database {} created with encryption key {}".format(
126+
database.name, database.encryption_config.kms_key_name))
127+
128+
129+
# [END spanner_create_database_with_encryption_key]
130+
131+
95132
# [START spanner_insert_data]
96133
def insert_data(instance_id, database_id):
97134
"""Inserts sample data into the given database.

‎samples/samples/snippets_test.py

+11
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ def unique_database_id():
3333

3434
INSTANCE_ID = unique_instance_id()
3535
DATABASE_ID = unique_database_id()
36+
CMEK_DATABASE_ID = unique_database_id()
3637

3738

3839
@pytest.fixture(scope="module")
@@ -63,6 +64,16 @@ def test_create_database(database):
6364
database.reload()
6465

6566

67+
def test_create_database_with_encryption_config(capsys, spanner_instance):
68+
kms_key_name = "projects/{}/locations/{}/keyRings/{}/cryptoKeys/{}".format(
69+
spanner_instance._client.project, "us-central1", "spanner-test-keyring", "spanner-test-key-1568497514"
70+
)
71+
snippets.create_database_with_encryption_key(INSTANCE_ID, CMEK_DATABASE_ID, kms_key_name)
72+
out, _ = capsys.readouterr()
73+
assert CMEK_DATABASE_ID in out
74+
assert kms_key_name in out
75+
76+
6677
def test_insert_data(capsys):
6778
snippets.insert_data(INSTANCE_ID, DATABASE_ID)
6879
out, _ = capsys.readouterr()

0 commit comments

Comments
 (0)