Skip to content

Commit f8d9bd3

Browse files
authored
feat: add samples for CMEK support (#275)
* feat: add samples for CMEK support * test: fix backups cleanup * test: correctly use database id for cmek restore * test: add clean up for databases * refactor: remove version time from sample * refactor: use user-provided key for creating encrypted backup message Co-authored-by: larkee <[email protected]>
1 parent 75f8340 commit f8d9bd3

File tree

4 files changed

+153
-1
lines changed

4 files changed

+153
-1
lines changed

samples/samples/backup_sample.py

Lines changed: 72 additions & 0 deletions
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, 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,42 @@ 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 using encryption key {}.".format(
151+
restore_info.backup_info.source_database,
152+
new_database_id,
153+
restore_info.backup_info.backup,
154+
new_database.encryption_config.kms_key_name,
155+
)
156+
)
157+
158+
159+
# [END spanner_restore_backup_with_encryption_key]
160+
161+
90162
# [START spanner_cancel_backup_create]
91163
def cancel_backup(instance_id, database_id, backup_id):
92164
spanner_client = spanner.Client()

samples/samples/backup_sample_test.py

Lines changed: 33 additions & 1 deletion
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_database_id()
44+
CMEK_BACKUP_ID = unique_backup_id()
45+
RETENTION_DATABASE_ID = unique_database_id()
4446
RETENTION_PERIOD = "7d"
4547

4648

@@ -54,6 +56,12 @@ def spanner_instance():
5456
op = instance.create()
5557
op.result(120) # block until completion
5658
yield instance
59+
for database_pb in instance.list_databases():
60+
database = instance.database(database_pb.name.split("/")[-1])
61+
database.drop()
62+
for backup_pb in instance.list_backups():
63+
backup = instance.backup(backup_pb.name.split("/")[-1])
64+
backup.delete()
5765
instance.delete()
5866

5967

@@ -77,6 +85,16 @@ def test_create_backup(capsys, database):
7785
assert BACKUP_ID in out
7886

7987

88+
def test_create_backup_with_encryption_key(capsys, spanner_instance, database):
89+
kms_key_name = "projects/{}/locations/{}/keyRings/{}/cryptoKeys/{}".format(
90+
spanner_instance._client.project, "us-central1", "spanner-test-keyring", "spanner-test-cmek"
91+
)
92+
backup_sample.create_backup_with_encryption_key(INSTANCE_ID, DATABASE_ID, CMEK_BACKUP_ID, kms_key_name)
93+
out, _ = capsys.readouterr()
94+
assert CMEK_BACKUP_ID in out
95+
assert kms_key_name in out
96+
97+
8098
# Depends on test_create_backup having run first
8199
@RetryErrors(exception=DeadlineExceeded, max_tries=2)
82100
def test_restore_database(capsys):
@@ -87,6 +105,20 @@ def test_restore_database(capsys):
87105
assert BACKUP_ID in out
88106

89107

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

samples/samples/snippets.py

Lines changed: 37 additions & 0 deletions
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

Lines changed: 11 additions & 0 deletions
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-cmek"
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)